1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
use crate::goal::runtime::GoalRuntimeHandle;
use crate::goal::service::GoalService;
use crate::goal::types::SessionGoal;
use crate::session::SessionStore;
use std::path::Path;
use std::sync::Arc;
/// Integrates the goal system with session and turn lifecycle.
///
/// Provides explicit hooks that callers should invoke at key points.
pub struct GoalExtension {
/// Shared goal service.
service: Arc<GoalService>,
/// Runtime handle for the current session.
runtime: Arc<GoalRuntimeHandle>,
}
impl GoalExtension {
/// Creates a new goal extension backed by the given service and runtime.
pub fn new(service: Arc<GoalService>, runtime: Arc<GoalRuntimeHandle>) -> Self {
Self { service, runtime }
}
/// Returns the shared goal service.
pub fn service(&self) -> &Arc<GoalService> {
&self.service
}
/// Returns the runtime handle.
pub fn runtime(&self) -> &Arc<GoalRuntimeHandle> {
&self.runtime
}
// ── Session lifecycle ──────────────────────────────────────
/// Called when a session starts. Registers the runtime with the service.
pub fn on_session_start(&self, session_id: &str) {
self.runtime.set_session_id(session_id);
self.service
.register_runtime(session_id.to_string(), Arc::clone(&self.runtime));
}
/// Called when a session is resumed from persistence.
/// Accepts a pre-loaded goal (typically from a `SessionSnapshot`) and
/// restores it into the runtime.
pub fn on_session_resume(
&self,
session_id: &str,
goal: Option<SessionGoal>,
) -> Option<SessionGoal> {
self.runtime.set_session_id(session_id);
if let Some(goal) = goal {
let should_auto_continue = goal.status.should_auto_continue();
self.runtime.update_goal(goal);
if should_auto_continue {
self.runtime.set_auto_continue(true);
}
return self.runtime.get_goal();
}
None
}
/// Called when a session ends. Unregisters the runtime and clears state.
pub fn on_session_end(&self, session_id: &str) {
self.service.unregister_runtime(session_id);
self.runtime.clear_goal();
}
/// Called when the thread/session becomes idle. Returns a continuation
/// prompt if the goal should auto-continue.
pub fn on_idle(&self) -> Option<String> {
self.runtime.continue_if_idle()
}
// ── Turn lifecycle ─────────────────────────────────────────
/// Called at the start of a turn. Initializes turn-level accounting.
pub fn on_turn_start(&self, _session_id: &str, _task: &str) {
self.runtime.start_turn();
}
/// Called at the end of a turn. Finalizes turn accounting.
pub fn on_turn_end(&self, _session_id: &str) {
self.runtime.finish_turn();
}
/// Called when a turn is aborted. Performs cleanup.
pub fn on_turn_abort(&self, _session_id: &str) {
self.runtime.finish_turn();
}
/// Called when a turn encounters an error.
pub fn on_turn_error(&self, error_message: &str) {
let lower = error_message.to_lowercase();
if lower.contains("usage limit") || lower.contains("rate limit") {
self.runtime.mark_usage_limited();
} else if lower.contains("fatal") || lower.contains("blocked") {
self.runtime.record_blocked_turn(error_message);
}
self.runtime.finish_turn();
}
// ── Tool lifecycle ─────────────────────────────────────────
/// Called after a tool call completes. Checks for budget limits.
/// Returns `Some(prompt)` if a budget-limit steering prompt should be injected.
pub fn on_tool_complete(&self) -> Option<String> {
self.runtime.budget_limit_prompt()
}
/// Called when token usage is reported. Records the delta.
/// Returns `true` if the budget was exceeded by this update.
pub fn on_token_usage(&self, input_tokens: u64, output_tokens: u64) -> bool {
let delta = (input_tokens + output_tokens) as i64;
self.runtime.record_tokens(delta)
}
// ── Persistence ────────────────────────────────────────────
/// Persists the current goal for direct `GoalService` callers.
pub fn persist_goal(
&self,
session_id: &str,
project_dir: &Path,
session_store: &SessionStore,
) -> anyhow::Result<()> {
self.service
.persist_goal(session_id, project_dir, session_store)
}
}