starweaver_runtime/
agent.rs

1//! Bare agent runtime.
2
3use std::sync::Arc;
4
5use starweaver_context::{AgentContext, AgentInfo, ModelConfig, ToolConfig};
6use starweaver_core::{AgentId, CancellationToken};
7use starweaver_model::{ModelAdapter, ModelRequestParameters, ModelSettings, OutputMode};
8use starweaver_tools::{DynToolset, ToolRegistry};
9
10use starweaver_usage::UsageLimits;
11
12use crate::{
13    agent::helpers::merge_request_params,
14    capability::{AgentCapability, CapabilityBundle, resolve_capability_order},
15    executor::{DirectAgentExecutor, DynAgentExecutor},
16    graph::{AgentGraphTrace, AgentNode, GraphError, inspect_graph},
17    instructions::DynDynamicInstruction,
18    output::{
19        DynOutputFunction, OutputPolicy, OutputSchema, OutputValidator, SchemaOutputFunction,
20    },
21    trace::{DynTraceRecorder, NoopTraceRecorder},
22};
23
24mod helpers;
25mod overrides;
26mod run_loop;
27mod run_loop_helpers;
28mod runtime_helpers;
29mod types;
30
31pub use overrides::AgentOverride;
32pub use types::{
33    AgentEndStrategy, AgentError, AgentInput, AgentResult, AgentRuntimePolicy,
34    AgentToolExecutionMode,
35};
36
37/// Minimal agent builder/runtime.
38#[derive(Clone)]
39pub struct Agent {
40    default_id: AgentId,
41    default_name: String,
42    model: Arc<dyn ModelAdapter>,
43    instructions: Vec<String>,
44    dynamic_instructions: Vec<DynDynamicInstruction>,
45    model_settings: Option<ModelSettings>,
46    request_params: ModelRequestParameters,
47    output_schema: Option<OutputSchema>,
48    output_validators: Vec<Arc<dyn OutputValidator>>,
49    output_functions: Vec<DynOutputFunction>,
50    usage_limits: Option<UsageLimits>,
51    tools: ToolRegistry,
52    toolsets: Vec<DynToolset>,
53    capabilities: Vec<Arc<dyn AgentCapability>>,
54    stream_observers: Vec<Arc<dyn AgentCapability>>,
55    cancellation_token: Option<CancellationToken>,
56    executor: DynAgentExecutor,
57    trace_recorder: DynTraceRecorder,
58    policy: AgentRuntimePolicy,
59    model_config: Option<ModelConfig>,
60    tool_config: Option<ToolConfig>,
61}
62
63impl Agent {
64    /// Create an agent with a model adapter.
65    #[must_use]
66    pub fn new(model: Arc<dyn ModelAdapter>) -> Self {
67        Self {
68            default_id: AgentId::default(),
69            default_name: AgentId::default().as_str().to_string(),
70            model,
71            instructions: Vec::new(),
72            dynamic_instructions: Vec::new(),
73            model_settings: None,
74            request_params: ModelRequestParameters::default(),
75            output_schema: None,
76            output_validators: Vec::new(),
77            output_functions: Vec::new(),
78            usage_limits: None,
79            tools: ToolRegistry::new(),
80            toolsets: Vec::new(),
81            capabilities: Vec::new(),
82            stream_observers: Vec::new(),
83            cancellation_token: None,
84            executor: Arc::new(DirectAgentExecutor),
85            trace_recorder: Arc::new(NoopTraceRecorder),
86            policy: AgentRuntimePolicy::default(),
87            model_config: None,
88            tool_config: None,
89        }
90    }
91
92    /// Return the default agent id used when this agent creates a context.
93    #[must_use]
94    pub const fn agent_id(&self) -> &AgentId {
95        &self.default_id
96    }
97
98    /// Return the default human-readable agent name.
99    #[must_use]
100    pub fn agent_name(&self) -> &str {
101        &self.default_name
102    }
103
104    /// Set the default agent id used when this agent creates a context.
105    #[must_use]
106    pub fn with_agent_id(mut self, agent_id: AgentId) -> Self {
107        self.default_id = agent_id;
108        self
109    }
110
111    /// Set the default human-readable agent name.
112    #[must_use]
113    pub fn with_agent_name(mut self, agent_name: impl Into<String>) -> Self {
114        self.default_name = agent_name.into();
115        self
116    }
117
118    /// Set both the default agent id and human-readable agent name.
119    #[must_use]
120    pub fn with_agent_identity(mut self, agent_id: AgentId, agent_name: impl Into<String>) -> Self {
121        self.default_id = agent_id;
122        self.default_name = agent_name.into();
123        self
124    }
125
126    /// Create a fresh context using this agent's configured identity.
127    #[must_use]
128    pub fn new_context(&self) -> AgentContext {
129        let mut context = AgentContext::new(self.default_id.clone());
130        self.apply_context_identity(&mut context);
131        context
132    }
133
134    fn apply_context_identity(&self, context: &mut AgentContext) {
135        context.agent_registry.insert(
136            self.default_id.as_str().to_string(),
137            AgentInfo::new(self.default_id.as_str(), self.default_name.clone()),
138        );
139        if self.default_name != self.default_id.as_str() {
140            context.metadata.insert(
141                "agent_name".to_string(),
142                serde_json::json!(self.default_name.as_str()),
143            );
144        }
145    }
146
147    /// Add a static instruction.
148    #[must_use]
149    pub fn with_instruction(mut self, instruction: impl Into<String>) -> Self {
150        self.instructions.push(instruction.into());
151        self
152    }
153
154    /// Add a dynamic instruction.
155    #[must_use]
156    pub fn with_dynamic_instruction(mut self, instruction: DynDynamicInstruction) -> Self {
157        self.dynamic_instructions.push(instruction);
158        self
159    }
160
161    /// Set default model settings.
162    #[must_use]
163    pub fn with_model_settings(mut self, settings: ModelSettings) -> Self {
164        self.model_settings = Some(settings);
165        self
166    }
167
168    /// Set default request parameters.
169    #[must_use]
170    pub fn with_request_params(mut self, params: ModelRequestParameters) -> Self {
171        self.request_params = params;
172        self
173    }
174
175    /// Set runtime tools.
176    #[must_use]
177    pub fn with_tools(mut self, tools: ToolRegistry) -> Self {
178        self.tools = tools;
179        self
180    }
181
182    /// Add one runtime toolset that is materialized for each agent context.
183    #[must_use]
184    pub fn with_toolset(mut self, toolset: DynToolset) -> Self {
185        self.toolsets.push(toolset);
186        self
187    }
188
189    /// Add many runtime toolsets that are materialized for each agent context.
190    #[must_use]
191    pub fn with_toolsets(mut self, toolsets: impl IntoIterator<Item = DynToolset>) -> Self {
192        self.toolsets.extend(toolsets);
193        self
194    }
195
196    /// Merge additional runtime tools into this agent.
197    #[must_use]
198    pub fn with_appended_tools(mut self, tools: &ToolRegistry) -> Self {
199        self.tools.insert_registry(tools);
200        self
201    }
202
203    /// Return a clone of the runtime tool registry.
204    #[must_use]
205    pub fn tools(&self) -> ToolRegistry {
206        let mut tools = self.tools.clone();
207        for toolset in &self.toolsets {
208            tools.insert_toolset(toolset);
209        }
210        tools
211    }
212
213    /// Prepare this agent's static tools and context-aware toolsets for a concrete context.
214    ///
215    /// This uses the same lifecycle-aware path as the normal run loop and is intended for
216    /// host operations that need to execute a previously suspended tool call.
217    ///
218    /// # Errors
219    ///
220    /// Returns an agent error when a context-aware toolset cannot be prepared.
221    pub async fn prepare_tools_for_context(
222        &self,
223        context: &mut AgentContext,
224    ) -> Result<ToolRegistry, AgentError> {
225        self.prepare_run_tools(context, true).await
226    }
227
228    /// Close context-aware toolsets after host-side execution outside the normal run loop.
229    pub async fn close_toolsets_for_context(&self, context: &mut AgentContext) {
230        self.close_run_toolsets(context).await;
231    }
232
233    /// Set the agent-level retry default for runtime tools.
234    #[must_use]
235    pub const fn with_tool_retries(mut self, max_retries: usize) -> Self {
236        self.tools.set_max_retries(max_retries);
237        self
238    }
239
240    /// Set structured output schema.
241    #[must_use]
242    pub fn with_output_schema(mut self, schema: OutputSchema) -> Self {
243        self.output_schema = Some(schema);
244        self
245    }
246
247    /// Apply a complete output policy.
248    #[must_use]
249    pub fn with_output_policy(mut self, policy: OutputPolicy) -> Self {
250        let (schema, validators, functions, retries, mode, allow_text_output, allow_image_output) =
251            policy.into_parts();
252        let schema_output_function = match (schema.as_ref(), mode) {
253            (Some(schema), Some(OutputMode::Tool | OutputMode::ToolOrText)) => {
254                Some(Arc::new(SchemaOutputFunction::new(schema.clone())) as DynOutputFunction)
255            }
256            _ => None,
257        };
258        if let Some(schema) = schema {
259            self.output_schema = Some(schema);
260        }
261        self.output_validators.extend(validators);
262        if let Some(function) = schema_output_function {
263            self.output_functions.push(function);
264        }
265        self.output_functions.extend(functions);
266        if let Some(retries) = retries {
267            self.policy.output_retries = retries;
268        }
269        if let Some(mode) = mode {
270            self.request_params.output_mode = Some(mode);
271        }
272        if let Some(allow_text_output) = allow_text_output {
273            self.request_params.allow_text_output = Some(allow_text_output);
274        }
275        if let Some(allow_image_output) = allow_image_output {
276            self.request_params.allow_image_output = Some(allow_image_output);
277        }
278        self
279    }
280
281    /// Add an output validator.
282    #[must_use]
283    pub fn with_output_validator(mut self, validator: Arc<dyn OutputValidator>) -> Self {
284        self.output_validators.push(validator);
285        self
286    }
287
288    /// Add an output function.
289    #[must_use]
290    pub fn with_output_function(mut self, function: DynOutputFunction) -> Self {
291        self.output_functions.push(function);
292        self
293    }
294
295    /// Set usage limits.
296    #[must_use]
297    pub const fn with_usage_limits(mut self, limits: UsageLimits) -> Self {
298        self.usage_limits = Some(limits);
299        self
300    }
301
302    /// Set the full model config exposed to `AgentContext`.
303    #[must_use]
304    pub fn with_model_config(mut self, model_config: ModelConfig) -> Self {
305        self.model_config = Some(model_config);
306        self
307    }
308
309    /// Set tool-level configuration exposed to runtime tools.
310    #[must_use]
311    pub fn with_tool_config(mut self, tool_config: ToolConfig) -> Self {
312        self.tool_config = Some(tool_config);
313        self
314    }
315
316    /// Set the model context window exposed to `AgentContext` runtime instructions.
317    #[must_use]
318    pub fn with_context_window(mut self, context_window: u64) -> Self {
319        let mut model_config = self.model_config.unwrap_or_default();
320        model_config.context_window = Some(context_window);
321        self.model_config = Some(model_config);
322        self
323    }
324
325    /// Add a capability hook.
326    #[must_use]
327    pub fn with_capability(mut self, capability: Arc<dyn AgentCapability>) -> Self {
328        self.capabilities.push(capability);
329        self
330    }
331
332    pub(super) fn ordered_capabilities(&self) -> Result<Vec<Arc<dyn AgentCapability>>, AgentError> {
333        resolve_capability_order(&self.capabilities).map_err(AgentError::from)
334    }
335
336    pub(super) fn ordered_stream_observers(
337        &self,
338    ) -> Result<Vec<Arc<dyn AgentCapability>>, AgentError> {
339        resolve_capability_order(&self.stream_observers).map_err(AgentError::from)
340    }
341
342    pub(super) fn ordered_capabilities_for_validation(
343        &self,
344    ) -> Result<Vec<Arc<dyn AgentCapability>>, crate::capability::CapabilityError> {
345        resolve_capability_order(&self.capabilities)
346            .map_err(|error| crate::capability::CapabilityError::Failed(error.to_string()))
347    }
348
349    /// Add a stream observer hook.
350    #[must_use]
351    pub fn with_stream_observer(mut self, observer: Arc<dyn AgentCapability>) -> Self {
352        self.stream_observers.push(observer);
353        self
354    }
355
356    /// Set a cooperative cancellation token used by streaming callers.
357    #[must_use]
358    pub fn with_cancellation_token(mut self, token: CancellationToken) -> Self {
359        self.cancellation_token = Some(token);
360        self
361    }
362
363    /// Apply a composable capability bundle.
364    #[must_use]
365    pub fn with_capability_bundle(mut self, bundle: &dyn CapabilityBundle) -> Self {
366        self.apply_capability_bundle(bundle);
367        self
368    }
369
370    /// Set durable execution checkpoint handler.
371    #[must_use]
372    pub fn with_executor(mut self, executor: DynAgentExecutor) -> Self {
373        self.executor = executor;
374        self
375    }
376
377    /// Set runtime trace recorder.
378    #[must_use]
379    pub fn with_trace_recorder(mut self, recorder: DynTraceRecorder) -> Self {
380        self.trace_recorder = recorder;
381        self
382    }
383
384    /// Set runtime policy.
385    #[must_use]
386    pub const fn with_policy(mut self, policy: AgentRuntimePolicy) -> Self {
387        self.policy = policy;
388        self
389    }
390
391    /// Inspect graph transitions from a state snapshot.
392    ///
393    /// # Errors
394    ///
395    /// Returns an error when the inspected transition is invalid for the provided state.
396    pub fn inspect_graph(
397        &self,
398        start: AgentNode,
399        state: &crate::run::AgentRunState,
400    ) -> Result<AgentGraphTrace, GraphError> {
401        inspect_graph(
402            start,
403            state,
404            self.policy.max_steps,
405            self.policy.max_steps.saturating_mul(4),
406        )
407    }
408
409    /// Create a scoped override builder for tests and alternate run contexts.
410    #[must_use]
411    pub fn override_config(&self) -> AgentOverride {
412        AgentOverride::new(self.clone())
413    }
414
415    fn apply_capability_bundle(&mut self, bundle: &dyn CapabilityBundle) {
416        self.capabilities.extend(bundle.hooks());
417        self.stream_observers.extend(bundle.stream_observers());
418        self.instructions.extend(bundle.get_instructions());
419        self.dynamic_instructions
420            .extend(bundle.dynamic_instructions());
421        if let Some(tools) = bundle.get_tools() {
422            self.tools.insert_registry(&tools);
423        }
424        if let Some(settings) = bundle.model_settings() {
425            self.model_settings = Some(match &self.model_settings {
426                Some(current) => current.merge(&settings),
427                None => settings,
428            });
429        }
430        if let Some(params) = bundle.request_params() {
431            self.request_params = merge_request_params(&self.request_params, &params);
432        }
433        self.output_functions.extend(bundle.output_functions());
434        self.output_validators.extend(bundle.output_validators());
435        if let Some(limits) = bundle.usage_limits() {
436            self.usage_limits = Some(limits);
437        }
438    }
439}
starweaver_runtime/agent.rs

starweaver_runtime/
agent.rs
