Skip to main content

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