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

starweaver_runtime/
agent.rs
