Skip to main content

theway_core/agent/assembly/
events.rs

1use std::sync::Arc;
2
3use tokio::sync::broadcast;
4
5use crate::types::AgentMessage;
6
7use super::AgentHarness;
8
9impl AgentHarness {
10    /// Register a harness-level lifecycle listener. Returns an unsubscriber closure.
11    ///
12    /// Listener panics are caught — see [`SessionEvent`] for the isolation contract. The
13    /// returned closure removes the listener; calling it twice is a no-op.
14    pub fn subscribe_harness(&self, listener: SessionListener) -> Box<dyn FnOnce() + Send> {
15        self.harness_listeners.lock().push(listener.clone());
16        let target = Arc::as_ptr(&listener) as *const () as usize;
17        let listeners = Arc::clone(&self.harness_listeners);
18        Box::new(move || {
19            let mut guard = listeners.lock();
20            if let Some(index) = guard
21                .iter()
22                .position(|item| (Arc::as_ptr(item) as *const () as usize) == target)
23            {
24                guard.remove(index);
25            }
26        })
27    }
28
29    /// Obtain a new [`tokio::sync::broadcast::Receiver`] for the [`SessionEvent`] broadcast
30    /// channel. The receiver sees all events emitted after subscription.
31    pub fn subscribe_session_broadcast(&self) -> broadcast::Receiver<SessionEvent> {
32        self.session_broadcast_tx.subscribe()
33    }
34
35    /// Dispatch to isolated synchronous callbacks, then publish to the broadcast channel.
36    pub(crate) fn emit_harness_event(&self, event: SessionEvent) {
37        let listeners = self.harness_listeners.lock().clone();
38        for listener in listeners {
39            let event = event.clone();
40            let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || listener(event)));
41        }
42        let _ = self.session_broadcast_tx.send(event);
43    }
44
45    pub(super) async fn ensure_session_start_emitted(&self) {
46        let should_emit = {
47            let mut emitted = self.session_start_emitted.lock();
48            if *emitted {
49                false
50            } else {
51                *emitted = true;
52                true
53            }
54        };
55        if !should_emit {
56            return;
57        }
58        let messages_replayed = self.agent.state().messages.len();
59        self.runtime_extensions.ensure_session_start().await;
60        self.emit_harness_event(SessionEvent::Started { messages_replayed });
61    }
62
63    /// Reconstruct the session-scoped extension runtime before it begins
64    /// serving prompts. Idempotent with the first prompt path.
65    pub async fn start_runtime_extensions(&self) {
66        self.ensure_session_start_emitted().await;
67    }
68
69    /// Cancel any active run, wait for its awaited cleanup, then publish the
70    /// extension session-shutdown lifecycle exactly once.
71    pub async fn shutdown_runtime_extensions(&self) {
72        self.abort();
73        // Bounded: a run wedged on a misbehaving listener must not block
74        // session activation/reap forever. The await-listener backstop in
75        // `emit` caps the wait in practice; this is defense in depth.
76        if tokio::time::timeout(SHUTDOWN_IDLE_TIMEOUT, self.agent.wait_until_idle())
77            .await
78            .is_err()
79        {
80            tracing::warn!("runtime-extension shutdown: run did not go idle in time; continuing");
81        }
82        self.runtime_extensions.shutdown().await;
83    }
84}
85
86/// Hard bound for `shutdown_runtime_extensions` waiting on the active run.
87const SHUTDOWN_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(45);
88
89/// Harness-level lifecycle events emitted in addition to the inner agent's per-turn events.
90#[derive(Clone, Debug)]
91pub enum SessionEvent {
92    /// First prompt entry after construction.
93    Started { messages_replayed: usize },
94    /// Auto- or manual compaction completed.
95    Compaction {
96        from_hook: bool,
97        summary: String,
98        tokens_before: u64,
99    },
100    /// The active session branch changed.
101    Branch {
102        from_entry_id: Option<String>,
103        to_entry_id: Option<String>,
104        summary_entry_id: Option<String>,
105    },
106    /// A best-effort persistence operation failed.
107    PersistenceError { context: String, message: String },
108    /// A turn-completion hook made a lifecycle decision.
109    TurnDecision {
110        decision: &'static str,
111        continuation_count: u32,
112        reason: Option<String>,
113        next_prompt_preview: Option<String>,
114    },
115    /// The skill catalog was hot-reloaded.
116    SkillsReloaded { total: usize },
117    /// An extension handled input without starting an LLM run.
118    ExtensionCommandOutcome {
119        outcome: theway_contract::extension::ExtensionCommandOutcome,
120    },
121}
122
123pub type SessionListener = Arc<dyn Fn(SessionEvent) + Send + Sync>;
124
125/// Snapshot passed into [`OnTurnEndHook`] after a prompt cycle reaches a natural stop.
126#[derive(Clone)]
127pub struct OnTurnEndContext {
128    pub transcript: Vec<AgentMessage>,
129    pub continuation_count: u32,
130    pub last_user_prompt: Option<String>,
131}
132
133/// What the runtime should do after [`OnTurnEndHook`] inspects a completed prompt cycle.
134#[derive(Clone, Debug)]
135pub enum TurnEndAction {
136    Noop,
137    Stop,
138    Pause { reason: String },
139    Continue { prompt: String },
140}
141
142impl TurnEndAction {
143    /// Stable value persisted in turn-end audit entries; `Noop` deliberately has no audit.
144    pub fn as_audit_str(&self) -> Option<&'static str> {
145        match self {
146            Self::Noop => None,
147            Self::Stop => Some("stop"),
148            Self::Pause { .. } => Some("pause"),
149            Self::Continue { .. } => Some("continue"),
150        }
151    }
152}
153
154/// Decision envelope returned from [`OnTurnEndHook`].
155#[derive(Clone, Debug)]
156pub struct TurnEndDecision {
157    pub action: TurnEndAction,
158    pub payload: Option<serde_json::Value>,
159}
160
161impl From<TurnEndAction> for TurnEndDecision {
162    fn from(action: TurnEndAction) -> Self {
163        Self {
164            action,
165            payload: None,
166        }
167    }
168}
169
170pub type OnTurnEndHook = Arc<
171    dyn Fn(
172            OnTurnEndContext,
173            tokio_util::sync::CancellationToken,
174        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = TurnEndDecision> + Send>>
175        + Send
176        + Sync,
177>;
178
179/// Default maximum number of continuation iterations per prompt cycle.
180pub const DEFAULT_TURN_CONTINUATION_CAP: u32 = 25;