Skip to main content

runifold_agent/
builder.rs

1use std::sync::Arc;
2
3use runifold_core::CapabilitySet;
4use runifold_effect::{EffectExecutor, EffectRecoveryPolicy};
5use runifold_model::{
6    ArtifactResolvingModel, ArtifactScope, ArtifactStore, FeaturePolicy, GenerationOptions,
7    Message, Model, ModelRef, OutputFormat, ProviderToolSpec, ResponseMode,
8};
9use runifold_retrieval::{Document, RetrievalError, Retriever};
10use runifold_tool::{Tool, ToolRegistrationError};
11use schemars::JsonSchema;
12use thiserror::Error;
13
14use crate::agent::DynamicContext;
15use crate::{
16    Agent, AgentConfig, AgentDescriptor, AgentError, AgentFuture, AgentOutcome,
17    AgentRegistrationError, AgentRoute, GatewayMiddleware, StructuredAgent, ToolErrorPolicy,
18};
19
20/// Failure while assembling an [`Agent`].
21#[derive(Clone, Debug, Error, Eq, PartialEq)]
22#[non_exhaustive]
23pub enum AgentBuildError {
24    /// Tool registration failed.
25    #[error("agent Tool registration failed: {0}")]
26    Tool(#[from] ToolRegistrationError),
27    /// Child Agent route registration failed.
28    #[error("agent route registration failed: {0}")]
29    Route(#[from] AgentRegistrationError),
30    /// One model-facing name was used by both a Tool and an Agent.
31    #[error("callable name `{0}` is registered as both a Tool and an Agent")]
32    CallableNameCollision(String),
33    /// The Agent name is blank.
34    #[error("agent name cannot be empty")]
35    EmptyName,
36    /// The configured turn limit cannot execute any model turn.
37    #[error("max_turns must be greater than zero")]
38    ZeroMaxTurns,
39    /// A static or dynamic context registration was invalid.
40    #[error("agent retrieval configuration failed: {0}")]
41    Retrieval(#[from] RetrievalError),
42}
43
44/// Failure while building and immediately prompting an Agent.
45#[derive(Debug, Error)]
46#[non_exhaustive]
47pub enum AgentPromptError {
48    /// Fluent Agent assembly failed before model execution.
49    #[error("failed to build agent: {0}")]
50    Build(#[from] AgentBuildError),
51    /// Canonical Agent execution failed.
52    #[error("agent prompt failed: {0}")]
53    Run(#[from] AgentError),
54}
55
56/// Fluent assembly of one canonical [`Agent`].
57///
58/// Registration failures are retained and returned by [`Self::build`] so
59/// Tool and child Agent calls remain chainable without silently replacing an
60/// existing name.
61pub struct AgentBuilder {
62    agent: Agent,
63    error: Option<AgentBuildError>,
64}
65
66impl AgentBuilder {
67    /// Creates a builder around the same execution path as [`Agent::new`].
68    pub fn new(name: impl Into<String>, model: Arc<dyn Model>, model_ref: ModelRef) -> Self {
69        Self {
70            agent: Agent::new(name, model, model_ref),
71            error: None,
72        }
73    }
74
75    /// Appends a system instruction.
76    #[must_use]
77    pub fn system(mut self, instruction: impl Into<String>) -> Self {
78        self.agent
79            .instructions
80            .push(Message::system(instruction.into()));
81        self
82    }
83
84    /// Adds one static document as untrusted user-level context.
85    ///
86    /// The document is never promoted to a system instruction. Use
87    /// [`Self::system`] for trusted application policy.
88    #[must_use]
89    pub fn context(self, text: impl Into<String>) -> Self {
90        let id = format!("static-context-{}", self.agent.context.len() + 1);
91        match Document::new(id, text) {
92            Ok(document) => self.context_document(document),
93            Err(error) => self.with_error(error.into()),
94        }
95    }
96
97    /// Adds one validated static context document.
98    #[must_use]
99    pub fn context_document(mut self, document: Document) -> Self {
100        if self.error.is_none() {
101            self.agent.context.push(document);
102        }
103        self
104    }
105
106    /// Configures reference-only artifact persistence for Tools and resolves
107    /// those references only at the final model transport boundary.
108    #[must_use]
109    pub fn artifacts(mut self, scope: ArtifactScope, store: Arc<dyn ArtifactStore>) -> Self {
110        self.agent.model = Arc::new(ArtifactResolvingModel::new(
111            self.agent.model.clone(),
112            scope.clone(),
113            store.clone(),
114        ));
115        self.agent.tools = self.agent.tools.clone().with_artifact_store(scope, store);
116        self
117    }
118
119    /// Adds an owned dynamic context source.
120    #[must_use]
121    pub fn dynamic_context<R>(self, limit: usize, retriever: R) -> Self
122    where
123        R: Retriever + 'static,
124    {
125        self.shared_dynamic_context(limit, Arc::new(retriever))
126    }
127
128    /// Adds a shared, type-erased dynamic context source.
129    #[must_use]
130    pub fn shared_dynamic_context(mut self, limit: usize, retriever: Arc<dyn Retriever>) -> Self {
131        if self.error.is_none() {
132            if limit == 0 {
133                self.error = Some(RetrievalError::ZeroLimit.into());
134            } else {
135                self.agent
136                    .dynamic_context
137                    .push(DynamicContext { limit, retriever });
138            }
139        }
140        self
141    }
142
143    /// Registers an owned Tool.
144    #[must_use]
145    pub fn tool<T>(self, tool: T) -> Self
146    where
147        T: Tool + 'static,
148    {
149        self.shared_tool(Arc::new(tool))
150    }
151
152    /// Registers a shared, type-erased Tool.
153    #[must_use]
154    pub fn shared_tool(mut self, tool: Arc<dyn Tool>) -> Self {
155        if self.error.is_none()
156            && let Err(error) = self.agent.tools.register(tool)
157        {
158            self.error = Some(error.into());
159        }
160        self
161    }
162
163    fn with_error(mut self, error: AgentBuildError) -> Self {
164        if self.error.is_none() {
165            self.error = Some(error);
166        }
167        self
168    }
169
170    /// Registers a child Agent route with explicit delegated capabilities.
171    #[must_use]
172    pub fn child(
173        mut self,
174        descriptor: AgentDescriptor,
175        child: Arc<Agent>,
176        capabilities: CapabilitySet,
177    ) -> Self {
178        if self.error.is_none() {
179            let route = AgentRoute::new(descriptor, child).with_capabilities(capabilities);
180            if let Err(error) = self.agent.agents.register(route) {
181                self.error = Some(error.into());
182            }
183        }
184        self
185    }
186
187    /// Appends Gateway around-middleware.
188    #[must_use]
189    pub fn gateway_layer(mut self, middleware: Arc<dyn GatewayMiddleware>) -> Self {
190        self.agent.agents.push_middleware(middleware);
191        self
192    }
193
194    /// Sets the maximum nested Agent delegation depth.
195    #[must_use]
196    pub fn max_delegation_depth(mut self, max_depth: u32) -> Self {
197        self.agent.agents = self.agent.agents.with_max_depth(max_depth);
198        self
199    }
200
201    /// Sets the local model-turn limit.
202    #[must_use]
203    pub const fn max_turns(mut self, max_turns: u32) -> Self {
204        self.agent.config.max_turns = max_turns;
205        self
206    }
207
208    /// Sets Tool failure behavior.
209    #[must_use]
210    pub const fn tool_error_policy(mut self, policy: ToolErrorPolicy) -> Self {
211        self.agent.config.tool_error_policy = policy;
212        self
213    }
214
215    /// Sets provider feature-degradation behavior.
216    #[must_use]
217    pub const fn feature_policy(mut self, policy: FeaturePolicy) -> Self {
218        self.agent.config.feature_policy = policy;
219        self
220    }
221
222    /// Sets the desired terminal model-output format.
223    #[must_use]
224    pub fn output_format(mut self, output_format: OutputFormat) -> Self {
225        self.agent.output_format = output_format;
226        self
227    }
228
229    /// Requests strict structured output described by the Rust type `T`.
230    #[must_use]
231    pub fn structured_output<T>(self, name: impl Into<String>) -> Self
232    where
233        T: JsonSchema,
234    {
235        self.output_format(OutputFormat::typed::<T>(name))
236    }
237
238    /// Adds a provider-hosted tool such as Ark web search.
239    #[must_use]
240    pub fn provider_tool(mut self, tool: ProviderToolSpec) -> Self {
241        self.agent.provider_tools.push(tool);
242        self
243    }
244
245    /// Replaces common model generation controls for every Agent turn.
246    #[must_use]
247    pub fn generation(mut self, generation: GenerationOptions) -> Self {
248        self.agent.generation = generation;
249        self
250    }
251
252    /// Sets the sampling temperature for every Agent turn.
253    #[must_use]
254    pub fn temperature(mut self, temperature: f64) -> Self {
255        self.agent.generation.temperature = Some(temperature);
256        self
257    }
258
259    /// Sets nucleus sampling for every Agent turn.
260    #[must_use]
261    pub fn top_p(mut self, top_p: f64) -> Self {
262        self.agent.generation.top_p = Some(top_p);
263        self
264    }
265
266    /// Sets the maximum number of output tokens for every Agent turn.
267    #[must_use]
268    pub fn max_output_tokens(mut self, max_output_tokens: u64) -> Self {
269        self.agent.generation.max_output_tokens = Some(max_output_tokens);
270        self
271    }
272
273    /// Selects streaming or complete response delivery for every Agent turn.
274    #[must_use]
275    pub const fn response_mode(mut self, response_mode: ResponseMode) -> Self {
276        self.agent.response_mode = response_mode;
277        self
278    }
279
280    /// Adds namespaced provider options to every Agent turn.
281    #[must_use]
282    pub fn provider_options(
283        mut self,
284        provider: impl Into<String>,
285        options: serde_json::Value,
286    ) -> Self {
287        self.agent.provider_options.insert(provider.into(), options);
288        self
289    }
290
291    /// Replaces all local Agent configuration.
292    #[must_use]
293    pub const fn config(mut self, config: AgentConfig) -> Self {
294        self.agent.config = config;
295        self
296    }
297
298    /// Shares a write-ahead effect coordinator with this Agent.
299    #[must_use]
300    pub fn effect_executor(mut self, effects: EffectExecutor) -> Self {
301        self.agent.effects = effects;
302        self
303    }
304
305    /// Sets recovery behavior for ambiguous callable effects.
306    #[must_use]
307    pub const fn effect_recovery_policy(mut self, policy: EffectRecoveryPolicy) -> Self {
308        self.agent.effect_recovery = policy;
309        self
310    }
311
312    /// Validates registrations and returns the canonical Agent.
313    ///
314    /// # Errors
315    ///
316    /// Returns [`AgentBuildError`] for blank identity, invalid turn bounds,
317    /// duplicate registrations, or Tool/Agent name collisions.
318    pub fn build(self) -> Result<Agent, AgentBuildError> {
319        if let Some(error) = self.error {
320            return Err(error);
321        }
322        if self.agent.name.trim().is_empty() {
323            return Err(AgentBuildError::EmptyName);
324        }
325        if self.agent.config.max_turns == 0 {
326            return Err(AgentBuildError::ZeroMaxTurns);
327        }
328        if let Some(collision) = self
329            .agent
330            .agents
331            .model_specs()
332            .into_iter()
333            .find(|spec| self.agent.tools.contains(&spec.name))
334        {
335            return Err(AgentBuildError::CallableNameCollision(collision.name));
336        }
337        Ok(self.agent)
338    }
339
340    /// Builds the Agent and runs one ergonomic prompt.
341    ///
342    /// This removes the explicit build step for one-shot usage while retaining
343    /// the complete canonical outcome. Use [`Self::build`] when the Agent will
344    /// be reused or executed with an explicit runtime context.
345    pub fn prompt(
346        self,
347        input: impl Into<String> + Send + 'static,
348    ) -> AgentFuture<'static, Result<AgentOutcome, AgentPromptError>> {
349        let input = input.into();
350        Box::pin(async move {
351            let agent = self.build()?;
352            Ok(agent.prompt(input).await?)
353        })
354    }
355
356    /// Builds the Agent, runs one ergonomic prompt, and returns only
357    /// model-visible text.
358    ///
359    /// This is the shortest path from provider configuration to a text answer.
360    /// Use [`Self::prompt`] when transcript, usage, warnings, and provider
361    /// events must be preserved.
362    pub fn prompt_text(
363        self,
364        input: impl Into<String> + Send + 'static,
365    ) -> AgentFuture<'static, Result<String, AgentPromptError>> {
366        let input = input.into();
367        Box::pin(async move {
368            let agent = self.build()?;
369            Ok(agent.prompt_text(input).await?)
370        })
371    }
372
373    /// Builds an Agent whose schema and local decoder are bound to `T`.
374    ///
375    /// This is the preferred terminal builder operation for structured output
376    /// because a different decode type cannot be selected later by mistake.
377    ///
378    /// # Errors
379    ///
380    /// Returns [`AgentBuildError`] under the same validation rules as
381    /// [`Self::build`].
382    pub fn build_structured<T>(
383        self,
384        name: impl Into<String>,
385    ) -> Result<StructuredAgent<T>, AgentBuildError>
386    where
387        T: JsonSchema,
388    {
389        self.structured_output::<T>(name)
390            .build()
391            .map(StructuredAgent::new)
392    }
393}
394
395impl std::fmt::Debug for AgentBuilder {
396    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
397        formatter
398            .debug_struct("AgentBuilder")
399            .field("agent", &self.agent.name)
400            .field("error", &self.error)
401            .finish_non_exhaustive()
402    }
403}
404
405#[cfg(test)]
406mod tests {
407    use std::{collections::BTreeMap, sync::Arc};
408
409    use runifold_core::{CapabilityId, CapabilitySet, EffectClass, RiskLevel};
410    use runifold_model::{
411        ContentPart, FinishReason, ModelRef, ModelStreamEvent, OutputFormat, ProviderToolSpec,
412        ResponseMode,
413    };
414    use runifold_testkit::ScriptedModel;
415    use runifold_tool::{Tool, ToolContext, ToolDescriptor, ToolError, ToolFuture, ToolOutput};
416    use schemars::JsonSchema;
417    use serde::Deserialize;
418    use serde_json::json;
419
420    use crate::{Agent, AgentBuildError, AgentDescriptor, AgentPromptError};
421
422    struct TestTool {
423        descriptor: ToolDescriptor,
424    }
425
426    #[derive(Deserialize, JsonSchema)]
427    struct TypedAnswer {
428        value: u32,
429    }
430
431    impl TestTool {
432        fn named(name: &str) -> Self {
433            Self {
434                descriptor: ToolDescriptor {
435                    id: CapabilityId::new(),
436                    name: name.into(),
437                    version: "1".into(),
438                    description: "test".into(),
439                    input_schema: json!({"type": "object"}),
440                    output_schema: json!({"type": "object"}),
441                    effect: EffectClass::Pure,
442                    risk: RiskLevel::Low,
443                    metadata: BTreeMap::new(),
444                },
445            }
446        }
447    }
448
449    impl Tool for TestTool {
450        fn descriptor(&self) -> &ToolDescriptor {
451            &self.descriptor
452        }
453
454        fn invoke(
455            &self,
456            input: serde_json::Value,
457            _context: ToolContext,
458        ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
459            Box::pin(async move { Ok(ToolOutput::model_visible(input)) })
460        }
461    }
462
463    #[test]
464    fn fluent_builder_assembles_the_canonical_agent() {
465        let model = Arc::new(ScriptedModel::new());
466        let agent = Agent::builder("worker", model, ModelRef::new("test", "scripted"))
467            .system("Be precise")
468            .tool(TestTool::named("lookup"))
469            .max_turns(4)
470            .build()
471            .unwrap();
472
473        assert_eq!(agent.name, "worker");
474        assert_eq!(agent.instructions.len(), 1);
475        assert!(agent.tools.contains("lookup"));
476        assert_eq!(agent.config.max_turns, 4);
477        assert_eq!(agent.callable_capabilities().len(), 1);
478    }
479
480    #[test]
481    fn builder_retains_generation_provider_and_delivery_controls() {
482        let provider_tool = ProviderToolSpec::new("ark", "web_search").unwrap();
483        let agent = Agent::builder(
484            "researcher",
485            Arc::new(ScriptedModel::new()),
486            ModelRef::new("ark", "doubao"),
487        )
488        .temperature(0.2)
489        .top_p(0.8)
490        .max_output_tokens(4_096)
491        .response_mode(ResponseMode::Complete)
492        .provider_tool(provider_tool)
493        .provider_options("ark", json!({"thinking": {"type": "enabled"}}))
494        .build()
495        .unwrap();
496
497        assert_eq!(agent.generation.temperature, Some(0.2));
498        assert_eq!(agent.generation.top_p, Some(0.8));
499        assert_eq!(agent.generation.max_output_tokens, Some(4_096));
500        assert_eq!(agent.response_mode, ResponseMode::Complete);
501        assert_eq!(agent.provider_tools[0].tool_type, "web_search");
502        assert_eq!(agent.provider_options["ark"]["thinking"]["type"], "enabled");
503    }
504
505    #[test]
506    fn build_rejects_tool_and_agent_name_collisions() {
507        let model = Arc::new(ScriptedModel::new());
508        let child = Arc::new(Agent::new(
509            "child",
510            model.clone(),
511            ModelRef::new("test", "child"),
512        ));
513        let error = Agent::builder("parent", model, ModelRef::new("test", "parent"))
514            .tool(TestTool::named("search"))
515            .child(
516                AgentDescriptor::new("search", "delegate search"),
517                child,
518                CapabilitySet::new(),
519            )
520            .build()
521            .unwrap_err();
522
523        assert_eq!(
524            error,
525            AgentBuildError::CallableNameCollision("search".into())
526        );
527    }
528
529    #[test]
530    fn builder_derives_a_strict_output_schema_from_a_rust_type() {
531        let example = TypedAnswer { value: 7 };
532        assert_eq!(example.value, 7);
533        let agent = Agent::builder(
534            "worker",
535            Arc::new(ScriptedModel::new()),
536            ModelRef::new("test", "scripted"),
537        )
538        .structured_output::<TypedAnswer>("typed_answer")
539        .build()
540        .unwrap();
541
542        let OutputFormat::JsonSchema {
543            name,
544            schema,
545            strict,
546        } = agent.output_format
547        else {
548            panic!("expected JSON-schema output");
549        };
550        assert_eq!(name, "typed_answer");
551        assert!(strict);
552        assert_eq!(schema["properties"]["value"]["type"], "integer");
553    }
554
555    #[test]
556    fn builder_prompt_text_is_a_single_use_golden_path() {
557        let model = ScriptedModel::new();
558        model.enqueue([
559            ModelStreamEvent::ResponseStarted {
560                id: Some("response-1".into()),
561                model: ModelRef::new("test", "scripted"),
562            },
563            ModelStreamEvent::ContentPartCompleted {
564                index: 0,
565                part: ContentPart::text("done"),
566            },
567            ModelStreamEvent::ResponseCompleted {
568                finish_reason: FinishReason::Stop,
569                provider_metadata: BTreeMap::new(),
570            },
571        ]);
572
573        let text = futures_executor::block_on(
574            Agent::builder("worker", Arc::new(model), ModelRef::new("test", "scripted"))
575                .system("Be precise")
576                .prompt_text("start"),
577        )
578        .unwrap();
579
580        assert_eq!(text, "done");
581    }
582
583    #[test]
584    fn builder_prompt_reports_build_failures_before_model_execution() {
585        let error = futures_executor::block_on(
586            Agent::builder(
587                "",
588                Arc::new(ScriptedModel::new()),
589                ModelRef::new("test", "scripted"),
590            )
591            .prompt("start"),
592        )
593        .unwrap_err();
594
595        assert!(matches!(
596            error,
597            AgentPromptError::Build(AgentBuildError::EmptyName)
598        ));
599    }
600}