From 4d00dbb27a314526558824730cea04bdae107c95 Mon Sep 17 00:00:00 2001 From: keepmeside Date: Tue, 22 Sep 2026 21:16:05 +0000 Subject: [PATCH] feat(context): add hybrid retrieval and typed agent decisions --- Cargo.toml | 2 +- README-vi.md | 23 +- README-zh.md | 14 +- README.md | 24 +- plans/2026-09-22-ctxe-system-one.md | 57 +++ plans/260922-2052-turso-prototype/plan.md | 26 ++ plans/260922-2052-typesafe-api/plan.md | 36 ++ .../2026-09-22-ctxe-system-one-research.md | 53 +++ plans/reports/turso-260922-2052-assessment.md | 57 +++ src/config.rs | 82 ++++- src/indexing/pipeline.rs | 15 +- src/lib.rs | 1 + src/main.rs | 22 +- src/mcp.rs | 326 +++++++++++++++++- src/mcp/tests.rs | 19 + src/prompts/mcp_ask_context.txt | 5 + src/query/decision.rs | 306 ++++++++++++++++ src/query/engine.rs | 165 ++++++++- src/query/hybrid.rs | 234 +++++++++++++ src/query/mod.rs | 2 + src/runtime/mod.rs | 1 + src/runtime/stdio.rs | 51 +++ src/typesafe.rs | 233 +++++++++++++ 23 files changed, 1701 insertions(+), 53 deletions(-) create mode 100644 plans/2026-09-22-ctxe-system-one.md create mode 100644 plans/260922-2052-turso-prototype/plan.md create mode 100644 plans/260922-2052-typesafe-api/plan.md create mode 100644 plans/reports/2026-09-22-ctxe-system-one-research.md create mode 100644 plans/reports/turso-260922-2052-assessment.md create mode 100644 src/prompts/mcp_ask_context.txt create mode 100644 src/query/decision.rs create mode 100644 src/query/hybrid.rs create mode 100644 src/runtime/stdio.rs create mode 100644 src/typesafe.rs diff --git a/Cargo.toml b/Cargo.toml index ec5d9d8..e3c4230 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -79,7 +79,7 @@ chrono = { version = "0.4", features = ["serde"] } regex = "1" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } -rmcp = { version = "1.7.0", features = ["server", "macros", "transport-streamable-http-server"] } +rmcp = { version = "1.7.0", features = ["server", "macros", "transport-io", "transport-streamable-http-server"] } tokio-util = "0.7" tokio-stream = { version = "0.1", features = ["sync"] } async-stream = "0.3" diff --git a/README-vi.md b/README-vi.md index bd1075e..eda655f 100644 --- a/README-vi.md +++ b/README-vi.md @@ -30,6 +30,17 @@ npm install -g vibervn-context-engine@latest vibervn-context-engine --port 6699 ``` +Để kết nối Codex, Claude Code, Cursor hoặc MCP client chạy dưới dạng subprocess, +dùng transport stdio: + +```bash +vibervn-context-engine --mcp-stdio +``` + +`ask-context` có thể gọi thêm System One của TypeSafe. Đặt `TYPESAFE_API_KEY` để +bật API `jev-latest`; có thể ghi đè bằng `TYPESAFE_API_URL`, `TYPESAFE_MODEL` và +`TYPESAFE_TIMEOUT_MS`. Nếu không có key, quyết định local vẫn hoạt động. + Nền tảng được hỗ trợ: Linux x64/arm64, macOS arm64, Windows x64. ## Tính năng @@ -37,6 +48,7 @@ Nền tảng được hỗ trợ: Linux x64/arm64, macOS arm64, Windows x64. | Tính năng | Mô tả | |-----------|-------| | Semantic code search | Tìm mã theo ý nghĩa thông qua embedding, không khớp văn bản thuần | +| Hybrid retrieval | Kết hợp vector semantic với lexical matching có giới hạn để bắt exact symbol, path và error string | | Multi-language parsing | Extract symbol bằng Tree-sitter cho 22 ngôn ngữ (xem bảng bên dưới) | | Call-graph expansion | Resolve caller/callee edge và BFS expand các symbol khớp khi query | | Import-path resolution | Trace import tới file thật cho TS/JS, Python, Go, Rust — resolve cross-module call mà name matching bỏ lỡ | @@ -48,9 +60,10 @@ Nền tảng được hỗ trợ: Linux x64/arm64, macOS arm64, Windows x64. | Real-time file watching | `notify` (debounce) tự động trigger re-index khi tệp thay đổi | | Voyage AI embedding | HTTP embedding client có disk cache để tránh gọi API thừa | | LLM rerank | Sắp xếp lại các candidate chunk bằng LLM (OpenAI / Google); tùy chọn, có thể tắt | +| Structured Ask decision | `ask-context` trả JSON gồm evidence, trạng thái source, confidence, uncertainty và cảnh báo degradation; TypeSafe bổ sung answer noul/choice/score khi được cấu hình | | Embedded SurrealDB | Lưu chunk, symbol và edge; một datastore cho mỗi repo | | HTTP API + Web UI | Quản lý cấu hình, index explorer và bảng điều khiển thử query | -| MCP server | Cung cấp `codebase-retrieval` và `file-retrieval` tool qua streamable HTTP | +| MCP server | Cung cấp retrieval dạng text và `ask-context` dạng cấu trúc qua streamable HTTP hoặc stdio | | SSE progress stream | Truyền sự kiện indexing progress trực tiếp tới UI | | Large-repo scaling | Bounded memory và không có đường O(n²) — xây dựng cho codebase quy mô Linux/Chromium | @@ -110,13 +123,15 @@ flowchart TD end Clients --> Q1[Embed query] - Q1 --> Q2[Vector search: top-k cosine] - Q2 -.reads.-> VecIndex + Q1 --> Q2[Hybrid retrieval: vector + lexical] + Q2 -.semantic reads.-> VecIndex + Q2 -.lexical reads.-> Store Q2 --> Q3[Graph expand: BFS caller/callee] Q3 --> Q4[Merge + dedup adjacent range] Q4 --> Q5[LLM rerank] Q5 --> Q6[Format: path#Lstart-end + dòng đánh số] - Q6 --> Result([Kết quả]) + Q6 --> Decision[Typed decision: action + confidence + evidence validity] + Decision --> Result([Kết quả]) ``` ## Đóng góp diff --git a/README-zh.md b/README-zh.md index 34d71f9..b11cb91 100644 --- a/README-zh.md +++ b/README-zh.md @@ -27,6 +27,16 @@ npm install -g vibervn-context-engine@latest vibervn-context-engine --port 6699 ``` +对于 Codex、Claude Code、Cursor 以及其他以子进程启动 MCP 的客户端，可使用 stdio： + +```bash +vibervn-context-engine --mcp-stdio +``` + +`ask-context` 可选调用 TypeSafe System One。设置 `TYPESAFE_API_KEY` 后启用 +`jev-latest`；也可通过 `TYPESAFE_API_URL`、`TYPESAFE_MODEL` 和 +`TYPESAFE_TIMEOUT_MS` 覆盖默认值。没有密钥时仍使用本地决策。 + 支持的平台：Linux x64/arm64、macOS arm64、Windows x64。 ## 功能特性 @@ -34,6 +44,7 @@ vibervn-context-engine --port 6699 | 功能 | 描述 | |------|------| | 语义代码搜索 | 通过嵌入向量按语义查找代码，而非字面文本匹配 | +| 混合检索 | 将向量检索与有界 lexical matching 结合，保留精确符号、路径和错误字符串 | | 多语言解析 | 使用 Tree-sitter 为 22 种语言提取符号（见下表） | | 调用图扩展 | 解析调用方/被调方边，并在查询时对匹配符号进行 BFS 扩展 | | Import 路径解析 | 为 TS/JS、Python、Go、Rust 追踪 import 到实际文件 —— 解析 name matching 遗漏的跨模块调用 | @@ -45,9 +56,10 @@ vibervn-context-engine --port 6699 | 实时文件监听 | `notify`（带去抖）在文件变更时自动触发重新索引 | | Voyage AI 嵌入 | 带磁盘缓存的 HTTP 嵌入客户端，避免重复 API 调用 | | LLM 重排序 | 使用 LLM（OpenAI / Google）对候选片段重新排序；可选，可禁用 | +| 结构化 Ask 决策 | `ask-context` 返回 JSON evidence、源码有效性、confidence、uncertainty 和降级警告；配置 TypeSafe 后还会返回 noul/choice/score typed answers | | 内嵌 SurrealDB | 存储片段、符号和边；每个仓库一个数据存储 | | HTTP API + Web 界面 | 配置管理、索引浏览器和查询测试控制台 | -| MCP 服务器 | 通过可流式 HTTP 暴露 `codebase-retrieval` 和 `file-retrieval` 工具 | +| MCP 服务器 | 通过可流式 HTTP 或 stdio 暴露文本检索和结构化 `ask-context` 工具 | | SSE 进度流 | 将实时索引进度事件流式传输到界面 | | 大型仓库扩展 | 内存有界且无 O(n²) 路径 —— 为 Linux/Chromium 规模的代码库而构建 | diff --git a/README.md b/README.md index 550b27a..69207ad 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,18 @@ npm install -g vibervn-context-engine@latest vibervn-context-engine --port 6699 ``` +For Codex, Claude Code, Cursor, and other subprocess-launched MCP clients, run +the stdio transport instead of the HTTP server: + +```bash +vibervn-context-engine --mcp-stdio +``` + +`ask-context` can optionally add TypeSafe System One judgments. Set +`TYPESAFE_API_KEY` to enable the real `jev-latest` API; `TYPESAFE_API_URL`, +`TYPESAFE_MODEL`, and `TYPESAFE_TIMEOUT_MS` are optional overrides. Without a +key, the local retrieval decision remains active. + Supported platforms: Linux x64/arm64, macOS arm64, Windows x64. ## Features @@ -36,6 +48,7 @@ Supported platforms: Linux x64/arm64, macOS arm64, Windows x64. | Feature | Description | |---------|-------------| | Semantic code search | Finds code by meaning via embeddings, not literal text matching | +| Hybrid retrieval | Fuses semantic vectors with bounded lexical matching for exact symbols, paths, and error strings | | Multi-language parsing | Tree-sitter symbol extraction for 22 languages (see table below) | | Call-graph expansion | Resolves caller/callee edges and BFS-expands matched symbols at query time | | Import-path resolution | Traces imports to actual files for TS/JS, Python, Go, and Rust — resolves cross-module calls that name matching misses | @@ -47,9 +60,10 @@ Supported platforms: Linux x64/arm64, macOS arm64, Windows x64. | Real-time file watching | `notify` (debounced) triggers re-index automatically on file changes | | Voyage AI embeddings | HTTP embedding client with an on-disk cache to avoid redundant API calls | | LLM reranking | Reorders candidate chunks with an LLM (OpenAI / Google); optional, can be disabled | +| Structured Ask decisions | `ask-context` returns JSON evidence, source validity, confidence, uncertainty, and degradation warnings; TypeSafe adds typed noul/choice/score answers when configured | | Embedded SurrealDB | Stores chunks, symbols, and edges; one datastore per repo | | HTTP API + Web UI | Settings management, index explorer, and a query test console | -| MCP server | Exposes `codebase-retrieval` and `file-retrieval` tools over streamable HTTP | +| MCP server | Exposes text retrieval plus structured `ask-context` over streamable HTTP or stdio | | SSE progress stream | Streams live indexing progress events to the UI | | Large-repo scaling | Bounded memory and no O(n²) paths — built for Linux/Chromium-scale codebases | @@ -111,14 +125,16 @@ flowchart TD Clients --> QF[Parse field filters: kind:/lang:/path:/name:] QF --> Q1[Embed remaining query text] - Q1 --> Q2[Vector search: top-k cosine] - Q2 -.reads.-> VecIndex + Q1 --> Q2[Hybrid retrieval: vector + lexical] + Q2 -.semantic reads.-> VecIndex + Q2 -.lexical reads.-> Store Q2 --> Q3[Apply filters + downrank generated files] Q3 --> Q4[Graph expand: BFS callers/callees] Q4 --> Q5[Merge + dedup adjacent ranges] Q5 --> Q6[LLM rerank] Q6 --> Q7[Format: path#Lstart-end + caller/callee names + numbered lines] - Q7 --> Result([Results]) + Q7 --> Decision[Typed decision: action + confidence + evidence validity] + Decision --> Result([Results]) ``` ## Contributing diff --git a/plans/2026-09-22-ctxe-system-one.md b/plans/2026-09-22-ctxe-system-one.md new file mode 100644 index 0000000..6ad6aaa --- /dev/null +++ b/plans/2026-09-22-ctxe-system-one.md @@ -0,0 +1,57 @@ +# CtxE and System One Integration + +Status: completed + +## Outcome + +Improve `vibervn-context-engine` with the strongest compatible ideas found in +CtxE and TypeSafe's System One announcement: hybrid semantic and lexical +retrieval, a goal-driven structured Ask surface, source-evidence validity, and +typed local confidence decisions that clients can consume without parsing prose. + +## Constraints + +- Preserve the existing HTTP UI, streamable HTTP MCP endpoint, indexing model, + and existing `codebase-retrieval` response contract. +- Keep all retrieval bounded and repository-scoped. +- Do not pretend to call Jev or reproduce proprietary calibration; expose an + explicit local heuristic decision with its uncertainty and signals. +- Keep optional LLM/network work outside database write transactions and degrade + to local retrieval when optional enrichment is unavailable. + +## Non-goals + +- Replacing the existing embedding or LLM providers. +- Adding a hosted TypeSafe dependency without a documented public API. +- Rebuilding the entire CtxE records database or daemon/service manager. + +## Acceptance criteria + +1. Queries combine vector candidates with bounded lexical candidates and keep + exact symbol/path matches useful when embeddings miss them. +2. Query JSON includes a typed decision (`action`, confidence, uncertainty, + evidence coverage, and reasons) and evidence records that distinguish + unchanged, changed, partial, missing, and unavailable source content. +3. MCP exposes a default `ask-context` tool returning structured JSON while + retaining the existing text retrieval tools. +4. The binary supports a stdio MCP mode suitable for Codex, Claude Code, + Cursor, and other subprocess-launched clients. +5. Focused unit/integration checks cover lexical scoring, decision outcomes, + evidence classification, config migration, and MCP tool wiring; the normal + Rust checks pass. + +## Sources + +- CtxE public contract: `https://ctxe.tlelabs.com/llms.txt` and its current docs + pages for hybrid retrieval, Ask, graph tools, evidence validity, and stdio MCP. +- TypeSafe AI, “Introducing System One Models & Jev”, published 2026-09-15: + `https://typesafe.ai/blog/introducing-system-one-models-and-jev`. + +## Verification + +- `cargo test --lib`: 665 passed, 0 failed, 6 ignored. +- The `ask-context` schema and native MCP `structuredContent` tests are included in that run. +- `cargo test --test integration --test router_integration --test mcp_session_restore`: 33 passed. +- `cargo check --all-targets` and `cargo run -- --help`: passed. +- A temp-home MCP stdio initialize handshake returned a valid JSON-RPC response. +- `cargo clippy --all-targets -- -D warnings` still reports five pre-existing findings in `src/embedding/cache.rs`, `src/store/ops.rs`, `src/router/plan.rs`, and `src/server.rs`; no new findings remain in the changed code. diff --git a/plans/260922-2052-turso-prototype/plan.md b/plans/260922-2052-turso-prototype/plan.md new file mode 100644 index 0000000..dab9cf2 --- /dev/null +++ b/plans/260922-2052-turso-prototype/plan.md @@ -0,0 +1,26 @@ +# Turso Experimental Backend Prototype + +Status: planned + +## Outcome + +Benchmark a Turso-backed code index against the existing SurrealDB/vector-shard +path without changing the production default or deleting existing data. + +## Boundaries + +- Additive feature flag and separate database file only. +- No automatic migration, destructive schema change, or default backend switch. +- Reuse the existing parser/chunker/symbol/call-edge output. +- Keep TypeSafe and MCP behavior independent of the storage choice. + +## Acceptance criteria + +1. A fixture repository can be imported idempotently into Turso tables for files, + chunks, symbols, and edges. +2. Turso FTS and exact vector queries return a common retrieval record shape. +3. A benchmark compares both backends on recall@k, p50/p95 latency, indexing + time, memory, and recovery after an interrupted update. +4. Existing SurrealDB tests remain unchanged and green. +5. Promotion remains a measured decision; the prototype never silently becomes + the default backend. diff --git a/plans/260922-2052-typesafe-api/plan.md b/plans/260922-2052-typesafe-api/plan.md new file mode 100644 index 0000000..5199180 --- /dev/null +++ b/plans/260922-2052-typesafe-api/plan.md @@ -0,0 +1,36 @@ +# TypeSafe API Integration + +Status: completed + +## Outcome + +Use the public TypeSafe System One HTTP API from the structured `ask-context` +workflow when `TYPESAFE_API_KEY` is configured, while preserving local retrieval +and degradation when the key or service is unavailable. + +## Scope + +- Add a small Rust client for `https://api.typesafe.ai/v1/systemone`. +- Ask bounded `noul`, `choice`, and `score` questions over retrieved evidence. +- Return the typed TypeSafe answers, model, and usage in `ask-context` JSON and + MCP `structuredContent` without exposing credentials. +- Keep all TypeSafe calls outside database transactions and cap request state. + +## Acceptance criteria + +1. `TYPESAFE_API_KEY` enables a real authenticated request using `jev-latest`. +2. Missing keys and non-success responses degrade to local Ask output with a + warning and no secret values in logs or responses. +3. Valid TypeSafe responses preserve typed answers and usage metadata. +4. Unit tests cover request shape, response decoding, error redaction, and the + no-key path; existing library/integration checks remain green. + +## Verification + +- `cargo test --lib`: 668 passed, 0 failed, 6 ignored, including the TypeSafe + request/response tests. +- `cargo test --test integration --test router_integration --test mcp_session_restore`: + 33 passed. +- `cargo check --all-targets`: passed. +- TypeSafe failures are logged only as status-level errors and return local Ask + output; response bodies and credentials are never copied into errors. diff --git a/plans/reports/2026-09-22-ctxe-system-one-research.md b/plans/reports/2026-09-22-ctxe-system-one-research.md new file mode 100644 index 0000000..0ffcf6b --- /dev/null +++ b/plans/reports/2026-09-22-ctxe-system-one-research.md @@ -0,0 +1,53 @@ +# CtxE and System One Research + +Conducted: 2026-09-22 UTC + +## Findings + +CtxE's public documentation describes a local, syntax-aware index with a +two-phase publish pipeline, hybrid semantic plus indexed-text retrieval, typed +symbol relationships, a goal-driven multi-round Ask flow, source snapshots that +report evidence validity, and explicit degradation/truncation diagnostics. Its +MCP integration is stdio-first and exposes separate status, tree, retrieval, +definition, usage, graph, impact, and record tools. + +The current repository already has strong local foundations: Tree-sitter +parsing, incremental indexing, a persisted call graph, vector shards with lazy +warming, optional LLM/agentic reranking, freshness gates, and a streamable +HTTP MCP service. The main gaps are exact lexical recall, machine-readable +confidence/evidence metadata, a structured Ask entry point, and stdio MCP +transport. + +TypeSafe's article presents System One models as fast decision functions: typed +outputs selected from a fixed schema, probabilities/confidence on every answer, +parallel sampling, and a preference for composing many small decisions inside +ordinary software workflows. The article itself did not include the API +contract, so the initial implementation used a typed local decision layer; +follow-up documentation verification now supports the optional API client +described below. + +## Implementation choice + +The change keeps vector search as the semantic backbone, adds bounded lexical +candidate lookup against the already-indexed chunk store, and fuses both scores +before graph expansion and reranking. Query results gain structured evidence +and decision metadata. `ask-context` returns that data as JSON, while the +existing text tool remains unchanged for clients that need concise snippets. + +## Follow-up API verification + +The current TypeSafe documentation now publishes a stable HTTP contract at +`POST https://api.typesafe.ai/v1/systemone` with `Bearer` authentication, +`jev-latest`, and the `noul`, `choice`, and `score` question primitives. The +follow-up implementation uses `TYPESAFE_API_KEY` and keeps the integration +optional: a missing key or non-success response falls back to the local typed +decision without copying credentials or provider response bodies into logs. + +## References + +- CtxE overview and docs: `https://ctxe.tlelabs.com/` +- CtxE machine-readable index: `https://ctxe.tlelabs.com/llms.txt` +- CtxE retrieval/evidence docs: `https://ctxe.tlelabs.com/docs/how-it-works` +- CtxE MCP contract: `https://ctxe.tlelabs.com/docs/mcp-integration` +- CtxE tools: `https://ctxe.tlelabs.com/docs/tools-reference` +- TypeSafe article: `https://typesafe.ai/blog/introducing-system-one-models-and-jev` diff --git a/plans/reports/turso-260922-2052-assessment.md b/plans/reports/turso-260922-2052-assessment.md new file mode 100644 index 0000000..3612633 --- /dev/null +++ b/plans/reports/turso-260922-2052-assessment.md @@ -0,0 +1,57 @@ +# Turso Integration Assessment + +Conducted: 2026-09-22 UTC + +## Decision + +Use Turso as an optional experimental storage and benchmark backend, not as an +immediate replacement for SurrealDB. + +The current engine already has large-repository behavior that would be risky to +discard: per-repository SurrealDB handles, persisted vector shards, mmap warmup, +quantized resident vectors, LRU eviction, mutation fences, and extensive +incremental/crash-recovery tests. The Turso path should prove equivalent +behavior and performance before it becomes the default. + +## Evidence + +Turso's current public documentation provides a code-indexing schema that maps +well to this repository: indexed files with content hashes, chunks with names, +signatures, line ranges and embeddings, FTS indexes over identifiers, and vector +distance queries. It also supports `BEGIN CONCURRENT` with MVCC for independent +writes. + +The important limits are equally clear: + +- FTS is powered by Tantivy and requires the experimental `index_method` feature. +- Vector similarity queries are exact linear scans; approximate vector indexing + remains on the roadmap. +- FTS and vector syntax/API may change while those features are experimental. +- Graph edges can be represented in relational tables, but traversal, cleanup, + and incremental blast-radius logic remain application-owned. + +Sources: + +- Turso code indexing guide: `https://docs.turso.tech/guides/code-indexing` +- Turso vector search: `https://docs.turso.tech/guides/vector-search` +- Turso FTS reference: `https://docs.turso.tech/sql-reference/functions/fts` +- Turso Rust reference: `https://docs.turso.tech/sdk/rust/reference` +- Turso concurrent writes: `https://docs.turso.tech/tursodb/concurrent-writes` +- Turso repository: `https://github.com/tursodatabase/turso` + +## Recommended prototype + +1. Add an opt-in `turso-experimental` feature and a separate local database + file; do not alter the existing SurrealDB path. +2. Import one indexed repository into `codebases`, `indexed_files`, `chunks`, + `symbols`, and `edges`, using `file_hash`/`chunk_key` for idempotent upserts. +3. Create weighted FTS over `name`, `signature`, and a bounded snippet field; + store vectors as `vector8` or `vector32` and compare exact-search latency to + the current resident shard. +4. Run the same query corpus through SurrealDB and Turso and record indexing + time, query p50/p95, recall@k, memory, FTS freshness, and crash recovery. +5. Promote only if Turso matches the current correctness suite and wins on a + measured workload; otherwise retain it as an export/portable-index option. + +No database migration has been applied. If the prototype becomes a migration, +the repository must first create a durable backup and define a rollback path. diff --git a/src/config.rs b/src/config.rs index e4153ed..7c31f27 100644 --- a/src/config.rs +++ b/src/config.rs @@ -9,7 +9,7 @@ use serde_json::Value; use tempfile::NamedTempFile; /// Bump this when a new migration is appended to MIGRATIONS. -pub const CURRENT_VERSION: u32 = 13; +pub const CURRENT_VERSION: u32 = 14; /// Migration function type: transforms a JSON Value from version N to version N+1. pub type MigrationFn = fn(Value) -> Result; @@ -29,6 +29,7 @@ pub const MIGRATIONS: &[MigrationFn] = &[ migrate_v10_to_v11, migrate_v11_to_v12, migrate_v12_to_v13, + migrate_v13_to_v14, ]; /// v1→v2: introduce `data_dir` (Option). The body is a no-op stamp — @@ -200,6 +201,35 @@ fn migrate_v12_to_v13(mut value: Value) -> Result { Ok(value) } +/// v13→v14: expose the structured Ask workflow by default. Existing users who +/// had an explicit MCP allowlist keep their choices plus the new read-only tool; +/// fresh installs get it from `default_enabled_mcp_tools`. +fn migrate_v13_to_v14(mut value: Value) -> Result { + if let Value::Object(ref mut obj) = value { + match obj.get_mut("enabled_mcp_tools") { + Some(Value::Array(tools)) => { + if !tools.is_empty() + && !tools + .iter() + .any(|tool| tool.as_str() == Some("ask-context")) + { + tools.push(Value::String("ask-context".to_string())); + } + } + _ => { + obj.insert( + "enabled_mcp_tools".to_string(), + Value::Array(vec![ + Value::String("codebase-retrieval".to_string()), + Value::String("ask-context".to_string()), + ]), + ); + } + } + } + Ok(value) +} + // ─── Settings ────────────────────────────────────────────────────────────── fn default_index_ignore_filenames() -> Vec { @@ -430,9 +460,9 @@ fn default_mcp_stale_after_days() -> u64 { } fn default_enabled_mcp_tools() -> Vec { - // Only `codebase-retrieval` is on by default. `file-retrieval` is an - // advanced, opt-in tool — new installs must enable it explicitly in the UI. - vec!["codebase-retrieval".to_string()] + // `file-retrieval` remains advanced and opt-in. The structured Ask tool is + // read-only and is the default machine-readable entry point for new installs. + vec!["codebase-retrieval".to_string(), "ask-context".to_string()] } /// A plan/key the user has bought (or claimed as a free trial) through the buy @@ -1843,8 +1873,8 @@ mod tests { assert_eq!(loaded.version, CURRENT_VERSION); assert_eq!( loaded.enabled_mcp_tools, - vec!["codebase-retrieval".to_string()], - "file-retrieval must be removed from enabled tools on upgrade" + vec!["codebase-retrieval".to_string(), "ask-context".to_string()], + "file-retrieval must be removed and structured Ask enabled on upgrade" ); let raw = fs::read_to_string(&path).expect("re-read"); @@ -1880,12 +1910,15 @@ mod tests { assert_eq!(tools, vec!["codebase-retrieval".to_string()]); } - /// New installs default to `codebase-retrieval` only; `file-retrieval` is - /// opt-in and must NOT appear in the default tool list. + /// New installs enable `codebase-retrieval` and structured Ask; + /// `file-retrieval` remains opt-in and must NOT appear in the default list. #[test] fn test_default_enabled_mcp_tools_excludes_file_retrieval() { let tools = default_enabled_mcp_tools(); - assert_eq!(tools, vec!["codebase-retrieval".to_string()]); + assert_eq!( + tools, + vec!["codebase-retrieval".to_string(), "ask-context".to_string()] + ); assert!( !tools.iter().any(|t| t == "file-retrieval"), "file-retrieval must not be enabled by default" @@ -1983,9 +2016,30 @@ mod tests { ); } + #[test] + fn test_v13_to_v14_adds_structured_ask_without_clobbering_allowlists() { + let value: Value = + serde_json::from_str(r#"{"enabled_mcp_tools":["codebase-retrieval","custom-tool"]}"#) + .expect("parse"); + let migrated = migrate_v13_to_v14(value).expect("migrate"); + let tools = migrated["enabled_mcp_tools"].as_array().expect("array"); + assert_eq!(tools[0], "codebase-retrieval"); + assert_eq!(tools[1], "custom-tool"); + assert_eq!(tools[2], "ask-context"); + + let empty: Value = serde_json::from_str(r#"{"enabled_mcp_tools":[]}"#).expect("parse"); + let empty = migrate_v13_to_v14(empty).expect("migrate empty"); + assert!( + empty["enabled_mcp_tools"] + .as_array() + .expect("array") + .is_empty() + ); + } + /// v12→v13 with a MISSING `enabled_mcp_tools` field: the migration is a - /// no-op (nothing to strip), and serde fills the new opt-in default - /// (codebase-retrieval only) on deserialize. + /// no-op (nothing to strip), and serde fills the structured Ask default on + /// deserialize. #[test] fn test_v12_to_v13_missing_field_defaults_to_codebase_only() { let home = TempDir::new().expect("tempdir"); @@ -2004,8 +2058,8 @@ mod tests { let loaded = ensure_dir_and_load(home.path()).expect("load v12 missing tools"); assert_eq!( loaded.enabled_mcp_tools, - vec!["codebase-retrieval".to_string()], - "missing enabled_mcp_tools must default to codebase-retrieval only" + vec!["codebase-retrieval".to_string(), "ask-context".to_string()], + "missing enabled_mcp_tools must include the default structured Ask tool" ); } @@ -2035,7 +2089,7 @@ mod tests { ); } - /// One-time-only guard: once a file is at CURRENT_VERSION (v13), loading it + /// One-time-only guard: once a file is at CURRENT_VERSION, loading it /// does NOT re-run the migration. A user who deliberately RE-ENABLES /// file-retrieval after the upgrade keeps it — the migration never strips it /// again. Also proves the on-disk file is untouched (no rewrite) on a diff --git a/src/indexing/pipeline.rs b/src/indexing/pipeline.rs index ed227da..792f962 100644 --- a/src/indexing/pipeline.rs +++ b/src/indexing/pipeline.rs @@ -1348,6 +1348,12 @@ impl IndexPipeline { return Err(e); } + // The changed-file edge set is resolved and unaffected calls remain + // intact, so the committed index has a complete graph marker. + set_meta(db, EDGES_RESOLVED_KEY, "1") + .await + .context("commit edges_resolved marker")?; + Ok((run_stats, vi_apply_ms)) } @@ -2972,11 +2978,10 @@ impl IndexPipeline { /// 3. Re-resolve raw_edge rows WHERE from_file IN resolve_set via keyset /// pagination (uses idx_raw_edge_from_file). /// - /// The `edges_resolved` crash-recovery marker is NOT written here — it is only - /// meaningful for a full rebuild where ALL raw_edge must be re-resolved on crash - /// recovery. Incremental is already idempotent: if it crashes before file_meta - /// is written (the crash-safe anchor in streaming_index), the whole incremental - /// re-runs on next trigger, including this method. + /// The `edges_resolved` marker is stamped by `incremental_run` only after this + /// scoped resolution succeeds. Incremental is idempotent: if it crashes before + /// file_meta is written (the crash-safe anchor in `streaming_index`), the whole + /// incremental re-runs on the next trigger, including this method. async fn resolve_edges_incremental( &self, db: &Surreal, diff --git a/src/lib.rs b/src/lib.rs index be6a9ba..5b67a73 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -18,6 +18,7 @@ pub mod query; pub mod router; pub mod server; pub mod store; +pub mod typesafe; pub mod vector; /// Check whether `file` belongs to the repository rooted at `repo`. diff --git a/src/main.rs b/src/main.rs index 8c71a23..c8d08a5 100644 --- a/src/main.rs +++ b/src/main.rs @@ -36,12 +36,26 @@ struct Cli { /// Worker idle window before scale-to-zero. Ignored in router mode. #[arg(long, env = "CONTEXT_ENGINE_WORKER_IDLE_SECS")] worker_idle_secs: Option, + + /// Run the MCP server over stdin/stdout for subprocess-launched clients. + /// This mode is mutually exclusive with the process-per-project worker. + #[arg(long)] + mcp_stdio: bool, } #[tokio::main] async fn main() { let cli = Cli::parse(); - init_tracing(cli.worker.is_some()); + if cli.mcp_stdio && cli.worker.is_some() { + eprintln!("error: --mcp-stdio cannot be combined with --worker"); + std::process::exit(2); + } + init_tracing(cli.worker.is_some(), cli.mcp_stdio); + + if cli.mcp_stdio { + runtime::stdio::run(&cli).await; + return; + } let bind = cli.bind.as_deref().unwrap_or("127.0.0.1").to_owned(); match cli.worker.clone() { @@ -50,11 +64,11 @@ async fn main() { } } -fn init_tracing(worker_mode: bool) { +fn init_tracing(worker_mode: bool, stdio_mode: bool) { let filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| EnvFilter::new("context_engine_rs=info,warn")); - if worker_mode { - // Worker stdout is reserved for the readiness handshake. + if worker_mode || stdio_mode { + // Worker/MCP stdout is reserved for a machine-readable protocol. tracing_subscriber::fmt() .with_env_filter(filter) .with_writer(std::io::stderr) diff --git a/src/mcp.rs b/src/mcp.rs index 5adb974..4f60a5d 100644 --- a/src/mcp.rs +++ b/src/mcp.rs @@ -308,6 +308,16 @@ pub struct FileRetrievalArgs { pub top_k: Option, } +#[derive(Debug, serde::Deserialize, schemars::JsonSchema)] +pub struct AskContextArgs { + /// Absolute path to the repository root. + pub workspace_full_path: String, + /// A concrete question or investigation objective. + pub query: String, + /// Bounded retrieval effort. Defaults to `min`. + pub effort: Option, +} + // ─── MCP handler ───────────────────────────────────────────────────────── #[derive(Clone)] @@ -337,7 +347,7 @@ impl McpHandler { settings: Arc>, enabled_tools: &[String], ) -> Self { - let all_tools: &[&str] = &["codebase-retrieval", "file-retrieval"]; + let all_tools: &[&str] = &["codebase-retrieval", "file-retrieval", "ask-context"]; let mut router = Self::tool_router(); for &name in all_tools { if !enabled_tools.iter().any(|e| e == name) { @@ -415,6 +425,34 @@ impl McpHandler { .await; Ok(CallToolResult::success(vec![Content::text(text)])) } + + #[doc = include_str!("prompts/mcp_ask_context.txt")] + #[tool(name = "ask-context")] + async fn ask_context( + &self, + Parameters(args): Parameters, + ctx: RequestContext, + ) -> Result { + let settings = self.settings.read().await.clone(); + let text = with_progress_heartbeat( + ctx.peer, + &ctx.meta, + ctx.ct, + MCP_PROGRESS_HEARTBEAT, + run_ask_context( + &self.home_dir, + &self.data_dir, + &self.index_engine, + &self.repo_dbs, + &settings, + &args.query, + &args.workspace_full_path, + args.effort.as_deref().unwrap_or("min"), + ), + ) + .await; + Ok(ask_context_call_result(text)) + } } #[tool_handler(router = self.tool_router)] @@ -446,6 +484,14 @@ pub struct RepoFileRetrievalArgs { pub top_k: Option, } +#[derive(Debug, serde::Deserialize, schemars::JsonSchema)] +pub struct RepoAskContextArgs { + /// A concrete question or investigation objective. + pub query: String, + /// Bounded retrieval effort. Defaults to `min`. + pub effort: Option, +} + #[derive(Clone)] pub struct RepoMcpHandler { home_dir: PathBuf, @@ -469,7 +515,7 @@ impl RepoMcpHandler { settings: Arc>, enabled_tools: &[String], ) -> Self { - let all_tools: &[&str] = &["codebase-retrieval", "file-retrieval"]; + let all_tools: &[&str] = &["codebase-retrieval", "file-retrieval", "ask-context"]; let mut router = Self::tool_router(); for &name in all_tools { if !enabled_tools.iter().any(|e| e == name) { @@ -540,6 +586,34 @@ impl RepoMcpHandler { .await; Ok(CallToolResult::success(vec![Content::text(text)])) } + + #[doc = include_str!("prompts/mcp_ask_context.txt")] + #[tool(name = "ask-context")] + async fn ask_context( + &self, + Parameters(args): Parameters, + ctx: RequestContext, + ) -> Result { + let settings = self.settings.read().await.clone(); + let text = with_progress_heartbeat( + ctx.peer, + &ctx.meta, + ctx.ct, + MCP_PROGRESS_HEARTBEAT, + run_ask_context( + &self.home_dir, + &self.data_dir, + &self.index_engine, + &self.repo_dbs, + &settings, + &args.query, + &self.repo_path, + args.effort.as_deref().unwrap_or("min"), + ), + ) + .await; + Ok(ask_context_call_result(text)) + } } #[tool_handler(router = self.tool_router)] @@ -727,6 +801,245 @@ fn build_augmented_query( } } +#[derive(Debug, serde::Serialize)] +struct AskContextResponse { + schema_version: u32, + workspace: String, + query: String, + effort: String, + results: Vec, + evidence: Vec, + decision: crate::query::decision::RetrievalDecision, + typesafe: Option, + timing: crate::query::engine::QueryTiming, + warnings: Vec, +} + +/// Preserve the human-readable JSON text while also filling MCP's native +/// `structuredContent` field for clients that can consume typed tool results. +fn ask_context_call_result(text: String) -> CallToolResult { + match serde_json::from_str::(&text) { + Ok(value) if value.get("error").is_some() => CallToolResult::structured_error(value), + Ok(value) => CallToolResult::structured(value), + Err(_) => CallToolResult::error(vec![Content::text(text)]), + } +} + +fn build_typesafe_state( + query: &str, + result: &crate::query::engine::QueryResult, +) -> serde_json::Value { + let results: Vec = result + .results + .iter() + .take(12) + .map(|item| { + serde_json::json!({ + "file": item.file, + "line_start": item.line_start, + "line_end": item.line_end, + "score": item.score, + "symbol": item.symbol, + "content": item.content.chars().take(1200).collect::(), + }) + }) + .collect(); + crate::typesafe::bounded_state( + serde_json::json!({ + "query": query, + "results": results, + "evidence": result.evidence, + "local_decision": result.decision, + }), + 24_000, + ) +} + +/// Structured, machine-readable Ask workflow. The response is intentionally a +/// fixed JSON schema so an agent can branch on `decision.action` and confidence +/// without parsing the legacy text output from `codebase-retrieval`. +#[allow(clippy::too_many_arguments)] +pub async fn run_ask_context( + home_dir: &Path, + data_dir: &Path, + index_engine: &Arc, + repo_dbs: &Arc>>>, + settings: &Settings, + query: &str, + workspace_full_path: &str, + effort: &str, +) -> String { + let query = query.trim(); + let repo = workspace_full_path.trim(); + if query.is_empty() { + return serde_json::json!({ "error": "query is required" }).to_string(); + } + if repo.is_empty() { + return serde_json::json!({ "error": "workspace_full_path is required" }).to_string(); + } + let effort = effort.trim().to_lowercase(); + let top_k = match effort.as_str() { + "min" => 12, + "medium" => 24, + "high" => 40, + _ => { + return serde_json::json!({ + "error": "effort must be one of: min, medium, high" + }) + .to_string(); + } + }; + + let repo = crate::store::normalize_repo_path(repo); + if !settings.repos.iter().any(|configured| configured == &repo) { + if !Path::new(&repo).is_dir() { + return serde_json::json!({ + "error": format!("workspace '{repo}' does not exist or is not a directory") + }) + .to_string(); + } + // Match codebase-retrieval's ergonomic auto-registration. The index + // watcher is started immediately; persistence is best-effort. + index_engine.register_repo(&repo).await; + if let Ok(mut disk) = crate::config::ensure_dir_and_load(home_dir) + && !disk.repos.iter().any(|configured| configured == &repo) + { + disk.repos.push(repo.clone()); + disk.version = crate::config::CURRENT_VERSION; + let target = crate::config::config_path(home_dir); + let _ = crate::config::write_settings_atomic(&target, &disk); + } + } + + if settings.embedding.api_keys.is_empty() { + return serde_json::json!({ + "error": "no embedding API keys configured" + }) + .to_string(); + } + + let (graph_mode, warm_budget) = + match readiness::await_index_ready(settings, index_engine, repo_dbs, data_dir, &repo).await + { + readiness::IndexReadiness::Ready { warm_budget } => (QueryGraphMode::Full, warm_budget), + readiness::IndexReadiness::ReadyVectorOnly { warm_budget } => { + (QueryGraphMode::VectorOnly, warm_budget) + } + readiness::IndexReadiness::Timeout => { + return serde_json::json!({ + "error": "index is still warming; retry this request", + "decision": { "action": "retry_index" } + }) + .to_string(); + } + readiness::IndexReadiness::Failed(error) => { + return serde_json::json!({ "error": format!("index readiness failed: {error:#}") }) + .to_string(); + } + }; + + let voyage_client = match VoyageClient::new_for_provider( + crate::embedding::voyage::Provider::parse(&settings.embedding.provider), + settings.embedding.model.clone(), + settings.embedding.api_keys.clone(), + settings.embedding.voyage_base_url.as_deref(), + settings.embedding.dimensions, + ) { + Ok(client) => client, + Err(error) => { + return serde_json::json!({ + "error": format!("failed to create embedding client: {error}") + }) + .to_string(); + } + }; + let llm_client = LlmClient::new(&settings.llm); + + let result = match crate::query::engine::run_query_with_filters_and_mode( + query, + top_k, + Some(&repo), + &voyage_client, + index_engine, + repo_dbs, + settings.llm.rerank_min_prune_lines, + llm_client.as_ref(), + warm_budget, + settings.llm.agentic_rag, + settings.llm.agentic_rag_max_turns, + settings.llm.agentic_rag_max_chunk_chars, + settings.llm.agentic_rag_grep_read, + None, + graph_mode, + ) + .await + { + Ok(result) => result, + Err(error) => { + return serde_json::json!({ "error": format!("query failed: {error}") }).to_string(); + } + }; + + let mut warnings = Vec::new(); + let typesafe = match crate::typesafe::TypeSafeClient::from_env() { + Ok(Some(client)) => match client + .evaluate( + build_typesafe_state(query, &result), + crate::typesafe::ask_questions(), + ) + .await + { + Ok(response) => Some(response), + Err(error) => { + tracing::warn!(error = %error, "TypeSafe enrichment unavailable; using local Ask decision"); + warnings.push( + "TypeSafe enrichment was unavailable; local decision was used".to_owned(), + ); + None + } + }, + Ok(None) => None, + Err(error) => { + tracing::warn!(error = %error, "TypeSafe configuration invalid; using local Ask decision"); + warnings.push("TypeSafe configuration was invalid; local decision was used".to_owned()); + None + } + }; + if result.warming { + warnings.push("the index is warming; retry for a complete result set".to_owned()); + } + if result.graph_pending { + warnings.push( + "call-graph enrichment is pending; vector and lexical evidence is still valid" + .to_owned(), + ); + } + if result + .rerank + .as_ref() + .is_some_and(|rerank| rerank.fallback_used) + { + warnings.push("optional reranking was unavailable; local ranking was used".to_owned()); + } + + let response = AskContextResponse { + schema_version: 1, + workspace: repo, + query: query.to_owned(), + effort, + results: result.results, + evidence: result.evidence, + decision: result.decision, + typesafe, + timing: result.timing, + warnings, + }; + serde_json::to_string_pretty(&response).unwrap_or_else(|error| { + serde_json::json!({ "error": format!("failed to encode Ask response: {error}") }) + .to_string() + }) +} + /// Format an enriched caller tag: `[callers: fn_a, fn_b, fn_c +N more]` /// When callers > 3, shows first 3 names + count of remaining. /// Returns empty string when no callers. @@ -893,11 +1206,18 @@ async fn do_query( fn build_db_key(workspace: &str, file_path: &str) -> String { let workspace = workspace.trim_end_matches(['/', '\\']); let file_path = file_path.trim_start_matches(['/', '\\']); - let file_path_native = if cfg!(windows) { + // Preserve the path family supplied by the client. MCP callers can ask for + // a Windows workspace while the service itself is running on another host + // (and the indexed database may have been produced by that client). + let windows_style = workspace.contains('\\') || workspace.as_bytes().get(1) == Some(&b':'); + let file_path_native = if windows_style { file_path.replace('/', "\\") } else { file_path.replace('\\', "/") }; + if windows_style { + return format!("{workspace}\\{file_path_native}"); + } let repo_path = std::path::Path::new(workspace); let abs_file = repo_path.join(&file_path_native); abs_file.to_string_lossy().to_string() diff --git a/src/mcp/tests.rs b/src/mcp/tests.rs index 4d67446..8c013d9 100644 --- a/src/mcp/tests.rs +++ b/src/mcp/tests.rs @@ -28,6 +28,25 @@ fn vector_only_fast_path_only_for_resolve_edges() { assert!(!readiness::is_resolve_edges_status(None)); } +#[test] +fn ask_context_arguments_are_machine_readable() { + let args: AskContextArgs = serde_json::from_value(serde_json::json!({ + "workspace_full_path": "/repo", + "query": "where is authentication handled?", + "effort": "medium" + })) + .expect("ask-context arguments should decode"); + assert_eq!(args.workspace_full_path, "/repo"); + assert_eq!(args.effort.as_deref(), Some("medium")); +} + +#[test] +fn ask_context_result_populates_native_structured_content() { + let result = ask_context_call_result(r#"{"schema_version":1}"#.to_owned()); + assert_eq!(result.is_error, Some(false)); + assert!(result.structured_content.is_some()); +} + #[test] fn file_retrieval_db_key_windows_backslash_input() { let repo = r"D:\projects\Python\local-context-engine"; diff --git a/src/prompts/mcp_ask_context.txt b/src/prompts/mcp_ask_context.txt new file mode 100644 index 0000000..69d6a4e --- /dev/null +++ b/src/prompts/mcp_ask_context.txt @@ -0,0 +1,5 @@ +Use this tool for broad codebase questions, architecture investigations, impact analysis, and comparisons. +It returns JSON with `results`, source `evidence`, and a typed `decision` containing an action, confidence, uncertainty, and the signals behind that decision. +Treat `confidence_kind=local_heuristic` as a routing signal rather than a calibrated probability. If the action is `need_more_evidence`, follow up with targeted retrieval or a direct file read. +The `effort` value controls the bounded candidate budget: `min`, `medium`, or `high`. +When configured with `TYPESAFE_API_KEY`, the response also includes TypeSafe's typed `noul`, `score`, and `choice` judgments under `typesafe`; this enrichment is optional and local evidence remains the fallback. diff --git a/src/query/decision.rs b/src/query/decision.rs new file mode 100644 index 0000000..1cdcb84 --- /dev/null +++ b/src/query/decision.rs @@ -0,0 +1,306 @@ +//! Typed retrieval decisions and source-evidence validity. +//! +//! This is the local equivalent of a System One decision boundary: the engine +//! emits a fixed schema with explicit uncertainty instead of forcing clients to +//! infer confidence from prose. The scores are intentionally labelled as local +//! heuristics; they are useful routing signals, not claims about Jev's or an +//! external model's calibration. + +use std::collections::HashMap; +use std::path::Path; + +use chrono::Utc; +use serde::Serialize; +use sha2::{Digest, Sha256}; + +use crate::query::engine::CodeResult; +use crate::query::merger::MergeChunk; + +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum EvidenceStatus { + Unchanged, + Changed, + Partial, + Missing, + Unavailable, +} + +#[derive(Debug, Clone, Serialize)] +pub struct EvidenceRecord { + pub file: String, + pub line_start: u32, + pub line_end: u32, + pub status: EvidenceStatus, + pub indexed_sha256: Option, + pub current_sha256: Option, + pub reason: String, + pub captured_at: String, +} + +#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum DecisionAction { + Answer, + NeedMoreEvidence, + RetryIndex, + NoRelevantEvidence, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RetrievalSignals { + pub result_count: usize, + pub top_score: f32, + pub evidence_coverage: f32, + pub graph_pending: bool, + pub warming: bool, + pub rerank_fallback: bool, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RetrievalDecision { + pub action: DecisionAction, + /// A bounded local routing score. It is not an externally calibrated + /// probability and is labelled as such by `confidence_kind`. + pub confidence: f32, + pub uncertainty: f32, + pub confidence_kind: &'static str, + pub signals: RetrievalSignals, + pub reasons: Vec, +} + +/// Build evidence records by comparing the indexed chunk body with the current +/// source lines. Reads are cached per `(file, line range)` so a result set with +/// adjacent snippets does not repeatedly hit the filesystem. +pub fn collect_evidence(results: &[CodeResult], chunks: &[MergeChunk]) -> Vec { + let mut current_cache: HashMap<(String, u32, u32), Option> = HashMap::new(); + let captured_at = Utc::now().to_rfc3339(); + + results + .iter() + .map(|result| { + let source_key = (result.file.clone(), result.line_start, result.line_end); + let current = current_cache + .entry(source_key) + .or_insert_with(|| read_numbered_source(result)) + .clone(); + let indexed_chunk = chunks.iter().find(|chunk| { + chunk.file == result.file + && chunk.line_start <= result.line_start + && chunk.line_end >= result.line_end + }); + let indexed = indexed_chunk.map(|chunk| normalize_indexed(&chunk.content)); + + let (status, current_sha256, reason) = match current.as_deref() { + None if !Path::new(&result.file).exists() => ( + EvidenceStatus::Missing, + None, + "source file is no longer present on disk".to_owned(), + ), + None => ( + EvidenceStatus::Unavailable, + None, + "source lines could not be read from disk".to_owned(), + ), + Some(current) => { + let current_hash = sha256(current); + let partial = indexed_chunk.is_some_and(|chunk| { + chunk.line_start != result.line_start || chunk.line_end != result.line_end + }); + if partial { + ( + EvidenceStatus::Partial, + Some(current_hash), + "the result contains a bounded slice of a larger indexed chunk" + .to_owned(), + ) + } else if indexed.as_deref() == Some(current) { + ( + EvidenceStatus::Unchanged, + Some(current_hash), + "source content matches the indexed snapshot".to_owned(), + ) + } else { + ( + EvidenceStatus::Changed, + Some(current_hash), + "source content differs from the indexed snapshot".to_owned(), + ) + } + } + }; + + EvidenceRecord { + file: result.file.clone(), + line_start: result.line_start, + line_end: result.line_end, + status, + indexed_sha256: indexed.as_deref().map(sha256), + current_sha256, + reason, + captured_at: captured_at.clone(), + } + }) + .collect() +} + +/// Emit a typed routing decision from deterministic retrieval signals. +pub fn decide( + results: &[CodeResult], + evidence: &[EvidenceRecord], + warming: bool, + graph_pending: bool, + rerank_fallback: bool, +) -> RetrievalDecision { + let top_score = results + .iter() + .map(|result| result.score) + .fold(0.0_f32, f32::max) + .clamp(0.0, 1.0); + let unchanged = evidence + .iter() + .filter(|record| matches!(record.status, EvidenceStatus::Unchanged)) + .count(); + let coverage = if evidence.is_empty() { + if results.is_empty() { 0.0 } else { 0.5 } + } else { + unchanged as f32 / evidence.len() as f32 + }; + + let (mut action, mut confidence, mut reasons) = if results.is_empty() && warming { + ( + DecisionAction::RetryIndex, + 0.95, + vec!["the vector shard is still warming or publishing".to_owned()], + ) + } else if results.is_empty() { + ( + DecisionAction::NoRelevantEvidence, + 0.65, + vec!["no candidate survived the retrieval and content fences".to_owned()], + ) + } else { + ( + DecisionAction::Answer, + (top_score * 0.65 + coverage * 0.35).clamp(0.0, 1.0), + vec!["retrieval returned content-verified candidates".to_owned()], + ) + }; + + if evidence.iter().any(|record| { + matches!( + record.status, + EvidenceStatus::Changed | EvidenceStatus::Missing | EvidenceStatus::Unavailable + ) + }) { + action = DecisionAction::NeedMoreEvidence; + confidence *= 0.55; + reasons.push("at least one cited source is stale, missing, or unreadable".to_owned()); + } else if evidence + .iter() + .any(|record| matches!(record.status, EvidenceStatus::Partial)) + { + confidence *= 0.85; + reasons.push("some evidence is a bounded slice of a larger chunk".to_owned()); + } + if graph_pending && !results.is_empty() { + action = DecisionAction::NeedMoreEvidence; + confidence *= 0.8; + reasons.push("call-graph resolution is still pending".to_owned()); + } + if rerank_fallback && !results.is_empty() { + confidence *= 0.9; + reasons.push("optional reranking was unavailable; local ranking was used".to_owned()); + } + + confidence = confidence.clamp(0.0, 1.0); + RetrievalDecision { + action, + confidence, + uncertainty: 1.0 - confidence, + confidence_kind: "local_heuristic", + signals: RetrievalSignals { + result_count: results.len(), + top_score, + evidence_coverage: coverage, + graph_pending, + warming, + rerank_fallback, + }, + reasons, + } +} + +fn read_numbered_source(result: &CodeResult) -> Option { + crate::query::engine::read_lines_from_fs(&result.file, result.line_start, result.line_end) + .ok() + .map(|text| normalize_numbered(&text)) +} + +fn normalize_indexed(content: &str) -> String { + content.replace("\r\n", "\n").trim_end().to_owned() +} + +fn normalize_numbered(content: &str) -> String { + content + .lines() + .map(|line| { + line.split_once(": ") + .and_then(|(number, text)| number.parse::().ok().map(|_| text)) + .unwrap_or(line) + }) + .collect::>() + .join("\n") + .trim_end() + .to_owned() +} + +fn sha256(value: &str) -> String { + let digest = Sha256::digest(value.as_bytes()); + format!("{digest:x}") +} + +#[cfg(test)] +mod tests { + use super::*; + + fn result(score: f32) -> CodeResult { + CodeResult { + file: "/tmp/example.rs".to_owned(), + line_start: 1, + line_end: 1, + score, + content: "1: fn main() {}".to_owned(), + symbol: Some("main".to_owned()), + callers: None, + caller_files: None, + caller_names: vec![], + callee_names: vec![], + callees: None, + } + } + + #[test] + fn empty_warming_result_requests_retry() { + let decision = decide(&[], &[], true, false, false); + assert_eq!(decision.action, DecisionAction::RetryIndex); + assert!(decision.confidence > 0.9); + } + + #[test] + fn stale_evidence_requires_more_evidence() { + let evidence = vec![EvidenceRecord { + file: "x.rs".to_owned(), + line_start: 1, + line_end: 2, + status: EvidenceStatus::Changed, + indexed_sha256: Some("a".to_owned()), + current_sha256: Some("b".to_owned()), + reason: "changed".to_owned(), + captured_at: "now".to_owned(), + }]; + let decision = decide(&[result(0.9)], &evidence, false, false, false); + assert_eq!(decision.action, DecisionAction::NeedMoreEvidence); + assert!(decision.uncertainty > 0.4); + } +} diff --git a/src/query/engine.rs b/src/query/engine.rs index 50bf2a2..2f88f13 100644 --- a/src/query/engine.rs +++ b/src/query/engine.rs @@ -13,8 +13,10 @@ use crate::embedding::voyage::VoyageClient; use crate::indexing::IndexEngine; use crate::llm::LlmClient; use crate::path_in_repo; +use crate::query::decision::{self, EvidenceRecord, RetrievalDecision}; use crate::query::find_db_for_file; use crate::query::graph_expand::graph_expand; +use crate::query::hybrid::{self, LexicalCandidate}; use crate::query::merger::{MergeChunk, merge_chunks}; use crate::query::reranker; @@ -82,6 +84,11 @@ impl QueryGraphMode { pub struct QueryResult { pub results: Vec, pub pre_rerank_results: Vec, + /// Source validity snapshots aligned with `results`. + #[serde(default)] + pub evidence: Vec, + /// Typed local routing decision for automation clients. + pub decision: RetrievalDecision, pub timing: QueryTiming, pub rerank: Option, /// True when the target repo's vector shard was not resident after the bounded @@ -234,7 +241,7 @@ pub(crate) async fn run_query_with_filters_and_mode( // Search for 2× top_k so graph expansion has candidates to work with. let current_identity = EmbeddingIdentity::from_client(voyage_client); let crate::indexing::VectorSearchOutcome { - results: raw_results, + results: mut raw_results, mut warming, } = index_engine .vector_search( @@ -245,12 +252,49 @@ pub(crate) async fn run_query_with_filters_and_mode( ¤t_identity, ) .await; + // Clone DB handles (Arc-wrapped, cheap) and release the RwLock immediately. + // This prevents holding the lock across lexical lookup and graph awaits. + let db_map: HashMap> = { + let guard = repo_dbs.read().await; + guard.clone() + }; // read lock dropped HERE — before any async DB queries + + // ── Step 2.5: Hybrid lexical retrieval ─────────────────────────────── + // Search the durable chunk store as a bounded second signal. This recovers + // exact names, paths, and error strings that a semantic-only top-k can miss. + let lexical_candidates = lexical_candidates(&db_map, repo_filter, embed_query, top_k * 2).await; let search_ms = search_start.elapsed().as_millis() as u64; - if raw_results.is_empty() { + // Apply repo filter to vector candidates, then fuse semantic and lexical + // scores. Lexical-only hits still enter the normal graph/merge/rerank path. + let mut filtered: Vec<_> = if let Some(repo) = repo_filter { + raw_results + .drain(..) + .filter(|r| path_in_repo(&r.chunk_id.file, repo)) + .collect() + } else { + raw_results + }; + fuse_retrieval_scores(&mut filtered, lexical_candidates); + filtered.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + filtered.truncate(top_k * 2); + + if filtered.is_empty() { return Ok(QueryResult { results: vec![], pre_rerank_results: vec![], + evidence: vec![], + decision: decision::decide( + &[], + &[], + warming, + matches!(graph_mode, QueryGraphMode::VectorOnly), + false, + ), timing: QueryTiming { embed_ms, search_ms, @@ -265,24 +309,9 @@ pub(crate) async fn run_query_with_filters_and_mode( }); } - // Apply repo filter. - let filtered: Vec<_> = if let Some(repo) = repo_filter { - raw_results - .into_iter() - .filter(|r| path_in_repo(&r.chunk_id.file, repo)) - .collect() - } else { - raw_results - }; - // ── Step 3: Fetch stored content for base chunks ────────────────────── // Clone DB handles (Arc-wrapped, cheap) and release the RwLock immediately. // This prevents holding the lock across graph expansion await points. - let db_map: HashMap> = { - let guard = repo_dbs.read().await; - guard.clone() - }; // read lock dropped HERE — before any async DB queries - let fenced = hydrate_candidates(&db_map, &filtered).await; if fenced.dropped > 0 { // A vector candidate without a durable chunk row is an index publication @@ -295,6 +324,14 @@ pub(crate) async fn run_query_with_filters_and_mode( return Ok(QueryResult { results: vec![], pre_rerank_results: vec![], + evidence: vec![], + decision: decision::decide( + &[], + &[], + true, + matches!(graph_mode, QueryGraphMode::VectorOnly), + false, + ), timing: QueryTiming { embed_ms, search_ms, @@ -541,10 +578,20 @@ pub(crate) async fn run_query_with_filters_and_mode( fallback_used: rerank_output.fallback_used, skip_reason: rerank_output.skip_reason, }; + let evidence = decision::collect_evidence(&results, res_chunks); + let retrieval_decision = decision::decide( + &results, + &evidence, + warming, + matches!(graph_mode, QueryGraphMode::VectorOnly), + rerank_info.fallback_used, + ); Ok(QueryResult { results, pre_rerank_results, + evidence, + decision: retrieval_decision, timing: QueryTiming { embed_ms, search_ms, @@ -1000,6 +1047,90 @@ pub fn format_callee_tag(stats: &CallerCalleeStats) -> String { } } +/// Search the repository's durable chunks for exact text matches. A lexical +/// lookup is best-effort: semantic retrieval and graph expansion remain valid +/// when a database is warming or an older index cannot answer the text query. +async fn lexical_candidates( + db_map: &HashMap>, + repo_filter: Option<&str>, + query: &str, + limit: usize, +) -> Vec { + let normalized_filter = repo_filter.map(crate::store::normalize_repo_path); + let targets: Vec<&Surreal> = db_map + .iter() + .filter(|(repo, _)| { + normalized_filter + .as_deref() + .is_none_or(|filter| repo.as_str() == filter) + }) + .map(|(_, db)| db) + .collect(); + + let mut candidates = Vec::new(); + for db in targets { + match hybrid::search_db(db, query, limit).await { + Ok(mut found) => candidates.append(&mut found), + Err(error) => { + warn!(error = %error, "hybrid lexical retrieval degraded to semantic results"); + } + } + } + candidates.sort_by(|a, b| { + b.lexical_score + .partial_cmp(&a.lexical_score) + .unwrap_or(std::cmp::Ordering::Equal) + }); + candidates.truncate(limit); + candidates +} + +/// Fuse vector and lexical signals before the existing graph and rerank stages. +/// The weights intentionally favor semantic similarity while giving exact text +/// matches enough weight to recover names and paths missed by embeddings. +fn fuse_retrieval_scores( + semantic: &mut Vec, + lexical: Vec, +) { + let mut lexical_by_key: HashMap<(String, u32, u32), f32> = lexical + .into_iter() + .map(|candidate| { + ( + ( + candidate.result.chunk_id.file.clone(), + candidate.result.chunk_id.line_start, + candidate.result.chunk_id.line_end, + ), + candidate.lexical_score, + ) + }) + .collect(); + + for result in semantic.iter_mut() { + let key = ( + result.chunk_id.file.clone(), + result.chunk_id.line_start, + result.chunk_id.line_end, + ); + if let Some(lexical_score) = lexical_by_key.remove(&key) { + result.score = (result.score * 0.65 + lexical_score * 0.35).clamp(0.0, 1.0); + } + } + + for ((file, line_start, line_end), lexical_score) in lexical_by_key { + semantic.push(crate::vector::SearchResult { + chunk_id: crate::vector::ChunkId { + file, + line_start, + line_end, + }, + // A lexical-only hit is deliberately strong enough to surface, but + // remains below a perfect semantic+lexical match. + score: (lexical_score * 0.85).clamp(0.0, 1.0), + }); + } +} + /// Fetch stored chunk content, symbol short name, full FQN, and symbol kind /// from whichever DB contains the file. #[allow(clippy::result_large_err)] diff --git a/src/query/hybrid.rs b/src/query/hybrid.rs new file mode 100644 index 0000000..f649bfb --- /dev/null +++ b/src/query/hybrid.rs @@ -0,0 +1,234 @@ +//! Bounded lexical retrieval used alongside semantic vector search. +//! +//! The vector index is still the primary semantic signal. This module adds a +//! small, deterministic text signal so exact symbol names, paths, and error +//! strings remain discoverable when an embedding is a poor match. Results are +//! capped before they enter the normal graph/merge/rerank pipeline. + +use anyhow::{Context, Result}; +use serde::Deserialize; +use surrealdb::Surreal; +use surrealdb::engine::local::Db; + +use crate::vector::{ChunkId, SearchResult}; + +const MIN_TERM_CHARS: usize = 2; + +#[derive(Debug, Clone)] +pub struct LexicalCandidate { + pub result: SearchResult, + pub lexical_score: f32, +} + +#[derive(Debug, Deserialize)] +struct LexicalRow { + file: String, + line_start: i64, + line_end: i64, + content: String, + symbol_ref: Option, +} + +/// Split a query into stable, case-insensitive terms suitable for bounded +/// substring lookup. Duplicate terms are removed while preserving order. +pub fn lexical_terms(query: &str) -> Vec { + let mut terms = Vec::new(); + for term in query + .split(|c: char| !c.is_alphanumeric() && c != '_') + .map(str::trim) + .filter(|term| term.chars().count() >= MIN_TERM_CHARS) + .map(str::to_lowercase) + { + if !terms.iter().any(|existing| existing == &term) { + terms.push(term); + } + } + terms +} + +/// Score one row using transparent, bounded signals. The value is always in +/// `[0, 1]`; it is a ranking feature, not a calibrated probability. +pub fn lexical_score( + query: &str, + terms: &[String], + content: &str, + file: &str, + symbol: Option<&str>, +) -> f32 { + if terms.is_empty() { + return 0.0; + } + + let content_lower = content.to_lowercase(); + let file_lower = file.to_lowercase(); + let query_lower = query.trim().to_lowercase(); + let matched = terms + .iter() + .filter(|term| content_lower.contains(term.as_str())) + .count(); + let term_coverage = matched as f32 / terms.len() as f32; + let phrase_match = + if query_lower.len() >= MIN_TERM_CHARS && content_lower.contains(&query_lower) { + 1.0 + } else { + 0.0 + }; + let path_match = if terms.iter().any(|term| file_lower.contains(term)) { + 1.0 + } else { + 0.0 + }; + let symbol_match = symbol + .map(|name| { + let name = name.to_lowercase(); + if terms + .iter() + .any(|term| name == *term || name.contains(term)) + { + 1.0 + } else { + 0.0 + } + }) + .unwrap_or(0.0); + + (term_coverage * 0.55 + phrase_match * 0.2 + symbol_match * 0.2 + path_match * 0.05) + .clamp(0.0, 1.0) +} + +/// Find text matches in one repository database. The query uses only bound +/// values; the dynamic part consists solely of generated parameter names. +pub async fn search_db( + db: &Surreal, + query: &str, + limit: usize, +) -> Result> { + let terms = lexical_terms(query); + if terms.is_empty() || limit == 0 { + return Ok(Vec::new()); + } + + let predicates = terms + .iter() + .enumerate() + .map(|(index, _)| { + format!( + "(string::lowercase(content) CONTAINS $term{index} OR \ + string::lowercase(file) CONTAINS $term{index})" + ) + }) + .collect::>(); + let statement = format!( + "SELECT file, line_start, line_end, content, symbol_ref FROM chunk WHERE {} LIMIT $limit", + predicates.join(" OR ") + ); + + let mut query_builder = db.query(statement).bind(("limit", limit.min(500) as i64)); + for (index, term) in terms.iter().enumerate() { + query_builder = query_builder.bind((format!("term{index}"), term.clone())); + } + let rows: Vec = query_builder + .await + .context("hybrid lexical chunk search")? + .take(0) + .context("decode hybrid lexical chunk search")?; + + let mut candidates = rows + .into_iter() + .map(|row| { + let symbol = row + .symbol_ref + .as_deref() + .and_then(|value| value.strip_prefix("symbol:⟨")) + .and_then(|value| value.strip_suffix('⟩')) + .and_then(|value| value.rsplit("::").next()); + let score = lexical_score(query, &terms, &row.content, &row.file, symbol); + LexicalCandidate { + result: SearchResult { + chunk_id: ChunkId { + file: row.file, + line_start: row.line_start.max(1) as u32, + line_end: row.line_end.max(row.line_start).max(1) as u32, + }, + score, + }, + lexical_score: score, + } + }) + .collect::>(); + candidates.sort_by(|a, b| { + b.lexical_score + .partial_cmp(&a.lexical_score) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.result.chunk_id.file.cmp(&b.result.chunk_id.file)) + .then_with(|| { + a.result + .chunk_id + .line_start + .cmp(&b.result.chunk_id.line_start) + }) + }); + candidates.truncate(limit.min(500)); + Ok(candidates) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn terms_are_stable_and_deduplicated() { + assert_eq!( + lexical_terms("Auth auth: UserService.create()"), + vec!["auth", "userservice", "create"] + ); + } + + #[test] + fn exact_phrase_and_symbol_match_rank_above_path_only() { + let terms = lexical_terms("UserService create"); + let exact = lexical_score( + "UserService create", + &terms, + "fn UserService create() {}", + "/repo/src/service.rs", + Some("create"), + ); + let path_only = lexical_score( + "UserService create", + &terms, + "fn unrelated() {}", + "/repo/UserService/create.rs", + None, + ); + assert!(exact > path_only); + assert!((0.0..=1.0).contains(&exact)); + } + + #[tokio::test] + async fn search_db_returns_bounded_text_matches() { + let dir = tempfile::TempDir::new().expect("tempdir"); + let db = crate::store::open_db(dir.path(), "/repo", 0) + .await + .expect("open db"); + db.query( + "CREATE chunk SET file = $file, line_start = 1, line_end = 2, \ + content = $content, symbol_ref = $symbol", + ) + .bind(("file", "/repo/src/auth.rs".to_owned())) + .bind(("content", "fn authenticate_user() {}".to_owned())) + .bind(( + "symbol", + "symbol:⟨/repo/src/auth.rs::authenticate_user⟩".to_owned(), + )) + .await + .expect("insert chunk"); + + let matches = search_db(&db, "authenticate_user", 5) + .await + .expect("lexical search"); + assert_eq!(matches.len(), 1); + assert_eq!(matches[0].result.chunk_id.file, "/repo/src/auth.rs"); + assert!(matches[0].lexical_score > 0.5); + } +} diff --git a/src/query/mod.rs b/src/query/mod.rs index dc26f74..c66efe8 100644 --- a/src/query/mod.rs +++ b/src/query/mod.rs @@ -1,7 +1,9 @@ pub mod content_fence; +pub mod decision; pub mod engine; pub mod filters; pub mod graph_expand; +pub mod hybrid; pub mod merger; pub mod reranker; diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs index 09cd322..67533f6 100644 --- a/src/runtime/mod.rs +++ b/src/runtime/mod.rs @@ -1,2 +1,3 @@ pub mod router; +pub mod stdio; pub mod worker; diff --git a/src/runtime/stdio.rs b/src/runtime/stdio.rs new file mode 100644 index 0000000..fc564c7 --- /dev/null +++ b/src/runtime/stdio.rs @@ -0,0 +1,51 @@ +//! Stdio MCP entry point for Claude Code, Codex, Cursor, and similar clients. + +use context_engine_rs::engine_boot::{BootOptions, boot_engine_with_home}; +use context_engine_rs::mcp::McpHandler; +use rmcp::ServiceExt; + +use crate::Cli; + +pub async fn run(cli: &Cli) { + let booted = match boot_engine_with_home( + BootOptions { + data_dir: cli.data_dir.clone(), + embeddings_dir: cli.embeddings_dir.clone(), + no_watchers: false, + only_repo: None, + }, + cli.home_dir.clone(), + ) + .await + { + Ok(booted) => booted, + Err(error) => { + eprintln!("error: {error:#}"); + std::process::exit(2); + } + }; + + let settings = booted.settings.read().await.clone(); + let handler = McpHandler::new( + booted.home_dir, + booted.data_dir, + booted.index_engine, + booted.repo_dbs, + booted.settings, + &settings.enabled_mcp_tools, + ); + + let transport = rmcp::transport::stdio(); + match handler.serve(transport).await { + Ok(service) => { + if let Err(error) = service.waiting().await { + eprintln!("error: MCP stdio service stopped: {error}"); + std::process::exit(1); + } + } + Err(error) => { + eprintln!("error: could not start MCP stdio service: {error}"); + std::process::exit(1); + } + } +} diff --git a/src/typesafe.rs b/src/typesafe.rs new file mode 100644 index 0000000..96e372d --- /dev/null +++ b/src/typesafe.rs @@ -0,0 +1,233 @@ +//! Minimal TypeSafe System One HTTP client. +//! +//! The client is intentionally optional. When `TYPESAFE_API_KEY` is absent, +//! callers keep the local retrieval decision path. When configured, the API is +//! used for narrow typed judgments over an already-bounded state snapshot. + +use std::time::Duration; + +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +pub const DEFAULT_ENDPOINT: &str = "https://api.typesafe.ai/v1/systemone"; +pub const DEFAULT_MODEL: &str = "jev-latest"; +const DEFAULT_TIMEOUT_MS: u64 = 5_000; +const MAX_TIMEOUT_MS: u64 = 30_000; + +#[derive(Debug, Clone)] +pub struct TypeSafeClient { + http: reqwest::Client, + endpoint: String, + api_key: String, + model: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TypeSafeUsage { + #[serde(default)] + pub input_tokens: Option, + #[serde(default)] + pub output_tokens: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TypeSafeResponse { + pub model: String, + #[serde(default)] + pub answers: Map, + #[serde(default)] + pub usage: Option, +} + +impl TypeSafeClient { + /// Build a client from the official environment variable convention. + /// Returns `Ok(None)` when no key is configured, without touching output. + pub fn from_env() -> Result> { + let Some(api_key) = std::env::var("TYPESAFE_API_KEY") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + return Ok(None); + }; + + let endpoint = std::env::var("TYPESAFE_API_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| DEFAULT_ENDPOINT.to_owned()); + let model = std::env::var("TYPESAFE_MODEL") + .ok() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| DEFAULT_MODEL.to_owned()); + let timeout_ms = std::env::var("TYPESAFE_TIMEOUT_MS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_TIMEOUT_MS) + .clamp(100, MAX_TIMEOUT_MS); + let http = reqwest::Client::builder() + .timeout(Duration::from_millis(timeout_ms)) + .build() + .context("build TypeSafe HTTP client")?; + + Ok(Some(Self { + http, + endpoint, + api_key, + model, + })) + } + + #[cfg(test)] + pub fn with_endpoint_for_test(endpoint: String, api_key: &str) -> Self { + Self { + http: reqwest::Client::new(), + endpoint, + api_key: api_key.to_owned(), + model: DEFAULT_MODEL.to_owned(), + } + } + + pub async fn evaluate(&self, state: Value, questions: Value) -> Result { + let body = serde_json::json!({ + "state": state, + "model": self.model, + "questions": questions, + }); + let response = self + .http + .post(&self.endpoint) + .bearer_auth(&self.api_key) + .json(&body) + .send() + .await + .context("TypeSafe API request")?; + let status = response.status(); + if !status.is_success() { + // Do not include the response body: provider errors can echo user + // state, and credentials must never become part of an error string. + bail!("TypeSafe API returned HTTP {status}"); + } + response + .json::() + .await + .context("decode TypeSafe API response") + } +} + +/// Questions used by `ask-context`. Each is a small decision that callers can +/// gate independently instead of asking an external model to write prose. +pub fn ask_questions() -> Value { + serde_json::json!({ + "answerable": { + "type": "noul", + "instructions": "Do the retrieved code snippets provide enough direct evidence to answer the user's query?", + "criteria": { + "true": "The evidence directly supports a reliable answer.", + "false": "The evidence is missing, stale, or too indirect." + } + }, + "evidence_quality": { + "type": "score", + "instructions": "How strong and current is the retrieved code evidence for the user's query?", + "criteria": [ + "Insufficient or stale evidence", + "Partial evidence that needs a follow-up", + "Strong, direct, current evidence" + ] + }, + "next_action": { + "type": "choice", + "instructions": "Choose the next software action from the retrieved evidence.", + "criteria": { + "answer": "Return an answer grounded in the current evidence.", + "need_more_evidence": "Search, inspect, or fetch more evidence before answering.", + "retry_index": "Wait for indexing or warming and retry.", + "no_relevant_evidence": "The indexed workspace does not contain relevant evidence." + } + } + }) +} + +/// Keep the state sent to the external decision service bounded and free of +/// full source dumps. The caller supplies an already-limited JSON snapshot. +pub fn bounded_state(state: Value, max_chars: usize) -> Value { + let Ok(serialized) = serde_json::to_string(&state) else { + return serde_json::json!({ "state": "unserializable" }); + }; + if serialized.chars().count() <= max_chars { + return state; + } + serde_json::json!({ + "truncated_state": serialized.chars().take(max_chars).collect::(), + "truncated": true + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::{Router, extract::Json, http::HeaderMap, routing::post}; + use serde_json::{Value, json}; + use tokio::net::TcpListener; + + #[test] + fn questions_use_all_system_one_primitives() { + let questions = ask_questions(); + assert_eq!(questions["answerable"]["type"], "noul"); + assert_eq!(questions["evidence_quality"]["type"], "score"); + assert_eq!(questions["next_action"]["type"], "choice"); + } + + #[test] + fn bounded_state_keeps_small_values_and_marks_large_values() { + let small = json!({"query":"auth"}); + assert_eq!(bounded_state(small.clone(), 100), small); + let large = bounded_state(json!({"text":"x".repeat(100)}), 20); + assert_eq!(large["truncated"], true); + } + + #[tokio::test] + async fn client_sends_bearer_request_and_decodes_typed_response() { + let app = Router::new().route( + "/systemone", + post(|headers: HeaderMap, Json(body): Json| async move { + assert_eq!( + headers + .get("authorization") + .and_then(|value| value.to_str().ok()), + Some("Bearer test-key") + ); + assert_eq!(body["model"], DEFAULT_MODEL); + assert!(body["questions"]["answerable"].is_object()); + ( + [("content-type", "application/json")], + json!({ + "model": "jev-1.13.0", + "answers": { + "answerable": {"type":"noul", "noul":0.9} + }, + "usage": {"input_tokens": 10, "output_tokens": 4} + }) + .to_string(), + ) + }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let address = listener.local_addr().expect("address"); + tokio::spawn(async move { + axum::serve(listener, app).await.expect("serve"); + }); + + let client = TypeSafeClient::with_endpoint_for_test( + format!("http://{address}/systemone"), + "test-key", + ); + let response = client + .evaluate(json!({"query":"auth"}), ask_questions()) + .await + .expect("response"); + assert_eq!(response.model, "jev-1.13.0"); + assert_eq!(response.answers["answerable"]["noul"], 0.9); + assert_eq!(response.usage.expect("usage").input_tokens, Some(10)); + } +}