Skip to main content

theway_core/agent/assembly/
mod.rs

1//! `AgentHarness` — opinionated assembly around the bare [`Agent`].
2//!
3//! The directory is split by domain: catalog mutation, lifecycle events, session
4//! persistence, and prompt execution. This root module owns the composed types and
5//! constructor so the public harness API remains one unit.
6//!
7//! ## Event planes
8//!
9//! - [`LoopEvent`](crate::agent::LoopEvent) covers single turn-loop internals.
10//! - [`SessionEvent`] covers cross-turn session lifecycle and configuration changes.
11//! - `SubagentJobEvent` covers multi-agent graph job telemetry through `SubagentJobRegistry`.
12//!
13//! External consumers subscribe through [`AgentHarness::subscribe_session_broadcast`].
14//! Synchronous callbacks registered with [`AgentHarness::subscribe_harness`] are isolated
15//! with `catch_unwind` and must remain memory-only. Persistence happens at its explicit
16//! call sites before matching session events are broadcast.
17
18mod catalog;
19mod events;
20mod run;
21mod runtime_extensions;
22mod session;
23
24use std::sync::Arc;
25
26use parking_lot::Mutex;
27use theway_llm_provider::Model;
28use tokio::sync::broadcast;
29
30use crate::agent::session::session::Session;
31use crate::agent::{Agent, AgentOptions, LoopListener};
32use crate::observability::{
33    ObservationContext, OperationId, RuntimeObserver, noop_runtime_observer,
34};
35use crate::types::*;
36
37use self::catalog::{build_system_prompt, deduplicate_tools};
38use super::compaction::algorithm::CompactAlgorithmRegistry;
39use super::compaction::compaction::{CompactionSettings, DEFAULT_COMPACTION_SETTINGS};
40use super::cost::{CostSnapshot, CostTracker};
41use super::runtime_extensions::{
42    ExtensionModelContextProjection, NoopRuntimeExtensionPort, RuntimeExtensionPort,
43};
44use super::types::{PromptTemplate, Skill};
45
46pub use self::events::{
47    DEFAULT_TURN_CONTINUATION_CAP, OnTurnEndContext, OnTurnEndHook, SessionEvent, SessionListener,
48    TurnEndAction, TurnEndDecision,
49};
50
51#[cfg(test)]
52use self::run::{extract_user_message_text, extract_user_prompt_text, preview_for_banner};
53#[cfg(test)]
54use self::session::{cap_control_plane_audit_label, finish_persisted_run, make_session_listener};
55#[cfg(test)]
56use crate::agent::AgentRunError;
57#[cfg(test)]
58use crate::agent::session::session::BranchSummaryInput;
59
60/// Capacity of the [`SessionEvent`] broadcast channel.
61pub const SESSION_EVENT_BROADCAST_CAPACITY: usize = 128;
62
63pub struct AgentHarnessOptions {
64    /// Base system prompt prepended to the rendered skill catalog.
65    pub system_prompt: String,
66    /// Active model. `None` when no model is configured yet (daemon launches
67    /// model-less and the client injects one per-session via `set_model`); the
68    /// turn loop errors with "Agent has no model set" until one is assigned.
69    pub model: Option<Model>,
70    pub thinking_level: ThinkingLevel,
71    pub skills: Vec<Skill>,
72    pub prompt_templates: Vec<PromptTemplate>,
73    pub tools: Vec<Arc<dyn AgentTool>>,
74    pub session: Session,
75    /// Embedder-owned runtime observer shared with the inner bare Agent.
76    pub observer: Arc<dyn RuntimeObserver>,
77    /// Session/run/job/node correlation inherited by inner operations.
78    pub observation_context: ObservationContext,
79    /// Optional parent operation for the inner agent run.
80    pub observation_parent: Option<OperationId>,
81    pub stream_fn: Option<StreamFn>,
82    /// Optional per-provider credential resolver applied to every model call.
83    pub get_api_key: Option<GetApiKey>,
84    /// Auto-compaction thresholds. Defaults to [`DEFAULT_COMPACTION_SETTINGS`].
85    pub compaction: CompactionSettings,
86    /// Custom compaction algorithm registry. The builtin algorithm is always available.
87    pub compact_algorithms: Arc<CompactAlgorithmRegistry>,
88    /// Optional tool hooks supplied by the embedding runtime.
89    pub before_tool_call: Option<BeforeToolCallHook>,
90    pub after_tool_call: Option<AfterToolCallHook>,
91    /// Control-plane prompt resolver. `None` keeps prompt-classified tool calls fail-closed.
92    pub on_control_plane_prompt: Option<OnControlPlanePromptHook>,
93    /// Per-session USD cap. `None` disables the check.
94    pub budget_cap_usd: Option<f64>,
95    /// Embedder-owned loader used by [`AgentHarness::reload_skills_from_disk`].
96    pub reload_skills_fn: Option<ReloadSkillsFn>,
97    /// Optional cross-prompt turn completion hook.
98    pub on_turn_end: Option<OnTurnEndHook>,
99    /// Maximum hook-driven continuations for one prompt cycle.
100    pub turn_continuation_cap: Option<u32>,
101    /// Hard cap on inner agent loop iterations. `None` is unbounded.
102    pub max_iterations: Option<u32>,
103    /// Engine-independent lifecycle port supplied by the embedding runtime.
104    pub runtime_extensions: Arc<dyn RuntimeExtensionPort>,
105    /// Reconstructed, de-duplicated persistent model context for this branch.
106    pub runtime_extension_model_context: ExtensionModelContextProjection,
107    /// Working directory included in runtime-extension lifecycle context.
108    pub runtime_extension_cwd: String,
109    /// Whether this harness is currently attached to an interactive client.
110    pub runtime_extension_has_interactive_client: bool,
111}
112
113impl AgentHarnessOptions {
114    /// `model` may be a concrete `Model` (wrapped into `Some`) or an
115    /// `Option<Model>` (kept as-is) so both the model-less daemon path and the
116    /// existing concrete-model callers build through one constructor.
117    pub fn new(model: impl Into<Option<Model>>, session: Session) -> Self {
118        Self {
119            system_prompt: String::new(),
120            model: model.into(),
121            thinking_level: ThinkingLevel::Off,
122            skills: Vec::new(),
123            prompt_templates: Vec::new(),
124            tools: Vec::new(),
125            session,
126            observer: noop_runtime_observer(),
127            observation_context: ObservationContext::default(),
128            observation_parent: None,
129            stream_fn: None,
130            get_api_key: None,
131            compaction: DEFAULT_COMPACTION_SETTINGS.clone(),
132            compact_algorithms: Arc::new(CompactAlgorithmRegistry::new()),
133            before_tool_call: None,
134            after_tool_call: None,
135            on_control_plane_prompt: None,
136            budget_cap_usd: None,
137            reload_skills_fn: None,
138            on_turn_end: None,
139            turn_continuation_cap: None,
140            max_iterations: None,
141            runtime_extensions: Arc::new(NoopRuntimeExtensionPort),
142            runtime_extension_model_context: ExtensionModelContextProjection::default(),
143            runtime_extension_cwd: ".".into(),
144            runtime_extension_has_interactive_client: false,
145        }
146    }
147}
148
149/// Async loader closure invoked by [`AgentHarness::reload_skills_from_disk`].
150pub type ReloadSkillsFn = Arc<
151    dyn Fn() -> std::pin::Pin<
152            Box<dyn std::future::Future<Output = super::skills::LoadSkillsOutput> + Send>,
153        > + Send
154        + Sync,
155>;
156
157#[derive(Debug, thiserror::Error)]
158pub enum ReloadSkillsError {
159    #[error("reload_skills_fn was not configured at harness construction")]
160    NotConfigured,
161}
162
163pub struct AgentHarness {
164    /// Compaction triggers live in a sibling module and require crate-level access.
165    pub(crate) agent: Arc<Agent>,
166    pub(crate) session: Session,
167    skills: Mutex<Vec<Skill>>,
168    base_system_prompt: String,
169    templates: Mutex<Vec<PromptTemplate>>,
170    pub(crate) compaction_settings: Mutex<CompactionSettings>,
171    pub(crate) compact_algorithms: Arc<CompactAlgorithmRegistry>,
172    pub(crate) stream_fn: Option<StreamFn>,
173    harness_listeners: Arc<Mutex<Vec<SessionListener>>>,
174    session_broadcast_tx: broadcast::Sender<SessionEvent>,
175    session_start_emitted: Mutex<bool>,
176    cost: CostTracker,
177    budget_cap_usd: Option<f64>,
178    reload_skills_fn: Option<ReloadSkillsFn>,
179    on_turn_end: Option<OnTurnEndHook>,
180    turn_continuation_cap: u32,
181    active_hook_cancel: Mutex<Option<tokio_util::sync::CancellationToken>>,
182    runtime_extensions: Arc<runtime_extensions::HarnessRuntimeExtensions>,
183}
184
185impl AgentHarness {
186    pub fn new(options: AgentHarnessOptions) -> Self {
187        let session_id = options.observation_context.session_id.clone();
188        let runtime_session_id = session_id
189            .clone()
190            .filter(|value| !value.trim().is_empty())
191            .unwrap_or_else(|| "local-session".into());
192        let runtime_extensions = Arc::new(runtime_extensions::HarnessRuntimeExtensions::new(
193            Arc::clone(&options.runtime_extensions),
194            runtime_session_id,
195            if options.runtime_extension_cwd.trim().is_empty() {
196                ".".into()
197            } else {
198                options.runtime_extension_cwd.clone()
199            },
200            options.runtime_extension_has_interactive_client,
201            options
202                .model
203                .as_ref()
204                .map(|m| theway_contract::extension::ExtensionModelRef {
205                    provider: m.provider.0.clone(),
206                    model: m.id.clone(),
207                }),
208            options.runtime_extension_model_context.clone(),
209        ));
210        let (tools, dropped_tools) = deduplicate_tools(options.tools);
211        if !dropped_tools.is_empty() {
212            tracing::warn!(
213                "dropped duplicate tool name(s) at harness construction (earlier registrations win): {}",
214                dropped_tools.join(", ")
215            );
216        }
217        let state = AgentState {
218            model: options.model,
219            thinking_level: Some(options.thinking_level),
220            tools,
221            system_prompt: build_system_prompt(&options.system_prompt, &options.skills),
222            ..Default::default()
223        };
224
225        let transform_runtime = Arc::clone(&runtime_extensions);
226        let transform_context: TransformContext = Arc::new(move |messages, cancel| {
227            let runtime = Arc::clone(&transform_runtime);
228            Box::pin(async move {
229                let messages = runtime.transform_context(messages, cancel).await;
230                crate::agent::context::virtualize_tool_results(messages)
231            })
232        });
233
234        let request_runtime = Arc::clone(&runtime_extensions);
235        let transform_model_request: TransformModelRequest = Arc::new(move |request, cancel| {
236            let runtime = Arc::clone(&request_runtime);
237            Box::pin(async move {
238                runtime
239                    .before_model_request(request, u32::MAX, cancel)
240                    .await
241            })
242        });
243        let provider_request_interceptor = runtime_extensions.provider_request_interceptor();
244
245        let message_runtime = Arc::clone(&runtime_extensions);
246        let transform_message: TransformMessage = Arc::new(move |message, cancel| {
247            let runtime = Arc::clone(&message_runtime);
248            Box::pin(async move { runtime.transform_message(message, cancel).await })
249        });
250
251        let configured_before_tool_call = options.before_tool_call.clone();
252        let before_tool_runtime = Arc::clone(&runtime_extensions);
253        let before_tool_call: BeforeToolCallHook = Arc::new(move |context, cancel| {
254            let runtime = Arc::clone(&before_tool_runtime);
255            let configured = configured_before_tool_call.clone();
256            Box::pin(async move {
257                let extension = runtime.before_tool_call(&context, &cancel).await;
258                if extension.block {
259                    return extension;
260                }
261                match configured {
262                    Some(hook) => hook(context, cancel).await,
263                    None => extension,
264                }
265            })
266        });
267
268        let after_tool_runtime = Arc::clone(&runtime_extensions);
269        let transform_tool_result: AfterToolCallHook = Arc::new(move |context, cancel| {
270            let runtime = Arc::clone(&after_tool_runtime);
271            Box::pin(async move { runtime.transform_tool_result(&context, &cancel).await })
272        });
273
274        let agent = Agent::new(AgentOptions {
275            initial_state: Some(state),
276            transform_context: Some(transform_context),
277            transform_model_request: Some(transform_model_request),
278            transform_message: Some(transform_message),
279            provider_request_interceptor: Some(provider_request_interceptor),
280            stream_fn: options.stream_fn.clone(),
281            get_api_key: options.get_api_key.clone(),
282            before_tool_call: Some(before_tool_call),
283            after_tool_call: options.after_tool_call.clone(),
284            transform_tool_result: Some(transform_tool_result),
285            on_control_plane_prompt: options.on_control_plane_prompt.clone(),
286            session_id,
287            observer: Arc::clone(&options.observer),
288            observation_context: options.observation_context.clone(),
289            observation_parent: options.observation_parent,
290            max_iterations: options.max_iterations,
291            ..Default::default()
292        });
293
294        let cost = CostTracker::new();
295        let _ = agent.subscribe_sync(cost.as_callback());
296        let (session_broadcast_tx, _) = broadcast::channel(SESSION_EVENT_BROADCAST_CAPACITY);
297
298        Self {
299            agent: Arc::new(agent),
300            session: options.session,
301            skills: Mutex::new(options.skills),
302            base_system_prompt: options.system_prompt,
303            templates: Mutex::new(options.prompt_templates),
304            compaction_settings: Mutex::new(options.compaction),
305            compact_algorithms: options.compact_algorithms,
306            stream_fn: options.stream_fn,
307            harness_listeners: Arc::new(Mutex::new(Vec::new())),
308            session_broadcast_tx,
309            session_start_emitted: Mutex::new(false),
310            cost,
311            budget_cap_usd: options.budget_cap_usd,
312            reload_skills_fn: options.reload_skills_fn,
313            on_turn_end: options.on_turn_end,
314            turn_continuation_cap: options
315                .turn_continuation_cap
316                .unwrap_or(DEFAULT_TURN_CONTINUATION_CAP),
317            active_hook_cancel: Mutex::new(None),
318            runtime_extensions,
319        }
320    }
321
322    pub fn cost(&self) -> CostSnapshot {
323        self.cost.snapshot()
324    }
325
326    pub fn reset_cost(&self) {
327        self.cost.reset();
328    }
329
330    pub fn agent_arc(&self) -> Arc<Agent> {
331        Arc::clone(&self.agent)
332    }
333
334    pub fn agent(&self) -> &Agent {
335        &self.agent
336    }
337
338    pub fn session(&self) -> &Session {
339        &self.session
340    }
341
342    pub fn runtime_extensions(&self) -> &Arc<dyn RuntimeExtensionPort> {
343        self.runtime_extensions.port()
344    }
345
346    pub fn abort(&self) {
347        self.agent.abort();
348        if let Some(token) = self.active_hook_cancel.lock().as_ref() {
349            token.cancel();
350        }
351    }
352
353    pub fn interrupt(&self) {
354        self.agent.interrupt();
355    }
356
357    pub fn enqueue_steering(&self, message: AgentMessage) {
358        self.agent.enqueue_steering(message);
359    }
360
361    pub fn enqueue_follow_up(&self, message: AgentMessage) {
362        self.agent.enqueue_follow_up(message);
363    }
364
365    pub fn subscribe(&self, listener: LoopListener) -> impl FnOnce() {
366        self.agent.subscribe(listener)
367    }
368}
369
370#[cfg(test)]
371tests_bridge_macro::tests_bridge!("agent/assembly");
372
373#[cfg(test)]
374mod assembly_linecov_tests {
375    tests_bridge_macro::tests_bridge!("agent/assembly/linecov");
376}