1mod 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
60pub const SESSION_EVENT_BROADCAST_CAPACITY: usize = 128;
62
63pub struct AgentHarnessOptions {
64 pub system_prompt: String,
66 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 pub observer: Arc<dyn RuntimeObserver>,
77 pub observation_context: ObservationContext,
79 pub observation_parent: Option<OperationId>,
81 pub stream_fn: Option<StreamFn>,
82 pub get_api_key: Option<GetApiKey>,
84 pub compaction: CompactionSettings,
86 pub compact_algorithms: Arc<CompactAlgorithmRegistry>,
88 pub before_tool_call: Option<BeforeToolCallHook>,
90 pub after_tool_call: Option<AfterToolCallHook>,
91 pub on_control_plane_prompt: Option<OnControlPlanePromptHook>,
93 pub budget_cap_usd: Option<f64>,
95 pub reload_skills_fn: Option<ReloadSkillsFn>,
97 pub on_turn_end: Option<OnTurnEndHook>,
99 pub turn_continuation_cap: Option<u32>,
101 pub max_iterations: Option<u32>,
103 pub runtime_extensions: Arc<dyn RuntimeExtensionPort>,
105 pub runtime_extension_model_context: ExtensionModelContextProjection,
107 pub runtime_extension_cwd: String,
109 pub runtime_extension_has_interactive_client: bool,
111}
112
113impl AgentHarnessOptions {
114 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
149pub 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 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}