context_harness/
traits.rs

1//! Extension traits for custom connectors and tools.
2//!
3//! This module provides the trait-based extension system for Context Harness.
4//! Users can implement [`Connector`] and [`Tool`] in Rust to create compiled
5//! extensions that run alongside built-in and Lua-scripted ones.
6//!
7//! # Architecture
8//!
9//! ```text
10//! ┌──────────────────────────────────────────┐
11//! │           ConnectorRegistry              │
12//! │  ┌─────────┐ ┌─────────┐ ┌────────────┐ │
13//! │  │Built-in │ │  Lua    │ │  Custom    │ │
14//! │  │FS/Git/S3│ │ Script  │ │ (Rust)     │ │
15//! │  └─────────┘ └─────────┘ └────────────┘ │
16//! └──────────────┬───────────────────────────┘
17//!                ▼
18//!          run_sync() → ingest pipeline
19//! ```
20//!
21//! ```text
22//! ┌──────────────────────────────────────────┐
23//! │              ToolRegistry                │
24//! │  ┌─────────┐ ┌─────────┐ ┌────────────┐ │
25//! │  │Built-in │ │  Lua    │ │  Custom    │ │
26//! │  │search   │ │ Script  │ │ (Rust)     │ │
27//! │  │get/src  │ │ Tools   │ │  Tools     │ │
28//! │  └─────────┘ └─────────┘ └────────────┘ │
29//! └──────────────┬───────────────────────────┘
30//!                ▼
31//!          run_server() → MCP HTTP API
32//! ```
33//!
34//! # Usage
35//!
36//! ```rust
37//! use context_harness::traits::{ConnectorRegistry, ToolRegistry};
38//!
39//! let mut connectors = ConnectorRegistry::new();
40//! // connectors.register(Box::new(MyConnector::new()));
41//!
42//! let mut tools = ToolRegistry::new();
43//! // tools.register(Box::new(MyTool::new()));
44//! ```
45//!
46//! See `docs/RUST_TRAITS.md` for the full specification and examples.
47
48use anyhow::Result;
49use async_trait::async_trait;
50use serde_json::Value;
51use std::sync::Arc;
52
53use crate::config::Config;
54use crate::get::{get_document, DocumentResponse};
55use crate::models::SourceItem;
56use crate::search::{search_documents, SearchResultItem};
57use crate::sources::{get_sources, SourceStatus};
58use crate::workspace::{
59    fan_out, FanOut, RouterError, ServerMode, WorkspaceRouter, ALL_SELECTOR, FAN_OUT_CONCURRENCY,
60};
61
62// ═══════════════════════════════════════════════════════════════════════
63// Connector Trait
64// ═══════════════════════════════════════════════════════════════════════
65
66/// A data source connector that produces documents for ingestion.
67///
68/// Implement this trait to create a custom connector in Rust. The
69/// connector is responsible for scanning an external data source and
70/// returning a list of [`SourceItem`]s that flow through the standard
71/// ingestion pipeline (normalization → chunking → embedding).
72///
73/// # Lifecycle
74///
75/// 1. The connector is registered via [`ConnectorRegistry::register`].
76/// 2. [`scan`](Connector::scan) is called during `ctx sync custom:<name>`.
77/// 3. Returned items are normalized, chunked, and indexed.
78///
79/// # Example
80///
81/// ```rust
82/// use async_trait::async_trait;
83/// use anyhow::Result;
84/// use context_harness::models::SourceItem;
85/// use context_harness::traits::Connector;
86/// use chrono::Utc;
87///
88/// pub struct DatabaseConnector {
89///     connection_string: String,
90/// }
91///
92/// #[async_trait]
93/// impl Connector for DatabaseConnector {
94///     fn name(&self) -> &str { "database" }
95///     fn description(&self) -> &str { "Ingest rows from a database table" }
96///     fn connector_type(&self) -> &str { "custom" }
97///
98///     async fn scan(&self) -> Result<Vec<SourceItem>> {
99///         // ... query database and return SourceItems
100///         Ok(vec![])
101///     }
102/// }
103/// ```
104#[async_trait]
105pub trait Connector: Send + Sync {
106    /// Returns the connector instance name (e.g. `"docs"`, `"platform"`).
107    ///
108    /// Combined with [`connector_type`](Connector::connector_type) to form
109    /// the source label: `"{type}:{name}"`.
110    fn name(&self) -> &str;
111
112    /// Returns a one-line description of what this connector does.
113    ///
114    /// Used in `ctx sources` output and documentation.
115    fn description(&self) -> &str;
116
117    /// Returns the connector type identifier (e.g. `"filesystem"`, `"git"`, `"s3"`, `"custom"`).
118    ///
119    /// Built-in connectors return their type name; custom (user-defined)
120    /// connectors default to `"custom"`.
121    fn connector_type(&self) -> &str {
122        "custom"
123    }
124
125    /// Returns the source label used to tag documents from this connector.
126    ///
127    /// Defaults to `"{connector_type}:{name}"` (e.g. `"git:platform"`).
128    fn source_label(&self) -> String {
129        format!("{}:{}", self.connector_type(), self.name())
130    }
131
132    /// Scan the data source and return all items to ingest.
133    ///
134    /// Called on the tokio async runtime. May perform I/O operations
135    /// (HTTP requests, database queries, file reads).
136    ///
137    /// # Returns
138    ///
139    /// A vector of [`SourceItem`]s. Each item flows through the standard
140    /// ingestion pipeline. Items with empty `body` or `source_id` are
141    /// skipped with a warning.
142    async fn scan(&self) -> Result<Vec<SourceItem>>;
143}
144
145/// Trusted classification for a tool execution that has already started.
146#[derive(Debug, Clone, Copy, PartialEq, Eq)]
147pub enum ToolExecutionErrorClass {
148    RecoverableError,
149    TerminalError,
150    UncertainSideEffect,
151}
152
153impl ToolExecutionErrorClass {
154    pub fn as_str(self) -> &'static str {
155        match self {
156            Self::RecoverableError => "recoverable_error",
157            Self::TerminalError => "terminal_error",
158            Self::UncertainSideEffect => "uncertain_side_effect",
159        }
160    }
161}
162
163/// A validated, model-safe error returned by trusted compiled tool code.
164///
165/// Formatting intentionally omits the public message so ordinary error logging
166/// cannot accidentally turn model-visible prose into an implementation error.
167pub struct ToolExecutionError {
168    class: ToolExecutionErrorClass,
169    code: String,
170    message: String,
171}
172
173impl ToolExecutionError {
174    pub fn recoverable(code: impl Into<String>, message: impl Into<String>) -> Result<Self> {
175        Self::new(ToolExecutionErrorClass::RecoverableError, code, message)
176    }
177
178    pub fn terminal(code: impl Into<String>, message: impl Into<String>) -> Result<Self> {
179        Self::new(ToolExecutionErrorClass::TerminalError, code, message)
180    }
181
182    pub fn uncertain_side_effect(
183        code: impl Into<String>,
184        message: impl Into<String>,
185    ) -> Result<Self> {
186        Self::new(ToolExecutionErrorClass::UncertainSideEffect, code, message)
187    }
188
189    fn new(
190        class: ToolExecutionErrorClass,
191        code: impl Into<String>,
192        message: impl Into<String>,
193    ) -> Result<Self> {
194        let code = code.into();
195        let message = message.into();
196        anyhow::ensure!(
197            !code.is_empty()
198                && code.len() <= 64
199                && code.bytes().enumerate().all(|(index, byte)| if index == 0 {
200                    byte.is_ascii_lowercase()
201                } else {
202                    byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'_'
203                }),
204            "invalid tool execution error code"
205        );
206        anyhow::ensure!(
207            !message.trim().is_empty() && message.len() <= 4096,
208            "invalid tool execution error message"
209        );
210        Ok(Self {
211            class,
212            code,
213            message,
214        })
215    }
216
217    pub fn class(&self) -> ToolExecutionErrorClass {
218        self.class
219    }
220
221    pub fn code(&self) -> &str {
222        &self.code
223    }
224
225    pub fn message(&self) -> &str {
226        &self.message
227    }
228}
229
230impl std::fmt::Debug for ToolExecutionError {
231    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
232        formatter
233            .debug_struct("ToolExecutionError")
234            .field("class", &self.class.as_str())
235            .field("code", &self.code)
236            .finish_non_exhaustive()
237    }
238}
239
240impl std::fmt::Display for ToolExecutionError {
241    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
242        write!(formatter, "{}:{}", self.class.as_str(), self.code)
243    }
244}
245
246impl std::error::Error for ToolExecutionError {}
247
248// ═══════════════════════════════════════════════════════════════════════
249// Tool Trait
250// ═══════════════════════════════════════════════════════════════════════
251
252/// A custom MCP tool that agents can discover and call.
253///
254/// Implement this trait to create a compiled Rust tool. Tools are
255/// registered at server startup and exposed via `GET /tools/list`
256/// for agent discovery and `POST /tools/{name}` for invocation.
257///
258/// # Lifecycle
259///
260/// 1. The tool is registered via [`ToolRegistry::register`].
261/// 2. [`name`](Tool::name), [`description`](Tool::description), and
262///    [`parameters_schema`](Tool::parameters_schema) are called at startup
263///    for the tool list.
264/// 3. [`execute`](Tool::execute) is called each time an agent invokes
265///    the tool.
266///
267/// # Example
268///
269/// ```rust
270/// use async_trait::async_trait;
271/// use anyhow::Result;
272/// use serde_json::{json, Value};
273/// use context_harness::traits::{Tool, ToolContext};
274///
275/// pub struct HealthCheckTool;
276///
277/// #[async_trait]
278/// impl Tool for HealthCheckTool {
279///     fn name(&self) -> &str { "health_check" }
280///     fn description(&self) -> &str { "Check connector health" }
281///
282///     fn parameters_schema(&self) -> Value {
283///         json!({
284///             "type": "object",
285///             "properties": {},
286///             "required": []
287///         })
288///     }
289///
290///     async fn execute(&self, _params: Value, ctx: &ToolContext) -> Result<Value> {
291///         let sources = ctx.sources()?;
292///         Ok(json!({ "sources": sources.len() }))
293///     }
294/// }
295/// ```
296#[async_trait]
297pub trait Tool: Send + Sync {
298    /// Runtime capability requirements. Unknown metadata is denied by the local
299    /// agent runtime; existing MCP tools retain their existing behavior.
300    fn capabilities(&self) -> Option<Vec<crate::agent_resource::Capability>> {
301        None
302    }
303
304    /// Sanitized immutable identity for declaratively bound tools. History and
305    /// recovery use this without loading implementations or resolving secrets.
306    fn binding_metadata(&self) -> Option<Value> {
307        None
308    }
309
310    /// Select runtime-owned execution context without coupling the dispatcher
311    /// to a public tool name. Most tools execute directly.
312    fn runtime_dispatch(&self) -> ToolRuntimeDispatch {
313        ToolRuntimeDispatch::Direct
314    }
315
316    /// Optional host-owned preparation that must complete through the durable
317    /// approval lifecycle before this tool can be advertised to a model.
318    fn preparation(&self) -> Option<ToolPreparation> {
319        None
320    }
321
322    async fn prepare(&self, _ctx: &ToolContext) -> Result<()> {
323        Ok(())
324    }
325
326    /// Returns the tool's name.
327    ///
328    /// Used as the route path (`POST /tools/{name}`) and in
329    /// `GET /tools/list` responses. Should be a lowercase
330    /// identifier with underscores (e.g., `"create_ticket"`).
331    fn name(&self) -> &str;
332
333    /// Returns a one-line description for agent discovery.
334    ///
335    /// Agents use this to decide whether to call the tool.
336    fn description(&self) -> &str;
337
338    /// Whether this tool is a built-in (true for search/get/sources).
339    ///
340    /// Built-in tools are marked with `"builtin": true` in the
341    /// `GET /tools/list` response. Defaults to `false`.
342    fn is_builtin(&self) -> bool {
343        false
344    }
345
346    /// Returns the OpenAI function-calling JSON Schema for parameters.
347    ///
348    /// Must be a valid JSON Schema object with `type: "object"`,
349    /// `properties`, and optionally `required`.
350    fn parameters_schema(&self) -> Value;
351
352    /// Validate model-visible arguments before authorization and tool start.
353    /// Bound adapters override this with their resolved public contract.
354    fn validate_arguments(&self, params: &Value) -> Result<()> {
355        anyhow::ensure!(params.is_object(), "tool arguments must be an object");
356        Ok(())
357    }
358
359    /// Sanitized effective arguments covered by approval. Bound tools override
360    /// this to include immutable fixed values and reject override attempts.
361    fn approval_arguments(&self, params: &Value) -> Result<Value> {
362        Ok(params.clone())
363    }
364
365    /// Execute the tool with validated parameters.
366    ///
367    /// Called each time an agent invokes the tool via `POST /tools/{name}`.
368    ///
369    /// # Arguments
370    ///
371    /// * `params` — JSON parameters (always a JSON object).
372    /// * `ctx` — Bridge to the Context Harness knowledge base.
373    ///
374    /// # Returns
375    ///
376    /// A JSON value that will be wrapped in `{ "result": ... }` in the
377    /// HTTP response.
378    async fn execute(&self, params: Value, ctx: &ToolContext) -> Result<Value>;
379}
380
381#[derive(Debug, Clone, Copy, PartialEq, Eq)]
382pub enum ToolRuntimeDispatch {
383    Direct,
384    AgentDelegation,
385    RunControl(RunControlKind),
386}
387
388#[derive(Debug, Clone, Copy, PartialEq, Eq)]
389pub enum RunControlKind {
390    Blocked,
391    RequestUserInput,
392}
393
394#[derive(Debug, Clone)]
395pub struct ToolPreparation {
396    pub name: String,
397    pub capabilities: Vec<crate::agent_resource::Capability>,
398    pub arguments: Value,
399}
400
401// ═══════════════════════════════════════════════════════════════════════
402// ToolContext
403// ═══════════════════════════════════════════════════════════════════════
404
405/// Options for [`ToolContext::search`].
406#[derive(Debug, Default)]
407pub struct SearchOptions {
408    /// Search mode: `"keyword"`, `"semantic"`, or `"hybrid"`.
409    pub mode: Option<String>,
410    /// Maximum number of results.
411    pub limit: Option<i64>,
412    /// Filter by source connector (e.g., `"git:platform"`).
413    pub source: Option<String>,
414}
415
416/// Context bridge for tool execution.
417///
418/// Provides tools with access to the Context Harness knowledge base
419/// during execution. Created by the server for each tool invocation.
420///
421/// All methods delegate to the same core functions used by the CLI
422/// and HTTP server, ensuring tools have identical capabilities.
423pub struct ToolContext {
424    /// The active workspace config used by the convenience methods below and by
425    /// compatibility-mode built-in tools. In multi mode this is the default
426    /// workspace's config; router-aware tools resolve per call instead.
427    config: Arc<Config>,
428    /// The workspace router. In compatibility mode this is a one-workspace
429    /// router; router-aware (multi-mode) tools dispatch through it.
430    router: Arc<WorkspaceRouter>,
431    /// Whether this context serves the pre-router flat shapes or labeled shapes.
432    mode: ServerMode,
433}
434
435impl ToolContext {
436    /// Create a single-workspace (compatibility) tool context from a config.
437    ///
438    /// Internally wraps the config in a one-workspace [`WorkspaceRouter`], so a
439    /// compatibility context is just "a router with one workspace".
440    pub fn new(config: Arc<Config>) -> Self {
441        let router = Arc::new(WorkspaceRouter::single(config.clone()));
442        Self {
443            config,
444            router,
445            mode: ServerMode::Compat,
446        }
447    }
448
449    /// Create a tool context backed by a (possibly multi-workspace) router.
450    ///
451    /// The convenience methods default to the router's default workspace; the
452    /// router-aware built-in tools resolve the request's `workspace` selector
453    /// against [`ToolContext::router`] per call.
454    pub fn routed(router: Arc<WorkspaceRouter>, mode: ServerMode) -> Self {
455        let config = router.default_config();
456        Self {
457            config,
458            router,
459            mode,
460        }
461    }
462
463    /// The server mode this context serves.
464    pub fn mode(&self) -> ServerMode {
465        self.mode
466    }
467
468    /// The workspace router, for router-aware tools.
469    pub fn router(&self) -> &Arc<WorkspaceRouter> {
470        &self.router
471    }
472
473    /// The active workspace config (the default workspace in multi mode).
474    pub fn config(&self) -> &Arc<Config> {
475        &self.config
476    }
477
478    /// Search the knowledge base.
479    ///
480    /// Equivalent to `POST /tools/search` or `ctx search`.
481    ///
482    /// # Example
483    ///
484    /// ```rust,no_run
485    /// # use context_harness::traits::{ToolContext, SearchOptions};
486    /// # async fn example(ctx: &ToolContext) -> anyhow::Result<()> {
487    /// let results = ctx.search("deployment runbook", SearchOptions {
488    ///     mode: Some("hybrid".to_string()),
489    ///     limit: Some(5),
490    ///     ..Default::default()
491    /// }).await?;
492    /// # Ok(())
493    /// # }
494    /// ```
495    pub async fn search(&self, query: &str, opts: SearchOptions) -> Result<Vec<SearchResultItem>> {
496        search_documents(
497            &self.config,
498            query,
499            opts.mode.as_deref().unwrap_or("keyword"),
500            opts.source.as_deref(),
501            None,
502            opts.limit,
503            false,
504        )
505        .await
506    }
507
508    /// Retrieve a document by UUID.
509    ///
510    /// Equivalent to `POST /tools/get` or `ctx get`.
511    pub async fn get(&self, id: &str) -> Result<DocumentResponse> {
512        get_document(&self.config, id).await
513    }
514
515    /// List all configured connectors and their status.
516    ///
517    /// Equivalent to `GET /tools/sources` or `ctx sources`.
518    pub fn sources(&self) -> Result<Vec<SourceStatus>> {
519        Ok(get_sources(&self.config))
520    }
521}
522
523// ═══════════════════════════════════════════════════════════════════════
524// Built-in Tool Implementations
525// ═══════════════════════════════════════════════════════════════════════
526
527/// Built-in search tool. Delegates to [`ToolContext::search`].
528pub struct SearchTool;
529
530#[async_trait]
531impl Tool for SearchTool {
532    fn name(&self) -> &str {
533        "search"
534    }
535
536    fn description(&self) -> &str {
537        "Search the knowledge base"
538    }
539
540    fn is_builtin(&self) -> bool {
541        true
542    }
543
544    fn parameters_schema(&self) -> Value {
545        serde_json::json!({
546            "type": "object",
547            "properties": {
548                "query": { "type": "string", "description": "Search query" },
549                "mode": { "type": "string", "enum": ["keyword", "semantic", "hybrid"], "default": "keyword" },
550                "limit": { "type": "integer", "description": "Max results", "default": 12 },
551                "filters": {
552                    "type": "object",
553                    "properties": {
554                        "source": { "type": "string", "description": "Filter by connector source" },
555                        "since": { "type": "string", "description": "Only results updated after this date (YYYY-MM-DD)" }
556                    }
557                }
558            },
559            "required": ["query"]
560        })
561    }
562
563    async fn execute(&self, params: Value, ctx: &ToolContext) -> Result<Value> {
564        let query = params["query"].as_str().unwrap_or("");
565        if query.trim().is_empty() {
566            anyhow::bail!("query must not be empty");
567        }
568
569        let mode = params["mode"].as_str().unwrap_or("keyword");
570        let limit = params["limit"].as_i64().unwrap_or(12);
571
572        let source = params
573            .get("filters")
574            .and_then(|f| f.get("source"))
575            .and_then(|s| s.as_str());
576        let since = params
577            .get("filters")
578            .and_then(|f| f.get("since"))
579            .and_then(|s| s.as_str());
580
581        let results =
582            search_documents(&ctx.config, query, mode, source, since, Some(limit), false).await?;
583
584        Ok(serde_json::json!({ "results": results }))
585    }
586}
587
588/// Built-in document retrieval tool. Delegates to [`get_document`].
589pub struct GetTool;
590
591#[async_trait]
592impl Tool for GetTool {
593    fn name(&self) -> &str {
594        "get"
595    }
596
597    fn description(&self) -> &str {
598        "Retrieve a document by UUID"
599    }
600
601    fn is_builtin(&self) -> bool {
602        true
603    }
604
605    fn parameters_schema(&self) -> Value {
606        serde_json::json!({
607            "type": "object",
608            "properties": {
609                "id": { "type": "string", "description": "Document UUID" }
610            },
611            "required": ["id"]
612        })
613    }
614
615    async fn execute(&self, params: Value, ctx: &ToolContext) -> Result<Value> {
616        let id = params["id"].as_str().unwrap_or("");
617        if id.trim().is_empty() {
618            anyhow::bail!("id must not be empty");
619        }
620
621        let doc = get_document(&ctx.config, id).await?;
622        Ok(serde_json::to_value(&doc)?)
623    }
624}
625
626/// Built-in sources listing tool. Delegates to [`get_sources`].
627pub struct SourcesTool;
628
629#[async_trait]
630impl Tool for SourcesTool {
631    fn name(&self) -> &str {
632        "sources"
633    }
634
635    fn description(&self) -> &str {
636        "List connector configuration and health status"
637    }
638
639    fn is_builtin(&self) -> bool {
640        true
641    }
642
643    fn parameters_schema(&self) -> Value {
644        serde_json::json!({
645            "type": "object",
646            "properties": {}
647        })
648    }
649
650    async fn execute(&self, _params: Value, ctx: &ToolContext) -> Result<Value> {
651        let sources = get_sources(&ctx.config);
652        Ok(serde_json::json!({ "sources": sources }))
653    }
654}
655
656// ═══════════════════════════════════════════════════════════════════════
657// Router-aware built-in tools (multi-workspace mode)
658// ═══════════════════════════════════════════════════════════════════════
659//
660// These mirror the compat built-ins but accept an optional `workspace`
661// selector, dispatch through the request's [`WorkspaceRouter`], and emit the
662// workspace-labeled (grouped) shapes (SPEC-0014 R18–R37). They are registered
663// only in multi-workspace mode; compatibility mode keeps the unit-struct tools
664// above unchanged (the additive invariant). `workspace = "all"` fan-out is
665// Phase 2 — the router currently rejects it with `unsupported_workspace_selector`.
666
667/// Add the optional `workspace` selector to a compat schema's `properties`.
668fn schema_with_workspace_selector(mut schema: Value) -> Value {
669    if let Some(props) = schema.get_mut("properties").and_then(|p| p.as_object_mut()) {
670        props.insert(
671            "workspace".to_string(),
672            serde_json::json!({
673                "type": "string",
674                "description": "Workspace id to target (omit for the default workspace)"
675            }),
676        );
677    }
678    schema
679}
680
681/// Build one `{ workspace, items }` group, tagging every item with `workspace`
682/// and a `qualified_id` (R29/R30/R36).
683fn search_group(workspace: &str, results: Vec<SearchResultItem>) -> Value {
684    let items: Vec<Value> = results
685        .into_iter()
686        .map(|item| {
687            let mut obj = serde_json::to_value(&item).unwrap_or(Value::Null);
688            if let Value::Object(map) = &mut obj {
689                let doc_id = map
690                    .get("id")
691                    .and_then(|v| v.as_str())
692                    .unwrap_or_default()
693                    .to_string();
694                map.insert(
695                    "workspace".to_string(),
696                    Value::String(workspace.to_string()),
697                );
698                map.insert(
699                    "qualified_id".to_string(),
700                    Value::String(format!("{workspace}:{doc_id}")),
701                );
702            }
703            obj
704        })
705        .collect();
706    serde_json::json!({ "workspace": workspace, "items": items })
707}
708
709/// A single-workspace grouped response: one group, no errors (R29/R30).
710fn shape_grouped_search(workspace: &str, results: Vec<SearchResultItem>) -> Value {
711    serde_json::json!({ "results": [ search_group(workspace, results) ], "errors": [] })
712}
713
714/// Build one `{ workspace, sources }` group for the `sources` tool.
715fn sources_group(workspace: &str, sources: Vec<SourceStatus>) -> Value {
716    serde_json::json!({ "workspace": workspace, "sources": sources })
717}
718
719/// Build one `errors[]` entry for a failed workspace (R37):
720/// `{ workspace, code, message }`.
721fn error_entry(workspace: &str, err: &RouterError) -> Value {
722    serde_json::json!({
723        "workspace": workspace,
724        "code": err.code(),
725        "message": err.to_string(),
726    })
727}
728
729/// Fan out a `search` across every enabled workspace (R32–R37): concurrent,
730/// per-workspace deadline, per-workspace `limit`, grouped results, and one
731/// `errors[]` entry per failed/timed-out workspace.
732async fn search_all(
733    ctx: &ToolContext,
734    query: String,
735    mode: String,
736    source: Option<String>,
737    since: Option<String>,
738    limit: i64,
739) -> Result<Value> {
740    let router = ctx.router();
741    let (healthy, seed_errors) = router.resolve_all();
742    let deadline = router.search_deadline();
743    let deadline_ms = deadline.as_millis() as u64;
744
745    let items: Vec<(String, Arc<Config>)> = healthy
746        .iter()
747        .map(|rt| (rt.id.clone(), rt.config.clone()))
748        .collect();
749
750    let outcomes = fan_out(
751        items,
752        deadline,
753        FAN_OUT_CONCURRENCY,
754        move |config: Arc<Config>| {
755            let query = query.clone();
756            let mode = mode.clone();
757            let source = source.clone();
758            let since = since.clone();
759            async move {
760                search_documents(
761                    &config,
762                    &query,
763                    &mode,
764                    source.as_deref(),
765                    since.as_deref(),
766                    Some(limit),
767                    false,
768                )
769                .await
770            }
771        },
772    )
773    .await;
774
775    let mut groups: Vec<Value> = Vec::new();
776    let mut error_pairs: Vec<(String, Value)> = seed_errors
777        .iter()
778        .map(|(id, err)| (id.clone(), error_entry(id, err)))
779        .collect();
780
781    for (id, outcome) in outcomes {
782        match outcome {
783            FanOut::Ok(results) => groups.push(search_group(&id, results)),
784            FanOut::Failed(msg) => error_pairs.push((
785                id.clone(),
786                error_entry(
787                    &id,
788                    &RouterError::WorkspaceUnavailable {
789                        id: id.clone(),
790                        reason: msg,
791                    },
792                ),
793            )),
794            FanOut::TimedOut => error_pairs.push((
795                id.clone(),
796                error_entry(
797                    &id,
798                    &RouterError::WorkspaceTimeout {
799                        id: id.clone(),
800                        deadline_ms,
801                    },
802                ),
803            )),
804        }
805    }
806
807    // errors[] follows the stable registry order (DESIGN-0008).
808    let ws_order: Vec<String> = ctx.router().list().iter().map(|rt| rt.id.clone()).collect();
809    error_pairs.sort_by_key(|(id, _)| ws_order.iter().position(|w| w == id).unwrap_or(usize::MAX));
810    let errors: Vec<Value> = error_pairs.into_iter().map(|(_, v)| v).collect();
811
812    Ok(serde_json::json!({ "results": groups, "errors": errors }))
813}
814
815/// Multi-workspace `search`: resolves the `workspace` selector and groups results.
816pub struct RoutedSearchTool;
817
818#[async_trait]
819impl Tool for RoutedSearchTool {
820    fn name(&self) -> &str {
821        "search"
822    }
823
824    fn description(&self) -> &str {
825        "Search the knowledge base (optionally scoped to a workspace)"
826    }
827
828    fn is_builtin(&self) -> bool {
829        true
830    }
831
832    fn parameters_schema(&self) -> Value {
833        schema_with_workspace_selector(SearchTool.parameters_schema())
834    }
835
836    async fn execute(&self, params: Value, ctx: &ToolContext) -> Result<Value> {
837        let query = params["query"].as_str().unwrap_or("");
838        if query.trim().is_empty() {
839            anyhow::bail!("query must not be empty");
840        }
841        let mode = params["mode"].as_str().unwrap_or("keyword");
842        let limit = params["limit"].as_i64().unwrap_or(12);
843        let source = params
844            .get("filters")
845            .and_then(|f| f.get("source"))
846            .and_then(|s| s.as_str());
847        let since = params
848            .get("filters")
849            .and_then(|f| f.get("since"))
850            .and_then(|s| s.as_str());
851
852        let selector = params["workspace"].as_str();
853        if selector == Some(ALL_SELECTOR) {
854            return search_all(
855                ctx,
856                query.to_string(),
857                mode.to_string(),
858                source.map(str::to_string),
859                since.map(str::to_string),
860                limit,
861            )
862            .await;
863        }
864        let runtime = ctx.router().resolve(selector)?;
865
866        let results = search_documents(
867            &runtime.config,
868            query,
869            mode,
870            source,
871            since,
872            Some(limit),
873            false,
874        )
875        .await?;
876        Ok(shape_grouped_search(&runtime.id, results))
877    }
878}
879
880/// Multi-workspace `get`: supports a qualified id (`<ws>:<doc>`) or an explicit
881/// `workspace` selector, with conflict detection (R38–R43).
882pub struct RoutedGetTool;
883
884#[async_trait]
885impl Tool for RoutedGetTool {
886    fn name(&self) -> &str {
887        "get"
888    }
889
890    fn description(&self) -> &str {
891        "Retrieve a document by id or qualified id (<workspace>:<id>)"
892    }
893
894    fn is_builtin(&self) -> bool {
895        true
896    }
897
898    fn parameters_schema(&self) -> Value {
899        schema_with_workspace_selector(GetTool.parameters_schema())
900    }
901
902    async fn execute(&self, params: Value, ctx: &ToolContext) -> Result<Value> {
903        let id = params["id"].as_str().unwrap_or("");
904        if id.trim().is_empty() {
905            anyhow::bail!("id must not be empty");
906        }
907        let field = params["workspace"].as_str();
908        let (qualified_prefix, raw_id) = ctx.router().split_qualified_id(id);
909
910        // A qualified id and an explicit workspace field must agree (R41/R42).
911        if let (Some(q), Some(f)) = (qualified_prefix, field) {
912            if q != f {
913                return Err(RouterError::WorkspaceIdConflict {
914                    field: f.to_string(),
915                    qualified: q.to_string(),
916                }
917                .into());
918            }
919        }
920
921        let selector = qualified_prefix.or(field);
922        let runtime = ctx.router().resolve(selector)?;
923
924        let doc = get_document(&runtime.config, raw_id).await?;
925        let mut obj = serde_json::to_value(&doc)?;
926        if let Value::Object(map) = &mut obj {
927            map.insert("workspace".to_string(), Value::String(runtime.id.clone()));
928            map.insert(
929                "qualified_id".to_string(),
930                Value::String(format!("{}:{}", runtime.id, raw_id)),
931            );
932        }
933        Ok(obj)
934    }
935}
936
937/// Multi-workspace `sources`: connector status for one workspace, grouped and
938/// redacted (R44/R48/R49).
939pub struct RoutedSourcesTool;
940
941#[async_trait]
942impl Tool for RoutedSourcesTool {
943    fn name(&self) -> &str {
944        "sources"
945    }
946
947    fn description(&self) -> &str {
948        "List connector status for a workspace (optionally scoped to a workspace)"
949    }
950
951    fn is_builtin(&self) -> bool {
952        true
953    }
954
955    fn parameters_schema(&self) -> Value {
956        schema_with_workspace_selector(SourcesTool.parameters_schema())
957    }
958
959    async fn execute(&self, params: Value, ctx: &ToolContext) -> Result<Value> {
960        let selector = params["workspace"].as_str();
961        if selector == Some(ALL_SELECTOR) {
962            let (healthy, seed_errors) = ctx.router().resolve_all();
963            let results: Vec<Value> = healthy
964                .iter()
965                .map(|rt| sources_group(&rt.id, get_sources(&rt.config)))
966                .collect();
967            let errors: Vec<Value> = seed_errors
968                .iter()
969                .map(|(id, err)| error_entry(id, err))
970                .collect();
971            return Ok(serde_json::json!({ "results": results, "errors": errors }));
972        }
973        let runtime = ctx.router().resolve(selector)?;
974        let sources = get_sources(&runtime.config);
975        Ok(serde_json::json!({
976            "results": [ sources_group(&runtime.id, sources) ],
977            "errors": []
978        }))
979    }
980}
981
982/// Multi-workspace discovery tool: lists registered workspaces, their health,
983/// and which is the default — without exposing any secret config values (R46–R49).
984pub struct WorkspacesTool;
985
986#[async_trait]
987impl Tool for WorkspacesTool {
988    fn name(&self) -> &str {
989        "workspaces"
990    }
991
992    fn description(&self) -> &str {
993        "List registered workspaces, their health, and the default workspace"
994    }
995
996    fn is_builtin(&self) -> bool {
997        true
998    }
999
1000    fn parameters_schema(&self) -> Value {
1001        serde_json::json!({ "type": "object", "properties": {} })
1002    }
1003
1004    async fn execute(&self, _params: Value, ctx: &ToolContext) -> Result<Value> {
1005        let router = ctx.router();
1006        let default = router.default_id();
1007        let workspaces: Vec<Value> = router
1008            .list()
1009            .iter()
1010            .map(|rt| {
1011                let health = match &rt.health {
1012                    crate::workspace::WorkspaceHealth::Ok => serde_json::json!({ "status": "ok" }),
1013                    crate::workspace::WorkspaceHealth::Unavailable(reason) => {
1014                        serde_json::json!({ "status": "unavailable", "reason": reason })
1015                    }
1016                };
1017                serde_json::json!({
1018                    "id": rt.id,
1019                    "root": rt.root.as_ref().map(|p| p.display().to_string()),
1020                    "enabled": rt.enabled,
1021                    "default": Some(rt.id.as_str()) == default,
1022                    "resolution": rt.resolution.as_str(),
1023                    "health": health,
1024                })
1025            })
1026            .collect();
1027        Ok(serde_json::json!({ "workspaces": workspaces }))
1028    }
1029}
1030
1031// ═══════════════════════════════════════════════════════════════════════
1032// Registries
1033// ═══════════════════════════════════════════════════════════════════════
1034
1035/// Registry for connectors (built-in and custom).
1036///
1037/// Use [`ConnectorRegistry::from_config`] to create a registry pre-loaded
1038/// with all built-in connectors from the config file, then optionally
1039/// call [`register`](ConnectorRegistry::register) to add custom ones.
1040///
1041/// # Example
1042///
1043/// ```rust
1044/// use context_harness::traits::ConnectorRegistry;
1045///
1046/// let mut connectors = ConnectorRegistry::new();
1047/// // connectors.register(Box::new(MyConnector::new()));
1048/// ```
1049pub struct ConnectorRegistry {
1050    connectors: Vec<Box<dyn Connector>>,
1051}
1052
1053impl ConnectorRegistry {
1054    /// Create an empty connector registry.
1055    pub fn new() -> Self {
1056        Self {
1057            connectors: Vec::new(),
1058        }
1059    }
1060
1061    /// Create a registry pre-loaded with all built-in connectors from the config.
1062    ///
1063    /// This resolves all filesystem, git, S3, and script connector instances
1064    /// from the TOML config and wraps them as trait objects.
1065    pub fn from_config(config: &Config) -> Self {
1066        use crate::connector_fs::FilesystemConnector;
1067        use crate::connector_git::GitConnector;
1068        use crate::connector_s3::S3Connector;
1069        use crate::connector_script::ScriptConnector;
1070
1071        let mut registry = Self::new();
1072
1073        for (name, cfg) in &config.connectors.filesystem {
1074            registry.register(Box::new(FilesystemConnector::new(
1075                name.clone(),
1076                cfg.clone(),
1077            )));
1078        }
1079        for (name, cfg) in &config.connectors.git {
1080            registry.register(Box::new(GitConnector::new(
1081                name.clone(),
1082                cfg.clone(),
1083                config.db.path.clone(),
1084            )));
1085        }
1086        for (name, cfg) in &config.connectors.s3 {
1087            registry.register(Box::new(S3Connector::new(name.clone(), cfg.clone())));
1088        }
1089        for (name, cfg) in &config.connectors.script {
1090            registry.register(Box::new(ScriptConnector::new(name.clone(), cfg.clone())));
1091        }
1092
1093        registry
1094    }
1095
1096    /// Register a connector.
1097    pub fn register(&mut self, connector: Box<dyn Connector>) {
1098        self.connectors.push(connector);
1099    }
1100
1101    /// Get all registered connectors.
1102    pub fn connectors(&self) -> &[Box<dyn Connector>] {
1103        &self.connectors
1104    }
1105
1106    /// Get connectors filtered by type (e.g. `"git"`, `"filesystem"`).
1107    pub fn connectors_by_type(&self, connector_type: &str) -> Vec<&dyn Connector> {
1108        self.connectors
1109            .iter()
1110            .filter(|c| c.connector_type() == connector_type)
1111            .map(|c| c.as_ref())
1112            .collect()
1113    }
1114
1115    /// Find a specific connector by type and name.
1116    pub fn find(&self, connector_type: &str, name: &str) -> Option<&dyn Connector> {
1117        self.connectors
1118            .iter()
1119            .find(|c| c.connector_type() == connector_type && c.name() == name)
1120            .map(|c| c.as_ref())
1121    }
1122
1123    /// Check if the registry is empty.
1124    pub fn is_empty(&self) -> bool {
1125        self.connectors.is_empty()
1126    }
1127
1128    /// Return the count of registered connectors.
1129    pub fn len(&self) -> usize {
1130        self.connectors.len()
1131    }
1132}
1133
1134impl Default for ConnectorRegistry {
1135    fn default() -> Self {
1136        Self::new()
1137    }
1138}
1139
1140/// Registry for tools (built-in, Lua, and custom Rust).
1141///
1142/// Use [`ToolRegistry::with_builtins`] to create a registry pre-loaded
1143/// with the core `search`, `get`, and `sources` tools, then optionally
1144/// call [`register`](ToolRegistry::register) to add custom ones.
1145///
1146/// # Example
1147///
1148/// ```rust
1149/// use context_harness::traits::ToolRegistry;
1150///
1151/// let mut tools = ToolRegistry::with_builtins();
1152/// // tools.register(Box::new(MyTool::new()));
1153/// ```
1154pub struct ToolRegistry {
1155    tools: Vec<Box<dyn Tool>>,
1156}
1157
1158impl ToolRegistry {
1159    /// Create an empty tool registry.
1160    pub fn new() -> Self {
1161        Self { tools: Vec::new() }
1162    }
1163
1164    /// Create a tool registry pre-loaded with built-in tools (search, get, sources).
1165    pub fn with_builtins() -> Self {
1166        let mut registry = Self::new();
1167        registry.register(Box::new(SearchTool));
1168        registry.register(Box::new(GetTool));
1169        registry.register(Box::new(SourcesTool));
1170        registry
1171    }
1172
1173    /// Create a tool registry with the router-aware built-ins for
1174    /// multi-workspace mode: `search`/`get`/`sources` accept a `workspace`
1175    /// selector and the `workspaces` discovery tool is added. Workspace-local
1176    /// Lua/Rust tools are intentionally not registered in multi mode in Phase 1
1177    /// (SPEC-0014 R54).
1178    pub fn with_builtins_multi() -> Self {
1179        let mut registry = Self::new();
1180        registry.register(Box::new(RoutedSearchTool));
1181        registry.register(Box::new(RoutedGetTool));
1182        registry.register(Box::new(RoutedSourcesTool));
1183        registry.register(Box::new(WorkspacesTool));
1184        registry
1185    }
1186
1187    /// Register a tool.
1188    pub fn register(&mut self, tool: Box<dyn Tool>) {
1189        self.tools.push(tool);
1190    }
1191
1192    /// Get all registered tools.
1193    pub fn tools(&self) -> &[Box<dyn Tool>] {
1194        &self.tools
1195    }
1196
1197    /// Consume the registry and return its tools for composition into another registry.
1198    pub fn into_tools(self) -> Vec<Box<dyn Tool>> {
1199        self.tools
1200    }
1201
1202    /// Find a tool by name.
1203    pub fn find(&self, name: &str) -> Option<&dyn Tool> {
1204        self.tools
1205            .iter()
1206            .find(|t| t.name() == name)
1207            .map(|t| t.as_ref())
1208    }
1209
1210    /// Check if the registry is empty.
1211    pub fn is_empty(&self) -> bool {
1212        self.tools.is_empty()
1213    }
1214
1215    /// Return the count of registered tools.
1216    pub fn len(&self) -> usize {
1217        self.tools.len()
1218    }
1219}
1220
1221impl Default for ToolRegistry {
1222    fn default() -> Self {
1223        Self::new()
1224    }
1225}
1226
1227#[cfg(test)]
1228mod phase2_tests {
1229    use super::*;
1230    use crate::workspace::RouterError;
1231
1232    #[test]
1233    fn search_group_tags_items_with_workspace_and_qualified_id() {
1234        let item = SearchResultItem {
1235            id: "01ABC".to_string(),
1236            score: 0.5,
1237            title: Some("T".to_string()),
1238            source: "filesystem".to_string(),
1239            source_id: "f".to_string(),
1240            updated_at: "2026-01-01T00:00:00Z".to_string(),
1241            snippet: "snip".to_string(),
1242            source_url: None,
1243            explain: None,
1244        };
1245        let group = search_group("beta", vec![item]);
1246        assert_eq!(group["workspace"], "beta");
1247        let items = group["items"].as_array().unwrap();
1248        assert_eq!(items[0]["workspace"], "beta");
1249        assert_eq!(items[0]["qualified_id"], "beta:01ABC");
1250    }
1251
1252    #[test]
1253    fn error_entry_has_workspace_code_message() {
1254        let err = RouterError::WorkspaceTimeout {
1255            id: "beta".to_string(),
1256            deadline_ms: 5000,
1257        };
1258        let entry = error_entry("beta", &err);
1259        assert_eq!(entry["workspace"], "beta");
1260        assert_eq!(entry["code"], "workspace_timeout");
1261        assert!(entry["message"].as_str().unwrap().contains("5000"));
1262    }
1263
1264    #[test]
1265    fn tool_execution_errors_validate_and_redact_public_messages() {
1266        let error = ToolExecutionError::recoverable("not_found", "Safe public message").unwrap();
1267        assert_eq!(error.class(), ToolExecutionErrorClass::RecoverableError);
1268        assert_eq!(error.code(), "not_found");
1269        assert_eq!(error.message(), "Safe public message");
1270        assert_eq!(error.to_string(), "recoverable_error:not_found");
1271        assert!(!format!("{error:?}").contains("Safe public message"));
1272
1273        for code in ["", "NotFound", "1bad", "bad-code", &"a".repeat(65)] {
1274            assert!(ToolExecutionError::terminal(code, "safe").is_err());
1275        }
1276        for message in ["", "   ", &"x".repeat(4097)] {
1277            assert!(ToolExecutionError::uncertain_side_effect("unknown", message).is_err());
1278        }
1279    }
1280}
context_harness/traits.rs

context_harness/
traits.rs
