Skip to main content

oxicode_sdk/
agent_builder.rs

1//! AgentBuilder — Fluent API for creating agents
2
3use std::path::PathBuf;
4use std::sync::Arc;
5
6use oxicode_agent::{
7    Agent, AgentConfig, AgentTool, AgentToolResult, ProviderResolver, ToolContext, ToolRegistry,
8};
9
10use crate::builder::Oxicode;
11use crate::middleware::{Middleware, MiddlewarePipeline};
12use crate::observability::{AuditLog, CostTracker, Tracer};
13use crate::security::{Authorizer, CapabilitySet};
14
15/// Closures owned by the cli (or any product) that participate in the
16/// agent hook chain. Passed into [`AgentBuilder::with_session_hooks`] so
17/// they are composed into the same `AgentHooks` that the middleware
18/// pipeline produces. The single-`set_hooks` invariant (only
19/// [`AgentBuilder::build`] calls `set_hooks`) is what keeps the
20/// before/after_tool_call slots alive across the cli session boot.
21pub struct SessionHookClosures {
22    /// Stop signal consulted at the end of every turn.
23    pub should_stop_after_turn:
24        std::sync::Arc<dyn Fn(&oxicode_agent::ShouldStopAfterTurnContext) -> bool + Send + Sync>,
25    /// Drain the steering queue on demand.
26    pub get_steering_messages: std::sync::Arc<dyn Fn() -> Vec<oxicode_ai::Message> + Send + Sync>,
27    /// Drain the follow-up queue on demand.
28    pub get_follow_up_messages: std::sync::Arc<dyn Fn() -> Vec<oxicode_ai::Message> + Send + Sync>,
29    /// Tool-execution mode (Sequential is the cli default).
30    pub tool_execution: oxicode_agent::ToolExecutionMode,
31}
32
33/// Builder for creating an agent with custom configuration.
34#[allow(dead_code)]
35pub struct AgentBuilder<'a> {
36    oxicode: &'a Oxicode,
37    config: AgentConfig,
38    tools: ToolRegistry,
39    workspace_dir: Option<PathBuf>,
40    system_prompt: Option<String>,
41    // ── Security ──
42    capabilities: Option<CapabilitySet>,
43    authorizer: Option<Arc<Authorizer>>,
44    // ── Observability ──
45    tracer: Option<Arc<Tracer>>,
46    audit_log: Option<Arc<AuditLog>>,
47    cost_tracker: Option<Arc<CostTracker>>,
48    // ── Middleware ──
49    middlewares: Vec<Arc<dyn Middleware>>,
50    // ── Hooks (port 16) ──
51    hooks_middleware: Option<crate::middleware::HookMiddleware>,
52    // ── Session-level closures (cli-owned stop flag + queues) ──
53    session_hooks: Option<SessionHookClosures>,
54    // ── Compaction ──
55    /// Custom compactor that replaces the default LLM compactor in every
56    /// agent run (see `AgentLoopConfig::compactor` for replace semantics).
57    compactor: Option<Arc<dyn oxicode_ai::Compactor>>,
58}
59
60impl<'a> AgentBuilder<'a> {
61    /// Create a new builder bound to the given [`Oxicode`] instance with the provided agent config.
62    pub fn new(oxicode: &'a Oxicode, config: AgentConfig) -> Self {
63        Self {
64            oxicode,
65            config,
66            tools: ToolRegistry::new(),
67            workspace_dir: None,
68            system_prompt: None,
69            capabilities: None,
70            authorizer: None,
71            tracer: None,
72            audit_log: None,
73            cost_tracker: None,
74            middlewares: Vec::new(),
75            hooks_middleware: None,
76            session_hooks: None,
77            compactor: None,
78        }
79    }
80
81    /// Set the working directory for file tools.
82    pub fn workspace(mut self, dir: impl Into<PathBuf>) -> Self {
83        self.workspace_dir = Some(dir.into());
84        self
85    }
86
87    /// Set a custom system prompt.
88    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
89        self.system_prompt = Some(prompt.into());
90        self
91    }
92
93    /// Set the agent's autonomy [`Mode`](oxicode_agent::Mode).
94    ///
95    /// In [`Mode::Auto`](oxicode_agent::Mode::Auto) the agent runs to
96    /// completion without asking the user questions (the `ask` tool is
97    /// short-circuited). Default: [`Mode::Default`](oxicode_agent::Mode::Default).
98    pub fn with_mode(mut self, mode: oxicode_agent::Mode) -> Self {
99        self.config.mode = mode;
100        self
101    }
102    /// Register a [`TodoStateProvider`](crate::TodoStateProvider) so the agent's `todo` tool works.
103    ///
104    /// The provider is shared between the agent (writer) and the host
105    /// application (reader), so you can observe phase changes in real time
106    /// by calling [`TodoStateProvider::get_phases()`](crate::TodoStateProvider::get_phases) periodically.
107    ///
108    /// Use [`InMemoryTodoState`](crate::inmem::InMemoryTodoState) for a
109    /// ready-to-go in-memory implementation:
110    ///
111    /// ```no_run
112    /// use std::sync::Arc;
113    /// use oxicode_sdk::{AgentConfig, OxicodeBuilder, inmem::InMemoryTodoState};
114    ///
115    /// let todo = Arc::new(InMemoryTodoState::new());
116    /// let oxicode = OxicodeBuilder::new().with_builtins().build();
117    /// let agent = oxicode.agent(AgentConfig {
118    ///     model_id: "anthropic/claude-sonnet-4-20250514".into(),
119    ///     ..Default::default()
120    /// })
121    /// .with_todo(todo.clone())
122    /// .build()
123    /// .unwrap();
124    ///
125    /// // Observe later:
126    /// let phases = todo.get_phases();
127    /// ```
128    pub fn with_todo(
129        mut self,
130        todo: std::sync::Arc<dyn oxicode_agent::tools::TodoStateProvider>,
131    ) -> Self {
132        self.config.todo = Some(todo);
133        self.tools.register(oxicode_agent::tools::todo::TodoTool);
134        self
135    }
136
137    /// Register a [`MemoryBackend`](oxicode_agent::tools::MemoryBackend) and the
138    /// four `memory_*` tools (`memory_recall`, `memory_reflect`,
139    /// `memory_retain`, `memory_edit`).
140    ///
141    /// Generic entry point: pass any `MemoryBackend`. For the common case of
142    /// bridging the engine's registered `MemoryStore` port, use
143    /// [`Self::with_port_memory`] instead.
144    pub fn with_memory_backend(
145        mut self,
146        backend: std::sync::Arc<dyn oxicode_agent::tools::MemoryBackend>,
147    ) -> Self {
148        self.config.memory = Some(backend);
149        self.tools.register(oxicode_agent::tools::MemoryRecallTool);
150        self.tools.register(oxicode_agent::tools::MemoryReflectTool);
151        self.tools.register(oxicode_agent::tools::MemoryRetainTool);
152        self.tools.register(oxicode_agent::tools::MemoryEditTool);
153        self
154    }
155
156    /// Bridge the engine's registered `MemoryStore` (+ `EmbeddingProvider`)
157    /// ports into this agent's `memory_*` tools via `PortMemoryBackend`.
158    ///
159    /// This is how a pure-SDK consumer makes memory functional end-to-end:
160    /// register the ports on [`OxicodeBuilder`](crate::OxicodeBuilder)
161    /// (`with_memory` / `with_embeddings`), then call this on the agent.
162    /// Without an `EmbeddingProvider`, `put` / `list` / `delete` work but
163    /// semantic `search` returns an error.
164    ///
165    /// Without this call (or [`Self::with_memory_backend`]), the
166    /// `memory_*` tools are absent and `ToolContext.memory` stays `None` —
167    /// the registered `MemoryStore` port is unused by the agent loop.
168    ///
169    /// ```no_run
170    /// use std::sync::Arc;
171    /// use oxicode_sdk::{OxicodeBuilder, inmem::InMemoryMemoryStore};
172    ///
173    /// let oxicode = OxicodeBuilder::new()
174    ///     .with_builtins()
175    ///     .with_memory(Arc::new(InMemoryMemoryStore::new()))
176    ///     .build();
177    /// let agent = oxicode.agent(oxicode_agent::AgentConfig {
178    ///     model_id: "anthropic/claude-sonnet-4-20250514".into(),
179    ///     ..Default::default()
180    /// })
181    /// .with_port_memory()
182    /// .build()
183    /// .unwrap();
184    /// ```
185    pub fn with_port_memory(self) -> Self {
186        let ports = self.oxicode.ports().clone();
187        let backend = crate::port_memory_backend::PortMemoryBackend::from_ports(
188            ports.memory.clone(),
189            Some(ports.embeddings.clone()),
190        );
191        self.with_memory_backend(std::sync::Arc::new(backend))
192    }
193
194    /// Set the URL resolver — enables internal-URL dispatch (`issue://`,
195    /// `skill://`, `memory://`, …) in the `read`/`grep`/`find` tools.
196    pub fn with_url_resolver(
197        mut self,
198        resolver: std::sync::Arc<dyn oxicode_agent::tools::UrlResolver>,
199    ) -> Self {
200        self.config.url_resolver = Some(resolver);
201        self
202    }
203
204    /// Bridge the engine's registered `InternalUrlRouter` port into this
205    /// agent's `read`/`grep`/`find` tools via [`SdkUrlResolver`](crate::url_resolver::SdkUrlResolver).
206    ///
207    /// This is how a pure-SDK consumer enables protocol-scheme URL
208    /// resolution: register scheme handlers on the router port, then call
209    /// this. Without it (or [`Self::with_url_resolver`]), URL-prefixed
210    /// paths are treated as regular file paths.
211    pub fn with_port_url_resolver(self) -> Self {
212        let resolver = std::sync::Arc::new(crate::url_resolver::SdkUrlResolver::new(
213            self.oxicode.ports().url_router.clone(),
214        ));
215        self.with_url_resolver(resolver)
216    }
217
218    /// Set the hashline snapshot store — enables line-anchored edit mode
219    /// (`read` emits `[path#TAG]` headers, `edit` validates against them).
220    ///
221    /// Without this, the `edit` tool falls back to plain text replacement.
222    /// Use [`oxicode_hashline::InMemorySnapshotStore`] for an ephemeral store, or
223    /// implement [`oxicode_hashline::SnapshotStore`] for persistence.
224    pub fn with_snapshot_store(
225        mut self,
226        store: std::sync::Arc<dyn oxicode_hashline::SnapshotStore>,
227    ) -> Self {
228        self.config.snapshot_store = Some(store);
229        self
230    }
231
232    /// Bridge the engine into an in-process subagent runner and register the
233    /// `subagent` tool.
234    ///
235    /// Uses [`SdkSubagentRunner`](crate::delegation::SdkSubagentRunner) so the `subagent`
236    /// tool runs isolated agents in-process (no CLI binary). Without this, the
237    /// `subagent` tool is absent from the agent's toolset.
238    pub fn with_port_subagent(mut self) -> Self {
239        let runner = std::sync::Arc::new(crate::delegation::SdkSubagentRunner::new(
240            self.oxicode.clone(),
241        ));
242        self.config.subagent_runner = Some(runner);
243        self.tools
244            .register(oxicode_agent::tools::SubagentTool::new());
245        self
246    }
247
248    /// Add the [`HookMiddleware`](crate::middleware::HookMiddleware) backed by the engine's registered
249    /// `HookRunner` port (see [`crate::OxicodeBuilder::with_hooks`]).
250    ///
251    /// When the port is `NoopHookRunner` (the default), this is a no-op.
252    /// The middleware composes into the existing pipeline at the
253    /// `audit → authorizer → hooks → user` position. `set_hooks` is called
254    /// exactly once in `build()` — see the single-`set_hooks` invariant.
255    pub fn with_port_hooks(mut self) -> Self {
256        let runner = std::sync::Arc::clone(&self.oxicode.ports().hooks);
257        self.hooks_middleware = Some(crate::middleware::HookMiddleware::new(runner));
258        self
259    }
260
261    /// Install session-level closures (stop flag + steering/follow_up
262    /// queues). These are composed into the same `AgentHooks` that the
263    /// middleware pipeline produces, so `set_hooks` is called exactly once.
264    /// This is the **only** way to install session hooks — never call
265    /// `agent.set_hooks(...)` elsewhere (it would wipe the middleware
266    /// pipeline's before/after_tool_call slots).
267    pub fn with_session_hooks(mut self, closures: SessionHookClosures) -> Self {
268        self.session_hooks = Some(closures);
269        self
270    }
271
272    /// Replace the default LLM compactor with a custom one.
273    ///
274    /// The compactor is threaded into every agent run (via
275    /// `AgentLoopConfig::compactor`) and replaces the default
276    /// `LlmCompactor` — the `CompactionManager` has a single compactor
277    /// slot. `None` (default) preserves the existing LLM-compactor
278    /// behavior.
279    ///
280    /// The SDK ships [`crate::snapcompact_compactor::SnapcompactCompactor`] — a PNG-frame
281    /// compactor that makes no LLM call:
282    ///
283    /// ```
284    /// # use oxicode_sdk::{AgentBuilder, snapcompact_compactor::SnapcompactCompactor};
285    /// # fn build(oxicode: &oxicode_sdk::Oxicode, config: oxicode_agent::AgentConfig)
286    /// #     -> anyhow::Result<oxicode_agent::Agent> {
287    /// AgentBuilder::new(oxicode, config)
288    ///     .with_compactor(std::sync::Arc::new(SnapcompactCompactor::new()))
289    ///     .build()
290    /// # }
291    /// ```
292    pub fn with_compactor(mut self, compactor: std::sync::Arc<dyn oxicode_ai::Compactor>) -> Self {
293        self.compactor = Some(compactor);
294        self
295    }
296
297    /// Register the standard coding tools (read, write, edit, bash, grep, find, ls, ...).
298    pub fn coding_tools(self) -> Self {
299        let cwd = self
300            .workspace_dir
301            .clone()
302            .unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
303        let tools = crate::tool_factory::coding_tools(&cwd);
304        for name in tools.names() {
305            if let Some(tool) = tools.get(&name) {
306                self.tools.register_arc(tool);
307            }
308        }
309        self
310    }
311
312    /// Register read-only tools (read, ls).
313    pub fn readonly_tools(self) -> Self {
314        let cwd = self
315            .workspace_dir
316            .clone()
317            .unwrap_or_else(|| std::env::current_dir().unwrap_or_else(|_| PathBuf::from(".")));
318        let tools = crate::tool_factory::readonly_tools(&cwd);
319        for name in tools.names() {
320            if let Some(tool) = tools.get(&name) {
321                self.tools.register_arc(tool);
322            }
323        }
324        self
325    }
326
327    /// Register a tool.
328    pub fn tool(self, tool: impl AgentTool + 'static) -> Self {
329        self.tools.register(tool);
330        self
331    }
332
333    /// Register a custom tool from a closure (synchronous handler).
334    ///
335    /// Creates a `ClosureTool` internally.
336    ///
337    /// # Example
338    /// ```rust
339    /// use oxicode_sdk::{ClosureTool, AgentToolResult};
340    ///
341    /// // custom_tool creates a tool from a closure
342    /// let tool = ClosureTool::new_sync(
343    ///     "memory_recall",
344    ///     "Search long-term memory",
345    ///     serde_json::json!({"type": "object", "properties": {"query": {"type": "string"}}}),
346    ///     |params, _ctx| {
347    ///         let query = params["query"].as_str().unwrap();
348    ///         Ok(AgentToolResult::success(format!("Recalled: {}", query)))
349    ///     },
350    /// );
351    /// ```
352    pub fn custom_tool(
353        self,
354        name: impl Into<String>,
355        description: impl Into<String>,
356        schema: serde_json::Value,
357        handler: impl Fn(
358            serde_json::Value,
359            &ToolContext,
360        ) -> Result<AgentToolResult, oxicode_agent::ToolError>
361        + Send
362        + Sync
363        + 'static,
364    ) -> Self {
365        self.tool(crate::closure_tool::ClosureTool::new_sync(
366            name,
367            description,
368            schema,
369            handler,
370        ))
371    }
372
373    /// Register multiple tools.
374    pub fn tools(self, tools: impl IntoIterator<Item = impl AgentTool + 'static>) -> Self {
375        for tool in tools {
376            self.tools.register(tool);
377        }
378        self
379    }
380
381    /// Register kernel tools from a [`KernelToolProvider`].
382    ///
383    /// This is the bridge for oxios kernel tools (exec, memory, browser, etc.).
384    /// The kernel implements `KernelToolProvider` and registers its tools
385    /// into the agent's tool registry.
386    ///
387    /// [`KernelToolProvider`]: crate::KernelToolProvider
388    pub fn kernel_tools(
389        self,
390        provider: &dyn crate::KernelToolProvider,
391        context: &crate::KernelToolContext,
392    ) -> Self {
393        provider.register_tools(&self.tools, context);
394        self
395    }
396
397    // ── Security ──────────────────────────────────────────
398
399    /// Set the capability set for this agent.
400    pub fn capabilities(mut self, caps: CapabilitySet) -> Self {
401        self.capabilities = Some(caps);
402        self
403    }
404
405    /// Use standard coding capabilities.
406    pub fn coding_capabilities(self) -> Self {
407        let ws = self
408            .workspace_dir
409            .clone()
410            .unwrap_or_else(|| PathBuf::from("."));
411        self.capabilities(CapabilitySet::coding(ws.to_str().unwrap_or(".")))
412    }
413
414    /// Use read-only capabilities.
415    pub fn readonly_capabilities(self) -> Self {
416        let ws = self
417            .workspace_dir
418            .clone()
419            .unwrap_or_else(|| PathBuf::from("."));
420        self.capabilities(CapabilitySet::read_only(ws.to_str().unwrap_or(".")))
421    }
422
423    /// Attach an authorizer for capability enforcement.
424    pub fn authorizer(mut self, authorizer: Arc<Authorizer>) -> Self {
425        self.authorizer = Some(authorizer);
426        self
427    }
428
429    // ── Observability ──────────────────────────────────────
430
431    /// Attach a tracer for distributed tracing.
432    pub fn tracer(mut self, tracer: Arc<Tracer>) -> Self {
433        self.tracer = Some(tracer);
434        self
435    }
436
437    /// Attach an audit log for security and tool audit trail.
438    pub fn audit_log(mut self, audit: Arc<AuditLog>) -> Self {
439        self.audit_log = Some(audit);
440        self
441    }
442
443    /// Attach a cost tracker for token and cost monitoring.
444    pub fn cost_tracker(mut self, tracker: Arc<CostTracker>) -> Self {
445        self.cost_tracker = Some(tracker);
446        self
447    }
448
449    // ── Middleware ─────────────────────────────────────────
450
451    /// Add a middleware to the pipeline.
452    pub fn middleware(mut self, mw: impl Middleware + 'static) -> Self {
453        self.middlewares.push(Arc::new(mw));
454        self
455    }
456
457    /// Add a rate limit middleware (convenience shortcut).
458    pub fn with_rate_limit(self, max_per_minute: usize) -> Self {
459        self.middleware(crate::middleware::RateLimitMiddleware::new(max_per_minute))
460    }
461
462    /// Add a token budget middleware (convenience shortcut).
463    pub fn with_token_budget(self, max_tokens: usize) -> Self {
464        self.middleware(crate::middleware::TokenBudgetMiddleware::new(max_tokens))
465    }
466
467    /// Add a logging middleware (convenience shortcut).
468    pub fn with_logging(self) -> Self {
469        self.middleware(crate::middleware::LoggingMiddleware::new(
470            tracing::Level::INFO,
471        ))
472    }
473
474    /// Build the agent.
475    ///
476    /// Uses the Oxicode engine's `ProviderResolver` for isolated provider/model
477    /// lookups, so `switch_model()` and compaction stay within the engine's
478    /// registry — no global state pollution.
479    pub fn build(mut self) -> anyhow::Result<Agent> {
480        // 1. Resolve model from Oxicode's instance registry
481        let model = self.oxicode.resolve_model(&self.config.model_id)?;
482
483        // 2. Create provider via Oxicode's engine (custom → built-in fallback)
484        let provider: Arc<dyn oxicode_ai::Provider> =
485            self.oxicode.create_provider(&model.provider)?;
486
487        // 3. Merge workspace_dir into config
488        let mut config = self.config.clone();
489        config.workspace_dir = self.workspace_dir.or(config.workspace_dir);
490        if let Some(ref prompt) = self.system_prompt {
491            config.system_prompt = Some(prompt.clone());
492        }
493
494        // 4. Use Oxicode directly as the resolver. Oxicode implements ProviderResolver
495        //    with catalog→static model resolution (builder.rs:106-123) and
496        //    credential-aware provider creation (builder.rs:131-148).
497        //
498        //    The previous hand-rolled OxicodeResolver/OxicodeCore closure only
499        //    consulted the static ModelRegistry via lookup(), silently
500        //    dropping the catalog port — so catalog-known models (e.g.
501        //    newer Z.AI models from models.dev) resolved at build() time
502        //    via Oxicode::resolve_model but failed inside the agent loop's
503        //    own resolve_model(), causing "Failed to resolve model" at
504        //    stream time.
505        let resolver: Arc<dyn ProviderResolver> = Arc::new(self.oxicode.clone());
506
507        // 4b. Capability gate: drop the `lsp` tool when no `LspProvider`
508        //     is configured on the agent config. This avoids the
509        //     "LSP not configured" runtime error path entirely — the
510        //     tool simply isn't visible to the model when LSP is off.
511        //     See docs/designs/2026-07-18-stub-completion.md §4.3.
512        if config.lsp.is_none() {
513            self.tools.unregister("lsp");
514        }
515
516        // 5. Create agent with the isolated resolver and optional custom
517        //    compactor (replaces the default LLM compactor when set).
518        let agent = Agent::new_with_compactor(
519            provider,
520            config,
521            Arc::new(self.tools),
522            resolver,
523            self.compactor,
524        );
525
526        // 6. Authorizer: grant capabilities.
527        //
528        // The authorizer middleware (`AuthorizerMiddleware`) checks
529        // `Capability::ToolUse { tool_name }` against the granted
530        // capabilities — type-specific, no cross-variant
531        // implication. Without a `ToolUse` grant, every tool
532        // call would be denied by the middleware regardless of
533        // whether the agent has fine-grained FileRead/Bash caps.
534        //
535        // Coarse-grant fallback: when the granted capability set
536        // contains no `ToolUse` variant, auto-add a wildcard
537        // `ToolUse { tool_name: "*" }`. This makes the SDK's
538        // authorizer integration usable out of the box with
539        // `CapabilitySet::coding()` / `read_only()` / `research()` /
540        // `browser()` (none of which contain `ToolUse`).
541        //
542        // Fine-grained enforcement (command/path restrictions)
543        // would require tool-specific arg parsing inside the
544        // middleware to derive `Bash`/`FileRead` capabilities
545        // from the call's JSON args. That's a follow-up; see
546        // design doc at
547        // docs/designs/2026-06-30-observability-wiring.md.
548        if let Some(authorizer) = &self.authorizer {
549            let agent_id = resolved_agent_id(&agent);
550            if let Some(mut caps) = self.capabilities.clone() {
551                let has_tool_use = caps
552                    .capabilities()
553                    .iter()
554                    .any(|c| matches!(c, crate::security::Capability::ToolUse { .. }));
555                if !has_tool_use {
556                    caps.add(crate::security::Capability::ToolUse {
557                        tool_name: "*".into(),
558                    });
559                }
560                let subject = crate::security::CapabilitySubject::Agent(agent_id);
561                authorizer.grant(subject, caps);
562            }
563        }
564
565        // 7. Build a single unified middleware pipeline that includes
566        //    user middlewares, the audit-log adapter, and the
567        //    authorizer adapter. Order matters: audit fires FIRST
568        //    (records all attempts), authorizer fires SECOND (denies
569        //    if needed — short-circuits before user mws run), user
570        //    middlewares fire LAST.
571        //
572        //    The pipeline is wrapped into AgentHooks via
573        //    `build_hooks` once, so `set_hooks()` is called exactly
574        //    once. This avoids the replace-semantics bug class
575        //    documented in docs/audits/2026-06-30-sdk-coverage.md
576        //    Gap-0 ("observability silently overwritten when composes
577        //    with user middlewares"). HookMiddleware slots and
578        //    session-level closures (via with_session_hooks) are composed
579        //    into the SAME AgentHooks instance — set_hooks remains the
580        //    single call site.
581        let has_observability_mws = self.audit_log.is_some() || self.authorizer.is_some();
582        let has_user_mws = !self.middlewares.is_empty();
583        let has_hooks = self.hooks_middleware.is_some();
584        let has_session_hooks = self.session_hooks.is_some();
585        if has_user_mws || has_observability_mws || has_hooks || has_session_hooks {
586            let agent_id = resolved_agent_id(&agent);
587            let mut pipeline = MiddlewarePipeline::new();
588
589            // Audit fires first so every attempt (allowed or denied) is logged.
590            if let Some(audit) = &self.audit_log {
591                pipeline = pipeline.add_arc(Arc::new(
592                    crate::middleware::observability_adapters::AuditLogMiddleware::new(
593                        Arc::clone(audit),
594                        agent_id.clone(),
595                    ),
596                ));
597            }
598
599            // Authorizer fires second — its denial short-circuits the
600            // pipeline via `MiddlewareAction::Block`, which the
601            // existing bridge maps to `BeforeToolCallResult { block: true }`.
602            if let Some(authorizer) = &self.authorizer {
603                let mut mw = crate::middleware::observability_adapters::AuthorizerMiddleware::new(
604                    Arc::clone(authorizer),
605                    agent_id.clone(),
606                );
607                if let Some(audit) = &self.audit_log {
608                    mw = mw.with_audit(Arc::clone(audit));
609                }
610                pipeline = pipeline.add_arc(Arc::new(mw));
611            }
612
613            // HookMiddleware fires AFTER authorizer (so authorizer denials
614            // still short-circuit) and BEFORE user middlewares (so user
615            // middlewares observe hook-driven blocks).
616            if let Some(hooks_mw) = self.hooks_middleware.take() {
617                pipeline = pipeline.add_arc(Arc::new(hooks_mw));
618            }
619
620            // User middlewares fire last so audit/auth observe their
621            // calls and Authorizer denials short-circuit before them.
622            for mw in self.middlewares.into_iter() {
623                pipeline = pipeline.add_arc(mw);
624            }
625
626            let pipeline = Arc::new(pipeline);
627            let terminate_flag = Arc::new(std::sync::atomic::AtomicBool::new(false));
628            let mut hooks = crate::middleware::build_hooks(pipeline, agent_id, terminate_flag);
629
630            // Session-level closures — overwrite the three slots the cli
631            // owns (should_stop_after_turn, steering, follow_up) on the
632            // SAME AgentHooks the pipeline just produced. before_tool_call
633            // and after_tool_call are preserved. This keeps set_hooks as a
634            // single call site and avoids the replace-semantics bug class
635            // (audit Gap-0).
636            if let Some(session) = self.session_hooks.take() {
637                hooks.should_stop_after_turn = Some(session.should_stop_after_turn);
638                hooks.get_steering_messages = Some(session.get_steering_messages);
639                hooks.get_follow_up_messages = Some(session.get_follow_up_messages);
640                hooks.tool_execution = session.tool_execution;
641            }
642
643            // SINGLE set_hooks call for the entire agent.
644            agent.set_hooks(hooks);
645        }
646
647        // 8. Tracer and CostTracker → event-tap path (accumulate, not replace).
648        if self.tracer.is_some() || self.cost_tracker.is_some() {
649            install_observability_dispatch(&agent, self.tracer.clone(), self.cost_tracker.clone());
650        }
651
652        Ok(agent)
653    }
654}
655
656/// Synthesize a stable agent id used as the principal in capability
657/// grants, audit-log entries, and observability dispatch. Matches the
658/// existing behavior at agent_builder.rs:443-447 (synthesize a UUID
659/// only when the config name is empty).
660fn resolved_agent_id(agent: &Agent) -> String {
661    let cfg = agent.get_config();
662    if cfg.name.is_empty() {
663        uuid::Uuid::new_v4().to_string()
664    } else {
665        cfg.name
666    }
667}
668
669/// Build the event-tap closure that records short lifecycle spans and drives
670/// `CostTracker` from the agent's emitted events.
671fn install_observability_dispatch(
672    agent: &Agent,
673    tracer: Option<Arc<Tracer>>,
674    cost_tracker: Option<Arc<crate::observability::CostTracker>>,
675) {
676    use crate::observability::{SpanKind, TokenUsage};
677    use oxicode_agent::AgentEvent;
678    if tracer.is_none() && cost_tracker.is_none() {
679        return;
680    }
681    // Use the same resolved agent_id as the middleware path so
682    // AuditLog / Authorizer / CostTracker observations all key
683    // by the same principal. Without this, a user-supplied
684    // AgentConfig with `name: ""` would create a divergence:
685    // `resolved_agent_id` falls back to a UUID for the
686    // middleware grants, but `agent.get_config().name`
687    // is the empty string — CostTracker would record under
688    // `""` while Authorizer grants under the UUID.
689    let agent_id = resolved_agent_id(agent);
690    let resolver = agent.resolver().clone();
691    let model_id = agent.get_config().model_id;
692    agent.add_observability_dispatch(move |event: AgentEvent| match event {
693        AgentEvent::AgentStart {
694            prompts,
695            session_id,
696        } => {
697            if let Some(tracer) = &tracer {
698                let mut span = tracer.start("run", SpanKind::Agent);
699                span.set_attribute("agent.id", serde_json::json!(agent_id));
700                span.set_attribute("model.id", serde_json::json!(model_id));
701                span.set_attribute("prompt.count", serde_json::json!(prompts.len()));
702                if let Some(session_id) = session_id {
703                    span.set_attribute("session.id", serde_json::json!(session_id));
704                }
705            }
706        }
707        AgentEvent::TurnStart { turn_number } => {
708            if let Some(tracer) = &tracer {
709                let mut span = tracer.start("turn_start", SpanKind::Agent);
710                span.set_attribute("turn.number", serde_json::json!(turn_number));
711            }
712        }
713        AgentEvent::TurnEnd {
714            turn_number,
715            tool_results,
716            ..
717        } => {
718            if let Some(tracer) = &tracer {
719                let mut span = tracer.start("turn_end", SpanKind::Agent);
720                span.set_attribute("turn.number", serde_json::json!(turn_number));
721                span.set_attribute("tool.result.count", serde_json::json!(tool_results.len()));
722            }
723        }
724        AgentEvent::ToolExecutionStart {
725            tool_call_id,
726            tool_name,
727            ..
728        } => {
729            if let Some(tracer) = &tracer {
730                let mut span = tracer.start("tool_start", SpanKind::Tool);
731                span.set_attribute("tool.call.id", serde_json::json!(tool_call_id));
732                span.set_attribute("tool.name", serde_json::json!(tool_name));
733            }
734        }
735        AgentEvent::ToolExecutionEnd {
736            tool_call_id,
737            tool_name,
738            is_error,
739            ..
740        } => {
741            if let Some(tracer) = &tracer {
742                let mut span = tracer.start("tool_end", SpanKind::Tool);
743                span.set_attribute("tool.call.id", serde_json::json!(tool_call_id));
744                span.set_attribute("tool.name", serde_json::json!(tool_name));
745                span.set_attribute("error", serde_json::json!(is_error));
746                if is_error {
747                    span.set_error("tool execution failed");
748                }
749            }
750        }
751        AgentEvent::Usage {
752            input_tokens,
753            output_tokens,
754        } => {
755            let Some(cost_tracker) = &cost_tracker else {
756                return;
757            };
758            if let Some(model) = resolver.resolve_model(&model_id) {
759                cost_tracker.record(
760                    &agent_id,
761                    &model,
762                    TokenUsage {
763                        input: input_tokens as u64,
764                        output: output_tokens as u64,
765                        cache_read: 0,
766                        cache_write: 0,
767                    },
768                );
769            }
770        }
771        _ => {}
772    });
773}
774
775#[cfg(test)]
776mod tests {
777    use super::*;
778    use crate::ports::catalog::{
779        CatalogEvent, CatalogModelEntry, CatalogProtocol, CatalogSource, ModelCatalog,
780    };
781    use crate::{OxicodeBuilder, SdkResult};
782    use std::future::Future;
783    use std::pin::Pin;
784    use tokio::sync::broadcast;
785
786    /// Minimal catalog with a single model that exists ONLY in the catalog
787    /// port — not in the static ModelRegistry. This reproduces the desync
788    /// where `Oxicode::resolve_model()` (catalog→static) finds the model but
789    /// the old `OxicodeResolver`'s `lookup()` (static-only) did not.
790    struct SingleModelCatalog {
791        entry: CatalogModelEntry,
792        tx: broadcast::Sender<CatalogEvent>,
793    }
794
795    impl std::fmt::Debug for SingleModelCatalog {
796        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
797            f.debug_struct("SingleModelCatalog").finish_non_exhaustive()
798        }
799    }
800
801    impl ModelCatalog for SingleModelCatalog {
802        fn list_providers(
803            &self,
804        ) -> Pin<Box<dyn Future<Output = SdkResult<Vec<String>>> + Send + '_>> {
805            let p = self.entry.provider.clone();
806            Box::pin(async move { Ok(vec![p]) })
807        }
808        fn get_provider(
809            &self,
810            _: &str,
811        ) -> Pin<
812            Box<
813                dyn Future<Output = SdkResult<Option<crate::ports::catalog::CatalogProviderEntry>>>
814                    + Send
815                    + '_,
816            >,
817        > {
818            Box::pin(async { Ok(None) })
819        }
820        fn list_models(
821            &self,
822            _: &str,
823        ) -> Pin<Box<dyn Future<Output = SdkResult<Vec<CatalogModelEntry>>> + Send + '_>> {
824            let e = self.entry.clone();
825            Box::pin(async move { Ok(vec![e]) })
826        }
827        fn get_model(
828            &self,
829            provider: &str,
830            model_id: &str,
831        ) -> Pin<Box<dyn Future<Output = SdkResult<Option<CatalogModelEntry>>> + Send + '_>>
832        {
833            let hit = self.get_model_sync(provider, model_id);
834            Box::pin(async move { Ok(hit) })
835        }
836        fn search(
837            &self,
838            _: &str,
839        ) -> Pin<Box<dyn Future<Output = SdkResult<Vec<CatalogModelEntry>>> + Send + '_>> {
840            let e = self.entry.clone();
841            Box::pin(async move { Ok(vec![e]) })
842        }
843        fn model_count(&self) -> Pin<Box<dyn Future<Output = SdkResult<usize>> + Send + '_>> {
844            Box::pin(async { Ok(1) })
845        }
846        fn refresh(
847            &self,
848        ) -> Pin<
849            Box<dyn Future<Output = SdkResult<crate::ports::catalog::RefreshOutcome>> + Send + '_>,
850        > {
851            Box::pin(async { Ok(crate::ports::catalog::RefreshOutcome::Unchanged) })
852        }
853        fn subscribe(&self) -> broadcast::Receiver<CatalogEvent> {
854            self.tx.subscribe()
855        }
856
857        // ── Sync overrides ──
858        fn get_model_sync(&self, provider: &str, model_id: &str) -> Option<CatalogModelEntry> {
859            if provider == self.entry.provider && model_id == self.entry.model_id {
860                Some(self.entry.clone())
861            } else {
862                None
863            }
864        }
865    }
866
867    /// Regression: AgentBuilder must wire the catalog-aware `Oxicode` as the
868    /// agent loop's resolver, not a static-only closure.
869    ///
870    /// Pre-fix, the resolver consulted only the static `ModelRegistry`
871    /// via `lookup()`, silently dropping the catalog port. Catalog-known
872    /// models (e.g. newer Z.AI models from models.dev) resolved at
873    /// `build()` time via `Oxicode::resolve_model` but failed inside the
874    /// agent loop's `resolve_model()` → "Failed to resolve model".
875    #[test]
876    fn agent_builder_resolver_consults_catalog() {
877        const MODEL_ID: &str = "anthropic/test-catalog-only-model";
878
879        let catalog = Arc::new(SingleModelCatalog {
880            entry: CatalogModelEntry {
881                provider: "anthropic".into(),
882                model_id: "test-catalog-only-model".into(),
883                name: "Test Catalog-Only Model".into(),
884                protocol: CatalogProtocol::AnthropicMessages,
885                source: CatalogSource::Embedded,
886                base_url: None,
887                reasoning: false,
888                supports_vision: false,
889                cost_input: 0.0,
890                cost_output: 0.0,
891                cost_cache_read: 0.0,
892                cost_cache_write: 0.0,
893                context_window: 200_000,
894                max_tokens: 8_192,
895                input_modalities: vec!["text".into()],
896                release_date: None,
897                status: Some("ga".into()),
898            },
899            tx: broadcast::channel(16).0,
900        });
901
902        let oxicode = OxicodeBuilder::new()
903            .with_builtins()
904            .with_catalog(catalog)
905            .build();
906
907        // Sanity: model resolves via Oxicode (catalog→static).
908        assert!(oxicode.resolve_model(MODEL_ID).is_ok());
909
910        // Sanity: model is NOT in the static registry (proves the desync).
911        assert!(
912            oxicode
913                .models_arc()
914                .lookup("anthropic", "test-catalog-only-model")
915                .is_none()
916        );
917
918        // Build an agent with the catalog-only model.
919        let config = AgentConfig {
920            model_id: MODEL_ID.to_string(),
921            ..Default::default()
922        };
923        let agent = oxicode.agent(config).build().unwrap();
924
925        // THE REGRESSION: the agent's loop resolver must also find the
926        // catalog-only model.
927        //   Pre-fix (OxicodeResolver): lookup() → None.
928        //   Post-fix (Oxicode clone): resolve_model() → catalog hit → Some.
929        assert!(
930            agent.resolver().resolve_model(MODEL_ID).is_some(),
931            "AgentBuilder's resolver must consult the catalog port, \
932             not just the static registry"
933        );
934    }
935
936    #[test]
937    fn with_compactor_wires_custom_compactor_into_agent() {
938        /// Compactor that returns a distinctive summary so the test can
939        /// prove the builder's compactor — not a default `LlmCompactor`
940        /// — reached the agent's `CompactionManager`.
941        struct BuilderCompactor;
942
943        impl oxicode_ai::Compactor for BuilderCompactor {
944            fn compact<'a>(
945                &'a self,
946                _messages: &'a [oxicode_ai::Message],
947                _instruction: Option<&'a str>,
948            ) -> Pin<
949                Box<
950                    dyn Future<
951                            Output = std::result::Result<
952                                oxicode_ai::CompactedContext,
953                                oxicode_ai::compaction::CompactionError,
954                            >,
955                        > + Send
956                        + 'a,
957                >,
958            > {
959                Box::pin(async move {
960                    Ok(oxicode_ai::CompactedContext {
961                        summary: "BUILDER-COMPACTOR".to_string(),
962                        kept_messages: Vec::new(),
963                        compacted_count: 1,
964                        metadata: Default::default(),
965                        frames: None,
966                    })
967                })
968            }
969        }
970
971        let oxicode = OxicodeBuilder::new().with_builtins().build();
972
973        let agent = oxicode
974            .agent(AgentConfig {
975                model_id: "anthropic/claude-sonnet-4-20250514".into(),
976                ..Default::default()
977            })
978            .with_compactor(std::sync::Arc::new(BuilderCompactor))
979            .build()
980            .unwrap();
981
982        let rt = tokio::runtime::Runtime::new().unwrap();
983        let ctx = rt
984            .block_on(agent.compaction_manager().compact_now(&[], None))
985            .expect("builder's compactor should be wired through to the agent");
986        assert_eq!(ctx.summary, "BUILDER-COMPACTOR");
987    }
988}