starweaver_runtime/
agent.rs

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

starweaver_runtime/
agent.rs
