Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d8d77824f4 | ||
|
|
4a04faf926 | ||
|
|
0782159ecd | ||
|
|
0d12ca9854 | ||
|
|
2639093fe4 | ||
|
|
16bd7df1b4 | ||
|
|
d6604ba0d9 | ||
|
|
07cdeef2ea | ||
|
|
f6a2b05045 | ||
|
|
84ab33a2d5 | ||
|
|
5d0e4e3456 | ||
|
|
c6643ca134 | ||
|
|
ba0da84d59 | ||
|
|
5014d4a183 | ||
|
|
6eb396720e | ||
|
|
afab4d0518 | ||
|
|
fa63901e51 | ||
|
|
0281a7ea32 | ||
|
|
9c4fc4dd39 | ||
|
|
0376604dea | ||
|
|
0d256f7dcc | ||
|
|
40c70b99f4 | ||
|
|
514de0eef1 | ||
|
|
0674961c05 | ||
|
|
6c7541756b | ||
|
|
3c90dbadef | ||
|
|
4d722c9765 | ||
|
|
1b77426a3f | ||
|
|
5a51860222 | ||
|
|
da01153eb5 | ||
|
|
65778084ba | ||
|
|
7c117aa9bd | ||
|
|
bf22c918d8 | ||
|
|
d357b2d90f | ||
|
|
5cc13d3e6e | ||
|
|
adda067f02 | ||
|
|
0a370e759d | ||
|
|
f16021a9be | ||
|
|
0381bb2fff | ||
|
|
8c22b6b96c | ||
|
|
61dd175815 | ||
|
|
4a95f5fb8b | ||
|
|
667cab4cad | ||
|
|
722a5dfa7b | ||
|
|
b1a64fbbf8 | ||
|
|
e536ca2c78 | ||
|
|
39c87d9d1d | ||
|
|
c2ce57efff | ||
|
|
c18e5b0816 | ||
|
|
36f8315ad9 | ||
|
|
74e8a55c1b | ||
|
|
46216f022f | ||
|
|
611fdb0bb4 | ||
|
|
036fd348a6 | ||
|
|
10a69d680d | ||
|
|
d263814335 | ||
|
|
7f91b926f6 | ||
|
|
fd5f59a6aa | ||
|
|
c0694b65e3 | ||
|
|
36823f42b7 | ||
|
|
9d0cd25681 | ||
|
|
51bd6e6c2b | ||
|
|
1494397960 | ||
|
|
96ec25e20f | ||
|
|
68d82637d5 | ||
|
|
dde05102cd | ||
|
|
929b2f7a0b | ||
|
|
e98d0b7dda | ||
|
|
afb98496ce | ||
|
|
d8dc4e7b82 | ||
|
|
9a5193ba4e | ||
|
|
1a8b419363 | ||
|
|
18c25f1dc3 | ||
|
|
b817af219d | ||
|
|
5cd49ca8fd | ||
|
|
97c1098146 | ||
|
|
46bcb75b34 | ||
|
|
739ef3812e | ||
|
|
ed48852636 | ||
|
|
7dba6e18bb | ||
|
|
d49fae6e86 | ||
|
|
0b378e7346 | ||
|
|
4b32118dbd | ||
|
|
4a72d62a88 | ||
|
|
9e817713f6 | ||
|
|
ab9d779ecb | ||
|
|
d086412f91 | ||
|
|
fbd4a9efed | ||
|
|
2f95dd1c2f | ||
|
|
e6ab0117f6 | ||
|
|
858c3302b1 | ||
|
|
e12f66bba2 | ||
|
|
c373574548 | ||
|
|
d5415baa6b | ||
|
|
9388b86f06 | ||
|
|
e4b288f096 | ||
|
|
e22d3d6061 | ||
|
|
bf2358fe67 | ||
|
|
b35ce59693 | ||
|
|
62703a9ed5 | ||
|
|
da4c31562f | ||
|
|
0c820804dc | ||
|
|
de60af332f | ||
|
|
2c60c3d37c | ||
|
|
bbe8da63ad | ||
|
|
803e67d864 | ||
|
|
8366d92d95 | ||
|
|
c9dc8c3aa1 | ||
|
|
a4200cb555 | ||
|
|
7d70fd9978 | ||
|
|
a8a0c999e5 | ||
|
|
d6e868c946 | ||
|
|
24037d0174 |
@@ -1 +0,0 @@
|
||||
{"messages":[],"version":1}
|
||||
@@ -1 +0,0 @@
|
||||
{"messages":[{"blocks":[{"text":"Say hello in one sentence","type":"text"}],"role":"user"},{"blocks":[{"text":"Hello! I'm Claude, an AI assistant ready to help you with software engineering tasks, code analysis, debugging, or any other programming challenges you might have.","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":11,"output_tokens":32}}],"version":1}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"permissions": {
|
||||
"defaultMode": "dontAsk"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
ANTHROPIC_API_KEY="9494feba6f7c45f48c3dfc35a85ffd89.2WUCscxcSp92ETNg"
|
||||
ANTHROPIC_BASE_URL="https://open.bigmodel.cn/api/anthropic"
|
||||
CLAW_MODEL="glm-5"
|
||||
@@ -1 +0,0 @@
|
||||
github: instructkr
|
||||
@@ -0,0 +1,36 @@
|
||||
name: CI
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
|
||||
jobs:
|
||||
rust:
|
||||
name: ${{ matrix.os }}
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os:
|
||||
- ubuntu-latest
|
||||
- macos-latest
|
||||
|
||||
steps:
|
||||
- name: Check out repository
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Install Rust toolchain
|
||||
uses: dtolnay/rust-toolchain@stable
|
||||
|
||||
- name: Run cargo check
|
||||
run: cargo check --workspace
|
||||
|
||||
- name: Run cargo test
|
||||
run: cargo test --workspace
|
||||
|
||||
- name: Run release build
|
||||
run: cargo build --release
|
||||
@@ -1,4 +1,6 @@
|
||||
__pycache__/
|
||||
archive/
|
||||
target/
|
||||
.omx/
|
||||
.clawd-agents/
|
||||
# Claw Code local artifacts
|
||||
.claw/
|
||||
.claude/
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
# CLAUDE.md
|
||||
|
||||
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
|
||||
|
||||
## Project Overview
|
||||
|
||||
Claw Code is a local coding-agent CLI tool written in safe Rust. It is a clean-room implementation inspired by Claude Code, providing an interactive REPL, one-shot prompts, workspace-aware tools, local agent workflows, and plugin support. The project name and references throughout use "claw" / "Claw Code".
|
||||
|
||||
## Build & Run Commands
|
||||
|
||||
```bash
|
||||
# Build release binary (produces target/release/claw)
|
||||
cargo build --release -p claw-cli
|
||||
|
||||
# Run from source (interactive REPL)
|
||||
cargo run --bin claw --
|
||||
|
||||
# Run one-shot prompt
|
||||
cargo run --bin claw -- prompt "summarize this workspace"
|
||||
|
||||
# Install locally
|
||||
cargo install --path crates/claw-cli --locked
|
||||
|
||||
# Run the HTTP server binary
|
||||
cargo run --bin claw-server
|
||||
```
|
||||
|
||||
## Verification Commands
|
||||
|
||||
```bash
|
||||
# Format check
|
||||
cargo fmt
|
||||
|
||||
# Lint (workspace-level clippy with deny warnings)
|
||||
cargo clippy --workspace --all-targets -- -D warnings
|
||||
|
||||
# Run all tests
|
||||
cargo test --workspace
|
||||
|
||||
# Run tests for a specific crate
|
||||
cargo test -p api
|
||||
cargo test -p runtime
|
||||
|
||||
# Run a single test by name
|
||||
cargo test -p <crate> -- <test_name>
|
||||
|
||||
# Integration tests in crates/api/tests/ use mock TCP servers (no network needed)
|
||||
# One smoke test is #[ignore] — run with: cargo test -p api -- --ignored
|
||||
```
|
||||
|
||||
## Workspace Architecture
|
||||
|
||||
Cargo workspace with `resolver = "2"`. All crates live under `crates/`.
|
||||
|
||||
### Crate Dependency Graph
|
||||
|
||||
```
|
||||
claw-cli ──→ api, runtime, tools, commands, plugins, compat-harness
|
||||
server ──→ api, runtime, tools, plugins, commands
|
||||
tools ──→ api, runtime, plugins
|
||||
commands ──→ runtime, plugins
|
||||
api ──→ runtime
|
||||
runtime ──→ lsp, plugins
|
||||
plugins ──→ (standalone, serde only)
|
||||
lsp ──→ (standalone)
|
||||
compat-harness ──→ commands, tools, runtime
|
||||
```
|
||||
|
||||
### Core Crates
|
||||
|
||||
- **`claw-cli`** — User-facing binary (`claw`). REPL loop with markdown rendering (pulldown-cmark + syntect), argument parsing, OAuth flow. Entry point: `crates/claw-cli/src/main.rs`.
|
||||
|
||||
- **`runtime`** — Session management, conversation runtime, permissions, system prompt construction, context compaction, MCP stdio management, and hook execution. Key types: `Session` (versioned message history), `ConversationRuntime<C, T>` (generic over `ApiClient` + `ToolExecutor` traits), `PermissionMode`, `McpServerManager`, `HookRunner`.
|
||||
|
||||
- **`api`** — HTTP client for LLM providers with SSE streaming. `ClawApiClient` (Anthropic-compatible), `OpenAiCompatClient`, and `Provider` trait. `ProviderKind` enum distinguishes ClawApi, Xai, OpenAi. Request/response types: `MessageRequest`, `StreamEvent`, `ToolDefinition`.
|
||||
|
||||
- **`tools`** — Built-in tool definitions and dispatch. `GlobalToolRegistry` is a lazy-static singleton. Tools: Read, Write, Edit, Glob, Grep, Bash, LSP, Task*, Cron*, Worktree*. Each tool has a `ToolSpec` with JSON schema.
|
||||
|
||||
- **`commands`** — Slash command registry and handlers (`/help`, `/config`, `/compact`, `/resume`, `/plugins`, `/agents`, `/doctor`, etc.). `SlashCommandSpec` defines each command's name, aliases, description, and category.
|
||||
|
||||
- **`plugins`** — Plugin discovery and lifecycle. `PluginManager` loads builtin, bundled, and external (from `~/.claw/plugins/`) plugins. Plugins can provide additional tools via `PluginTool`.
|
||||
|
||||
- **`server`** — Axum-based HTTP server (`claw-server`). REST endpoints for session CRUD + SSE event streaming. `AppState` holds shared session store.
|
||||
|
||||
- **`lsp`** — Language Server Protocol types and process management for code intelligence features.
|
||||
|
||||
- **`compat-harness`** — Extracts command/tool/bootstrap-plan manifests from upstream TypeScript source files (for compatibility tracking). Uses `CLAUDE_CODE_UPSTREAM` env var to locate the upstream repo.
|
||||
|
||||
## Key Architectural Patterns
|
||||
|
||||
- **Trait-based abstraction**: `ApiClient`, `ToolExecutor`, `Provider` traits enable swappable implementations. `ConversationRuntime` is generic over client and executor.
|
||||
- **Static registries**: `GlobalToolRegistry` and slash command specs use lazy-static initialization with compile-time definitions.
|
||||
- **SSE streaming**: API responses stream through `MessageStream` (async iterator) to the terminal renderer or server SSE endpoints.
|
||||
- **Permission model**: `PermissionMode` enum — ReadOnly, WorkspaceWrite, DangerFullAccess. Configurable via `.claw.json` (`permissions.defaultMode`).
|
||||
- **Hook system**: Pre/post tool execution hooks via `HookRunner` in the runtime crate.
|
||||
|
||||
## Configuration & Environment
|
||||
|
||||
- `.claw.json` — Project-level config (permissions, etc.)
|
||||
- `.claw/` — Project directory for hooks, plugins, local settings
|
||||
- `~/.claw/` — User-level config directory
|
||||
- `.env` — API keys (gitignored): `ANTHROPIC_API_KEY`, `ANTHROPIC_BASE_URL`, `XAI_API_KEY`, `XAI_BASE_URL`
|
||||
- `CLAW.md` — Workspace instructions loaded into the system prompt (analogous to CLAUDE.md but for claw itself)
|
||||
|
||||
## Lint Rules
|
||||
|
||||
- `unsafe_code` is **forbidden** at workspace level
|
||||
- Clippy `all` + `pedantic` lints are warnings; some pedantic lints are allowed (`module_name_repetitions`, `missing_panics_doc`, `missing_errors_doc`)
|
||||
- CI runs: `cargo check --workspace`, `cargo test --workspace`, `cargo build --release` on Ubuntu and macOS
|
||||
|
||||
## Language
|
||||
|
||||
Code comments, commit messages, and documentation are primarily in Chinese (中文). UI strings and exported symbol names are in English.
|
||||
@@ -0,0 +1,15 @@
|
||||
# CLAW.md
|
||||
|
||||
This file provides guidance to Claw Code (clawcode.dev) when working with code in this repository.
|
||||
|
||||
## Detected stack
|
||||
- Languages: Rust.
|
||||
- Frameworks: none detected from the supported starter markers.
|
||||
|
||||
## Verification
|
||||
- Run Rust verification from the repo root: `cargo fmt`, `cargo clippy --workspace --all-targets -- -D warnings`, `cargo test --workspace`
|
||||
|
||||
## Working agreement
|
||||
- Prefer small, reviewable changes and keep generated bootstrap files aligned with actual repo workflows.
|
||||
- Keep shared defaults in `.claw.json`; reserve `.claw/settings.local.json` for machine-local overrides.
|
||||
- Do not overwrite existing `CLAW.md` content automatically; update it intentionally when repo workflows change.
|
||||
@@ -0,0 +1,43 @@
|
||||
# 贡献指南
|
||||
|
||||
感谢你为 Claw Code 做出贡献。
|
||||
|
||||
## 开发设置
|
||||
|
||||
- 安装稳定的 Rust 工具链。
|
||||
- 在此 Rust 工作区的仓库根目录下进行开发。如果你从父仓库根目录开始,请先执行 `cd rust/`。
|
||||
|
||||
## 构建
|
||||
|
||||
```bash
|
||||
cargo build
|
||||
cargo build --release
|
||||
```
|
||||
|
||||
## 测试与验证
|
||||
|
||||
在开启 Pull Request 之前,请运行完整的 Rust 验证集:
|
||||
|
||||
```bash
|
||||
cargo fmt --all --check
|
||||
cargo clippy --workspace --all-targets -- -D warnings
|
||||
cargo check --workspace
|
||||
cargo test --workspace
|
||||
```
|
||||
|
||||
如果你更改了行为,请在同一个 Pull Request 中添加或更新相关的测试。
|
||||
|
||||
## 代码风格
|
||||
|
||||
- 遵循所修改 crate 中的现有模式,而不是引入新的风格。
|
||||
- 使用 `rustfmt` 格式化代码。
|
||||
- 确保你修改的工作区目标的 `clippy` 检查通过。
|
||||
- 优先采用针对性的 diff,而不是顺便进行的重构。
|
||||
|
||||
## Pull Request
|
||||
|
||||
- 从 `main` 分支拉取新分支。
|
||||
- 确保每个 Pull Request 的范围仅限于一个明确的更改。
|
||||
- 说明更改动机、实现摘要以及你运行的验证。
|
||||
- 在请求审查之前,确保本地检查已通过。
|
||||
- 如果审查反馈导致行为更改,请重新运行相关的验证命令。
|
||||
@@ -17,6 +17,15 @@ dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "android_system_properties"
|
||||
version = "0.1.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311"
|
||||
dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "api"
|
||||
version = "0.1.0"
|
||||
@@ -25,15 +34,96 @@ dependencies = [
|
||||
"runtime",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"telemetry",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476"
|
||||
dependencies = [
|
||||
"async-stream-impl",
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-stream-impl"
|
||||
version = "0.3.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "atomic-waker"
|
||||
version = "1.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
|
||||
|
||||
[[package]]
|
||||
name = "autocfg"
|
||||
version = "1.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
|
||||
|
||||
[[package]]
|
||||
name = "axum"
|
||||
version = "0.8.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8"
|
||||
dependencies = [
|
||||
"axum-core",
|
||||
"bytes",
|
||||
"form_urlencoded",
|
||||
"futures-util",
|
||||
"http",
|
||||
"http-body",
|
||||
"http-body-util",
|
||||
"hyper",
|
||||
"hyper-util",
|
||||
"itoa",
|
||||
"matchit",
|
||||
"memchr",
|
||||
"mime",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"serde_core",
|
||||
"serde_json",
|
||||
"serde_path_to_error",
|
||||
"serde_urlencoded",
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tower",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "axum-core"
|
||||
version = "0.5.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"http",
|
||||
"http-body",
|
||||
"http-body-util",
|
||||
"mime",
|
||||
"pin-project-lite",
|
||||
"sync_wrapper",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "base64"
|
||||
version = "0.22.1"
|
||||
@@ -49,6 +139,12 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "1.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
|
||||
|
||||
[[package]]
|
||||
name = "bitflags"
|
||||
version = "2.11.0"
|
||||
@@ -98,11 +194,54 @@ version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
|
||||
|
||||
[[package]]
|
||||
name = "chrono"
|
||||
version = "0.4.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0"
|
||||
dependencies = [
|
||||
"iana-time-zone",
|
||||
"js-sys",
|
||||
"num-traits",
|
||||
"wasm-bindgen",
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "claw-cli"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"api",
|
||||
"commands",
|
||||
"compat-harness",
|
||||
"crossterm",
|
||||
"dotenvy",
|
||||
"plugins",
|
||||
"pulldown-cmark",
|
||||
"runtime",
|
||||
"rustyline",
|
||||
"serde_json",
|
||||
"syntect",
|
||||
"tokio",
|
||||
"tools",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "clipboard-win"
|
||||
version = "5.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "bde03770d3df201d4fb868f2c9c59e66a3e4e2bd06692a0fe701e7103c7e84d4"
|
||||
dependencies = [
|
||||
"error-code",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "commands"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"plugins",
|
||||
"runtime",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -114,6 +253,12 @@ dependencies = [
|
||||
"tools",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation-sys"
|
||||
version = "0.8.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
|
||||
|
||||
[[package]]
|
||||
name = "cpufeatures"
|
||||
version = "0.2.17"
|
||||
@@ -138,11 +283,11 @@ version = "0.28.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "829d955a0bb380ef178a640b91779e3987da38c9aea133b20614cfed8cdea9c6"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"crossterm_winapi",
|
||||
"mio",
|
||||
"parking_lot",
|
||||
"rustix",
|
||||
"rustix 0.38.44",
|
||||
"signal-hook",
|
||||
"signal-hook-mio",
|
||||
"winapi",
|
||||
@@ -197,6 +342,30 @@ dependencies = [
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "dotenvy"
|
||||
version = "0.15.7"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b"
|
||||
|
||||
[[package]]
|
||||
name = "either"
|
||||
version = "1.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
|
||||
|
||||
[[package]]
|
||||
name = "endian-type"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d"
|
||||
|
||||
[[package]]
|
||||
name = "env_home"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe"
|
||||
|
||||
[[package]]
|
||||
name = "equivalent"
|
||||
version = "1.0.2"
|
||||
@@ -213,6 +382,23 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "error-code"
|
||||
version = "3.3.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59"
|
||||
|
||||
[[package]]
|
||||
name = "fd-lock"
|
||||
version = "4.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ce92ff622d6dadf7349484f42c93271a0d49b7cc4d466a936405bacbe10aa78"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"rustix 1.1.4",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "find-msvc-tools"
|
||||
version = "0.1.9"
|
||||
@@ -229,6 +415,15 @@ dependencies = [
|
||||
"miniz_oxide",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fluent-uri"
|
||||
version = "0.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "17c704e9dbe1ddd863da1e6ff3567795087b1eb201ce80d8fa81162e1516500d"
|
||||
dependencies = [
|
||||
"bitflags 1.3.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "fnv"
|
||||
version = "1.0.7"
|
||||
@@ -266,6 +461,17 @@ version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
|
||||
|
||||
[[package]]
|
||||
name = "futures-macro"
|
||||
version = "0.3.32"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-sink"
|
||||
version = "0.3.32"
|
||||
@@ -286,6 +492,7 @@ checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-io",
|
||||
"futures-macro",
|
||||
"futures-sink",
|
||||
"futures-task",
|
||||
"memchr",
|
||||
@@ -351,6 +558,15 @@ version = "0.16.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100"
|
||||
|
||||
[[package]]
|
||||
name = "home"
|
||||
version = "0.5.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d"
|
||||
dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "http"
|
||||
version = "1.4.0"
|
||||
@@ -390,6 +606,12 @@ version = "1.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87"
|
||||
|
||||
[[package]]
|
||||
name = "httpdate"
|
||||
version = "1.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
|
||||
|
||||
[[package]]
|
||||
name = "hyper"
|
||||
version = "1.9.0"
|
||||
@@ -403,6 +625,7 @@ dependencies = [
|
||||
"http",
|
||||
"http-body",
|
||||
"httparse",
|
||||
"httpdate",
|
||||
"itoa",
|
||||
"pin-project-lite",
|
||||
"smallvec",
|
||||
@@ -450,6 +673,30 @@ dependencies = [
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iana-time-zone"
|
||||
version = "0.1.65"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470"
|
||||
dependencies = [
|
||||
"android_system_properties",
|
||||
"core-foundation-sys",
|
||||
"iana-time-zone-haiku",
|
||||
"js-sys",
|
||||
"log",
|
||||
"wasm-bindgen",
|
||||
"windows-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "iana-time-zone-haiku"
|
||||
version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f"
|
||||
dependencies = [
|
||||
"cc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "icu_collections"
|
||||
version = "2.1.1"
|
||||
@@ -614,6 +861,12 @@ version = "0.4.15"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab"
|
||||
|
||||
[[package]]
|
||||
name = "linux-raw-sys"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
|
||||
|
||||
[[package]]
|
||||
name = "litemap"
|
||||
version = "0.8.1"
|
||||
@@ -641,12 +894,48 @@ version = "0.1.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
|
||||
|
||||
[[package]]
|
||||
name = "lsp"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"lsp-types",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lsp-types"
|
||||
version = "0.97.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "53353550a17c04ac46c585feb189c2db82154fc84b79c7a66c96c2c644f66071"
|
||||
dependencies = [
|
||||
"bitflags 1.3.2",
|
||||
"fluent-uri",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_repr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "matchit"
|
||||
version = "0.8.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3"
|
||||
|
||||
[[package]]
|
||||
name = "memchr"
|
||||
version = "2.8.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
|
||||
|
||||
[[package]]
|
||||
name = "mime"
|
||||
version = "0.3.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a"
|
||||
|
||||
[[package]]
|
||||
name = "miniz_oxide"
|
||||
version = "0.8.9"
|
||||
@@ -669,12 +958,51 @@ dependencies = [
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "mock-anthropic-service"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"api",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nibble_vec"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43"
|
||||
dependencies = [
|
||||
"smallvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "nix"
|
||||
version = "0.29.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
"cfg-if",
|
||||
"cfg_aliases",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-conv"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c6673768db2d862beb9b39a78fdcb1a69439615d5794a1be50caa9bc92c81967"
|
||||
|
||||
[[package]]
|
||||
name = "num-traits"
|
||||
version = "0.2.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
|
||||
dependencies = [
|
||||
"autocfg",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "once_cell"
|
||||
version = "1.21.4"
|
||||
@@ -687,7 +1015,7 @@ version = "6.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "336b9c63443aceef14bea841b899035ae3abe89b7c486aaf4c5bd8aafedac3f0"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"libc",
|
||||
"once_cell",
|
||||
"onig_sys",
|
||||
@@ -757,6 +1085,14 @@ dependencies = [
|
||||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "plugins"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "potential_utf"
|
||||
version = "0.1.4"
|
||||
@@ -796,7 +1132,7 @@ version = "0.13.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7c3a14896dfa883796f1cb410461aef38810ea05f2b2c33c5aded3649095fdad"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"getopts",
|
||||
"memchr",
|
||||
"pulldown-cmark-escape",
|
||||
@@ -888,6 +1224,16 @@ version = "5.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
|
||||
|
||||
[[package]]
|
||||
name = "radix_trie"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd"
|
||||
dependencies = [
|
||||
"endian-type",
|
||||
"nibble_vec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rand"
|
||||
version = "0.9.2"
|
||||
@@ -923,7 +1269,7 @@ version = "0.5.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -985,12 +1331,14 @@ dependencies = [
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-rustls",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams",
|
||||
"web-sys",
|
||||
"webpki-roots",
|
||||
]
|
||||
@@ -1014,12 +1362,15 @@ name = "runtime"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"glob",
|
||||
"plugins",
|
||||
"regex",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"telemetry",
|
||||
"tokio",
|
||||
"walkdir",
|
||||
"which",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1034,11 +1385,24 @@ version = "0.38.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys",
|
||||
"windows-sys 0.52.0",
|
||||
"linux-raw-sys 0.4.15",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustix"
|
||||
version = "1.1.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
"errno",
|
||||
"libc",
|
||||
"linux-raw-sys 0.12.1",
|
||||
"windows-sys 0.61.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1090,17 +1454,40 @@ dependencies = [
|
||||
"commands",
|
||||
"compat-harness",
|
||||
"crossterm",
|
||||
"mock-anthropic-service",
|
||||
"plugins",
|
||||
"pulldown-cmark",
|
||||
"reqwest",
|
||||
"runtime",
|
||||
"rustyline",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"syntect",
|
||||
"tokio",
|
||||
"tools",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustyline"
|
||||
version = "15.0.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2ee1e066dc922e513bda599c6ccb5f3bb2b0ea5870a579448f2622993f0a9a2f"
|
||||
dependencies = [
|
||||
"bitflags 2.11.0",
|
||||
"cfg-if",
|
||||
"clipboard-win",
|
||||
"fd-lock",
|
||||
"home",
|
||||
"libc",
|
||||
"log",
|
||||
"memchr",
|
||||
"nix",
|
||||
"radix_trie",
|
||||
"unicode-segmentation",
|
||||
"unicode-width",
|
||||
"utf8parse",
|
||||
"windows-sys 0.59.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ryu"
|
||||
version = "1.0.23"
|
||||
@@ -1165,6 +1552,28 @@ dependencies = [
|
||||
"zmij",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_path_to_error"
|
||||
version = "0.1.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"serde",
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_repr"
|
||||
version = "0.1.20"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "175ee3e80ae9982737ca543e96133087cbd9a485eecc3bc4de9c1a37b47ea59c"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_urlencoded"
|
||||
version = "0.7.1"
|
||||
@@ -1177,6 +1586,27 @@ dependencies = [
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "server"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"api",
|
||||
"async-stream",
|
||||
"axum",
|
||||
"chrono",
|
||||
"commands",
|
||||
"dotenvy",
|
||||
"plugins",
|
||||
"reqwest",
|
||||
"runtime",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
"tools",
|
||||
"tower",
|
||||
"tower-http",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.10.9"
|
||||
@@ -1317,6 +1747,14 @@ dependencies = [
|
||||
"yaml-rust",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "telemetry"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "2.0.18"
|
||||
@@ -1430,11 +1868,27 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-util"
|
||||
version = "0.7.18"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"futures-sink",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tools"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"api",
|
||||
"commands",
|
||||
"flate2",
|
||||
"plugins",
|
||||
"reqwest",
|
||||
"runtime",
|
||||
"serde",
|
||||
@@ -1455,6 +1909,7 @@ dependencies = [
|
||||
"tokio",
|
||||
"tower-layer",
|
||||
"tower-service",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -1463,7 +1918,7 @@ version = "0.6.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8"
|
||||
dependencies = [
|
||||
"bitflags",
|
||||
"bitflags 2.11.0",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"http",
|
||||
@@ -1493,6 +1948,7 @@ version = "0.1.44"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100"
|
||||
dependencies = [
|
||||
"log",
|
||||
"pin-project-lite",
|
||||
"tracing-core",
|
||||
]
|
||||
@@ -1530,6 +1986,12 @@ version = "1.0.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-segmentation"
|
||||
version = "1.13.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-width"
|
||||
version = "0.2.2"
|
||||
@@ -1560,6 +2022,12 @@ version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be"
|
||||
|
||||
[[package]]
|
||||
name = "utf8parse"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
|
||||
|
||||
[[package]]
|
||||
name = "version_check"
|
||||
version = "0.9.5"
|
||||
@@ -1655,6 +2123,19 @@ dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-streams"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "web-sys"
|
||||
version = "0.3.93"
|
||||
@@ -1684,6 +2165,18 @@ dependencies = [
|
||||
"rustls-pki-types",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "which"
|
||||
version = "7.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "24d643ce3fd3e5b54854602a080f34fb10ab75e0b813ee32d00ca2b44fa74762"
|
||||
dependencies = [
|
||||
"either",
|
||||
"env_home",
|
||||
"rustix 1.1.4",
|
||||
"winsafe",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winapi"
|
||||
version = "0.3.9"
|
||||
@@ -1715,12 +2208,65 @@ version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
|
||||
|
||||
[[package]]
|
||||
name = "windows-core"
|
||||
version = "0.62.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb"
|
||||
dependencies = [
|
||||
"windows-implement",
|
||||
"windows-interface",
|
||||
"windows-link",
|
||||
"windows-result",
|
||||
"windows-strings",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-implement"
|
||||
version = "0.60.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-interface"
|
||||
version = "0.59.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-link"
|
||||
version = "0.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
|
||||
|
||||
[[package]]
|
||||
name = "windows-result"
|
||||
version = "0.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-strings"
|
||||
version = "0.5.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091"
|
||||
dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.52.0"
|
||||
@@ -1730,6 +2276,15 @@ dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.59.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
|
||||
dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.60.2"
|
||||
@@ -1877,6 +2432,12 @@ version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650"
|
||||
|
||||
[[package]]
|
||||
name = "winsafe"
|
||||
version = "0.0.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904"
|
||||
|
||||
[[package]]
|
||||
name = "wit-bindgen"
|
||||
version = "0.51.0"
|
||||
@@ -0,0 +1,55 @@
|
||||
[workspace]
|
||||
members = ["crates/*"]
|
||||
resolver = "2"
|
||||
|
||||
[workspace.package]
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
publish = false
|
||||
|
||||
[workspace.dependencies]
|
||||
lsp-types = "0.97"
|
||||
serde_json = "1"
|
||||
|
||||
[workspace.lints.rust]
|
||||
unsafe_code = "forbid"
|
||||
|
||||
[workspace.lints.clippy]
|
||||
all = { level = "warn", priority = -1 }
|
||||
pedantic = { level = "warn", priority = -1 }
|
||||
module_name_repetitions = "allow"
|
||||
missing_panics_doc = "allow"
|
||||
missing_errors_doc = "allow"
|
||||
uninlined_format_args = "allow"
|
||||
map_unwrap_or = "allow"
|
||||
doc_markdown = "allow"
|
||||
redundant_pattern_matching = "allow"
|
||||
items_after_statements = "allow"
|
||||
must_use_candidate = "allow"
|
||||
io_other_error = "allow"
|
||||
implicit_clone = "allow"
|
||||
redundant_closure_for_method_calls = "allow"
|
||||
unnecessary_map_or = "allow"
|
||||
manual_let_else = "allow"
|
||||
format_push_string = "allow"
|
||||
match_same_arms = "allow"
|
||||
similar_names = "allow"
|
||||
needless_continue = "allow"
|
||||
assigning_clones = "allow"
|
||||
cast_possible_truncation = "allow"
|
||||
cast_possible_wrap = "allow"
|
||||
cast_sign_loss = "allow"
|
||||
cmp_owned = "allow"
|
||||
collapsible_if = "allow"
|
||||
too_many_lines = "allow"
|
||||
wildcard_in_or_patterns = "allow"
|
||||
explicit_counter_loop = "allow"
|
||||
manual_repeat_n = "allow"
|
||||
manual_str_repeat = "allow"
|
||||
needless_borrow = "allow"
|
||||
needless_pass_by_value = "allow"
|
||||
single_match_else = "allow"
|
||||
too_many_arguments = "allow"
|
||||
unnested_or_patterns = "allow"
|
||||
unused_self = "allow"
|
||||
@@ -1,191 +1,121 @@
|
||||
# Rewriting Project Claw Code
|
||||
# Claw Code
|
||||
|
||||
<p align="center">
|
||||
<strong>⭐ The fastest repo in history to surpass 50K stars, reaching the milestone in just 2 hours after publication ⭐</strong>
|
||||
</p>
|
||||
Claw Code 是一个使用安全 Rust 实现的本地编程代理(coding-agent)命令行工具。它的设计灵感来自 **Claude Code**,并作为一个**净室实现(clean-room implementation)**开发:旨在提供强大的本地代理体验,但它**不是** Claude Code 的直接移植或复制。
|
||||
|
||||
<p align="center">
|
||||
<a href="https://star-history.com/#instructkr/claw-code&Date">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" srcset="https://api.star-history.com/svg?repos=instructkr/claw-code&type=Date&theme=dark" />
|
||||
<source media="(prefers-color-scheme: light)" srcset="https://api.star-history.com/svg?repos=instructkr/claw-code&type=Date" />
|
||||
<img alt="Star History Chart" src="https://api.star-history.com/svg?repos=instructkr/claw-code&type=Date" width="600" />
|
||||
</picture>
|
||||
</a>
|
||||
</p>
|
||||
Rust 工作区是当前主要的产品界面。`claw` 二进制文件在单个工作区内提供交互式会话、单次提示、工作区感知工具、本地代理工作流以及支持插件的操作。
|
||||
|
||||
<p align="center">
|
||||
<img src="assets/clawd-hero.jpeg" alt="Claw" width="300" />
|
||||
</p>
|
||||
## 当前状态
|
||||
|
||||
<p align="center">
|
||||
<strong>Better Harness Tools, not merely storing the archive of leaked Claude Code</strong>
|
||||
</p>
|
||||
- **版本:** `0.1.0`
|
||||
- **发布阶段:** 初始公开发布,源码编译分发
|
||||
- **主要实现:** 本仓库中的 Rust 工作区
|
||||
- **平台焦点:** macOS 和 Linux 开发工作站
|
||||
|
||||
<p align="center">
|
||||
<a href="https://github.com/sponsors/instructkr"><img src="https://img.shields.io/badge/Sponsor-%E2%9D%A4-pink?logo=github&style=for-the-badge" alt="Sponsor on GitHub" /></a>
|
||||
</p>
|
||||
## 安装、构建与运行
|
||||
|
||||
> [!IMPORTANT]
|
||||
> **Rust port is now in progress** on the [`dev/rust`](https://github.com/instructkr/claw-code/tree/dev/rust) branch and is expected to be merged into main today. The Rust implementation aims to deliver a faster, memory-safe harness runtime. Stay tuned — this will be the definitive version of the project.
|
||||
### 准备工作
|
||||
|
||||
> If you find this work useful, consider [sponsoring @instructkr on GitHub](https://github.com/sponsors/instructkr) to support continued open-source harness engineering research.
|
||||
- Rust 稳定版工具链
|
||||
- Cargo
|
||||
- 你想使用的模型的提供商凭据
|
||||
|
||||
---
|
||||
你可以通过环境变量或在项目根目录创建 **`.env`** 文件来配置 API 密钥:
|
||||
|
||||
## Backstory
|
||||
|
||||
At 4 AM on March 31, 2026, I woke up to my phone blowing up with notifications. The Claude Code source had been exposed, and the entire dev community was in a frenzy. My girlfriend in Korea was genuinely worried I might face legal action from Anthropic just for having the code on my machine — so I did what any engineer would do under pressure: I sat down, ported the core features to Python from scratch, and pushed it before the sun came up.
|
||||
|
||||
The whole thing was orchestrated end-to-end using [oh-my-codex (OmX)](https://github.com/Yeachan-Heo/oh-my-codex) by [@bellman_ych](https://x.com/bellman_ych) — a workflow layer built on top of OpenAI's Codex ([@OpenAIDevs](https://x.com/OpenAIDevs)). I used `$team` mode for parallel code review and `$ralph` mode for persistent execution loops with architect-level verification. The entire porting session — from reading the original harness structure to producing a working Python tree with tests — was driven through OmX orchestration.
|
||||
|
||||
The result is a clean-room Python rewrite that captures the architectural patterns of Claude Code's agent harness without copying any proprietary source. I'm now actively collaborating with [@bellman_ych](https://x.com/bellman_ych) — the creator of OmX himself — to push this further. The basic Python foundation is already in place and functional, but we're just getting started. **Stay tuned — a much more capable version is on the way.**
|
||||
|
||||
https://github.com/instructkr/claw-code
|
||||
|
||||

|
||||
|
||||
## The Creators Featured in Wall Street Journal For Avid Claude Code Fans
|
||||
|
||||
I've been deeply interested in **harness engineering** — studying how agent systems wire tools, orchestrate tasks, and manage runtime context. This isn't a sudden thing. The Wall Street Journal featured my work earlier this month, documenting how I've been one of the most active power users exploring these systems:
|
||||
|
||||
> AI startup worker Sigrid Jin, who attended the Seoul dinner, single-handedly used 25 billion of Claude Code tokens last year. At the time, usage limits were looser, allowing early enthusiasts to reach tens of billions of tokens at a very low cost.
|
||||
>
|
||||
> Despite his countless hours with Claude Code, Jin isn't faithful to any one AI lab. The tools available have different strengths and weaknesses, he said. Codex is better at reasoning, while Claude Code generates cleaner, more shareable code.
|
||||
>
|
||||
> Jin flew to San Francisco in February for Claude Code's first birthday party, where attendees waited in line to compare notes with Cherny. The crowd included a practicing cardiologist from Belgium who had built an app to help patients navigate care, and a California lawyer who made a tool for automating building permit approvals using Claude Code.
|
||||
>
|
||||
> "It was basically like a sharing party," Jin said. "There were lawyers, there were doctors, there were dentists. They did not have software engineering backgrounds."
|
||||
>
|
||||
> — *The Wall Street Journal*, March 21, 2026, [*"The Trillion Dollar Race to Automate Our Entire Lives"*](https://lnkd.in/gs9td3qd)
|
||||
|
||||

|
||||
|
||||
---
|
||||
|
||||
## Porting Status
|
||||
|
||||
The main source tree is now Python-first.
|
||||
|
||||
- `src/` contains the active Python porting workspace
|
||||
- `tests/` verifies the current Python workspace
|
||||
- the exposed snapshot is no longer part of the tracked repository state
|
||||
|
||||
The current Python workspace is not yet a complete one-to-one replacement for the original system, but the primary implementation surface is now Python.
|
||||
|
||||
## Why this rewrite exists
|
||||
|
||||
I originally studied the exposed codebase to understand its harness, tool wiring, and agent workflow. After spending more time with the legal and ethical questions—and after reading the essay linked below—I did not want the exposed snapshot itself to remain the main tracked source tree.
|
||||
|
||||
This repository now focuses on Python porting work instead.
|
||||
|
||||
## Repository Layout
|
||||
|
||||
```text
|
||||
.
|
||||
├── src/ # Python porting workspace
|
||||
│ ├── __init__.py
|
||||
│ ├── commands.py
|
||||
│ ├── main.py
|
||||
│ ├── models.py
|
||||
│ ├── port_manifest.py
|
||||
│ ├── query_engine.py
|
||||
│ ├── task.py
|
||||
│ └── tools.py
|
||||
├── tests/ # Python verification
|
||||
├── assets/omx/ # OmX workflow screenshots
|
||||
├── 2026-03-09-is-legal-the-same-as-legitimate-ai-reimplementation-and-the-erosion-of-copyleft.md
|
||||
└── README.md
|
||||
**配置 Claude (推荐):**
|
||||
```bash
|
||||
ANTHROPIC_API_KEY="..."
|
||||
# 使用兼容的端点时可选
|
||||
export ANTHROPIC_BASE_URL="https://api.anthropic.com"
|
||||
```
|
||||
|
||||
## Python Workspace Overview
|
||||
|
||||
The new Python `src/` tree currently provides:
|
||||
|
||||
- **`port_manifest.py`** — summarizes the current Python workspace structure
|
||||
- **`models.py`** — dataclasses for subsystems, modules, and backlog state
|
||||
- **`commands.py`** — Python-side command port metadata
|
||||
- **`tools.py`** — Python-side tool port metadata
|
||||
- **`query_engine.py`** — renders a Python porting summary from the active workspace
|
||||
- **`main.py`** — a CLI entrypoint for manifest and summary output
|
||||
|
||||
## Quickstart
|
||||
|
||||
Render the Python porting summary:
|
||||
Grok 模型:
|
||||
|
||||
```bash
|
||||
python3 -m src.main summary
|
||||
export XAI_API_KEY="..."
|
||||
# 使用兼容的端点时可选
|
||||
export XAI_BASE_URL="https://api.x.ai"
|
||||
```
|
||||
|
||||
Print the current Python workspace manifest:
|
||||
也可以使用 OAuth 登录:
|
||||
|
||||
```bash
|
||||
python3 -m src.main manifest
|
||||
cargo run --bin claw -- login
|
||||
```
|
||||
|
||||
List the current Python modules:
|
||||
### 本地安装
|
||||
|
||||
```bash
|
||||
python3 -m src.main subsystems --limit 16
|
||||
cargo install --path crates/claw-cli --locked
|
||||
```
|
||||
|
||||
Run verification:
|
||||
### 从源码构建
|
||||
|
||||
```bash
|
||||
python3 -m unittest discover -s tests -v
|
||||
cargo build --release -p claw-cli
|
||||
```
|
||||
|
||||
Run the parity audit against the local ignored archive (when present):
|
||||
### 运行
|
||||
|
||||
在工作区内运行:
|
||||
|
||||
```bash
|
||||
python3 -m src.main parity-audit
|
||||
cargo run --bin claw -- --help
|
||||
cargo run --bin claw --
|
||||
cargo run --bin claw -- prompt "总结此工作区"
|
||||
cargo run --bin claw -- --model sonnet "审查最新更改"
|
||||
```
|
||||
|
||||
Inspect mirrored command/tool inventories:
|
||||
运行发布版本:
|
||||
|
||||
```bash
|
||||
python3 -m src.main commands --limit 10
|
||||
python3 -m src.main tools --limit 10
|
||||
./target/release/claw
|
||||
./target/release/claw prompt "解释 crates/runtime"
|
||||
```
|
||||
|
||||
## Current Parity Checkpoint
|
||||
## 支持的功能
|
||||
|
||||
The port now mirrors the archived root-entry file surface, top-level subsystem names, and command/tool inventories much more closely than before. However, it is **not yet** a full runtime-equivalent replacement for the original TypeScript system; the Python tree still contains fewer executable runtime slices than the archived source.
|
||||
- 交互式 REPL 和单次提示执行
|
||||
- 已保存会话的检查和恢复流程
|
||||
- 内置工作区工具:shell、文件读/写/编辑、搜索、网页获取/搜索、待办事项和笔记本更新
|
||||
- 斜杠命令:状态、压缩、配置检查、差异(diff)、导出、会话管理和版本报告
|
||||
- 本地代理和技能发现:通过 `claw agents` 和 `claw skills`
|
||||
- 通过命令行和斜杠命令界面发现并管理插件
|
||||
- OAuth 登录/注销,以及从命令行选择模型/提供商
|
||||
- 工作区感知的指令/配置加载(`CLAW.md`、配置文件、权限、插件设置)
|
||||
|
||||
## 当前限制
|
||||
|
||||
## Built with `oh-my-codex`
|
||||
- 目前公开发布**仅限源码构建**;此工作区尚未设置 crates.io 发布
|
||||
- GitHub CI 验证 `cargo check`、`cargo test` 和发布构建,但尚未提供自动化的发布打包
|
||||
- 当前 CI 目标为 Ubuntu 和 macOS;Windows 的发布就绪性仍待建立
|
||||
- 一些实时提供商集成覆盖是可选的,因为它们需要外部凭据 and 网络访问
|
||||
- 命令界面可能会在 `0.x` 系列期间继续演进
|
||||
|
||||
The restructuring and documentation work on this repository was AI-assisted and orchestrated with Yeachan Heo's [oh-my-codex (OmX)](https://github.com/Yeachan-Heo/oh-my-codex), layered on top of Codex.
|
||||
## 实现现状
|
||||
|
||||
- **`$team` mode:** used for coordinated parallel review and architectural feedback
|
||||
- **`$ralph` mode:** used for persistent execution, verification, and completion discipline
|
||||
- **Codex-driven workflow:** used to turn the main `src/` tree into a Python-first porting workspace
|
||||
Rust 工作区是当前的产品实现。目前包含以下 crate:
|
||||
|
||||
### OmX workflow screenshots
|
||||
- `claw-cli` — 面向用户的二进制文件
|
||||
- `api` — 提供商客户端和流式处理
|
||||
- `runtime` — 会话、配置、权限、提示词和运行时循环
|
||||
- `tools` — 内置工具实现
|
||||
- `commands` — 斜杠命令注册和处理程序
|
||||
- `plugins` — 插件发现、注册和生命周期支持
|
||||
- `lsp` — 语言服务器协议支持类型和进程助手
|
||||
- `server` 和 `compat-harness` — 支持服务和兼容性工具
|
||||
|
||||

|
||||
## 路线图
|
||||
|
||||
*Ralph/team orchestration view while the README and essay context were being reviewed in terminal panes.*
|
||||
- 发布打包好的构件,用于公共安装
|
||||
- 添加可重复的发布工作流和长期维护的变更日志(changelog)规范
|
||||
- 将平台验证扩展到当前 CI 矩阵之外
|
||||
- 添加更多以任务为中心的示例和操作员文档
|
||||
- 继续加强 Rust 实现的功能覆盖并磨炼用户体验(UX)
|
||||
|
||||

|
||||
## 发行版本说明
|
||||
|
||||
*Split-pane review and verification flow during the final README wording pass.*
|
||||
- 0.1.0 发行说明草案:[`docs/releases/0.1.0.md`](docs/releases/0.1.0.md)
|
||||
|
||||
## Community
|
||||
## 许可
|
||||
|
||||
<p align="center">
|
||||
<a href="https://instruct.kr/"><img src="assets/instructkr.png" alt="instructkr" width="400" /></a>
|
||||
</p>
|
||||
|
||||
Join the [**instructkr Discord**](https://instruct.kr/) — the best Korean language model community. Come chat about LLMs, harness engineering, agent workflows, and everything in between.
|
||||
|
||||
[](https://instruct.kr/)
|
||||
|
||||
## Star History
|
||||
|
||||
See the chart at the top of this README.
|
||||
|
||||
## Ownership / Affiliation Disclaimer
|
||||
|
||||
- This repository does **not** claim ownership of the original Claude Code source material.
|
||||
- This repository is **not affiliated with, endorsed by, or maintained by Anthropic**.
|
||||
有关许可详情,请参阅仓库根目录。
|
||||
|
||||
|
Before Width: | Height: | Size: 233 KiB |
|
Before Width: | Height: | Size: 4.8 KiB |
|
Before Width: | Height: | Size: 1.3 MiB |
|
Before Width: | Height: | Size: 1.4 MiB |
|
Before Width: | Height: | Size: 312 KiB |
|
Before Width: | Height: | Size: 812 KiB |
|
Before Width: | Height: | Size: 873 KiB |
@@ -9,7 +9,8 @@ publish.workspace = true
|
||||
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
|
||||
runtime = { path = "../runtime" }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
serde_json.workspace = true
|
||||
telemetry = { path = "../telemetry" }
|
||||
tokio = { version = "1", features = ["io-util", "macros", "net", "rt-multi-thread", "time"] }
|
||||
|
||||
[lints]
|
||||
@@ -0,0 +1,53 @@
|
||||
# API 模块 (api)
|
||||
|
||||
本模块提供了与大型语言模型 (LLM) 服务提供商(主要是 Anthropic 的 Claude 和兼容 OpenAI 的服务)进行交互的高层抽象和客户端。
|
||||
|
||||
## 概览
|
||||
|
||||
`api` 模块负责以下职责:
|
||||
- 标准化与不同 AI 提供商的通信。
|
||||
- 通过服务器发送事件 (SSE) 处理流式响应。
|
||||
- 管理身份验证源(API 密钥、OAuth 令牌)。
|
||||
- 提供消息、工具和使用情况跟踪的共享数据结构。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **提供商抽象 (Provider Abstraction)**:支持多种 AI 后端,包括:
|
||||
- `ClawApiClient`: Claude 模型的主要提供商。
|
||||
- `OpenAiCompatClient`: 支持兼容 OpenAI 的 API(如本地模型、专门的提供商)。
|
||||
- **流式支持 (Streaming Support)**:健壮的 SSE 解析实现 (`SseParser`),用于处理实时的内容生成。
|
||||
- **工具集成 (Tool Integration)**:为 `ToolDefinition`、`ToolChoice` 和 `ToolResultContentBlock` 提供强类型定义,支持智能代理 (Agentic) 工作流。
|
||||
- **身份验证管理 (Auth Management)**:用于解析启动身份验证源和管理 OAuth 令牌的实用工具。
|
||||
- **模型智能 (Model Intelligence)**:解析模型别名和计算最大标记 (Token) 限制的元数据及辅助函数。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`client.rs`**: 定义了 `ProviderClient` 特性 (Trait) 和基础客户端逻辑。它使用 `reqwest` 处理 HTTP 请求,并管理消息流的生命周期。
|
||||
- **`types.rs`**: 包含 API 的核心数据模型,如 `InputMessage`、`OutputContentBlock` 以及 `MessageRequest`/`MessageResponse`。
|
||||
- **`sse.rs`**: 实现了一个状态化的 SSE 解析器,能够处理分段的数据块并发出类型化的 `StreamEvent`。
|
||||
- **`providers/`**: 包含针对不同 LLM 端点的特定逻辑,将它们的独特格式映射到本模块使用的共享类型。
|
||||
|
||||
### 数据流
|
||||
|
||||
1. 构建包含模型详情、消息和工具定义的 `MessageRequest`。
|
||||
2. `ApiClient` 将此请求转换为提供商特定的 HTTP 请求。
|
||||
3. 如果启用了流式传输,客户端返回一个 `MessageStream`,该流使用 `SseParser` 来产生 `StreamEvent`。
|
||||
4. 最终响应包含用于跟踪 Token 消耗的 `Usage` 信息。
|
||||
|
||||
## 使用示例
|
||||
|
||||
```rust
|
||||
use api::{ApiClient, MessageRequest, InputMessage};
|
||||
|
||||
// 示例初始化(已简化)
|
||||
let client = ApiClient::new(auth_source);
|
||||
let request = MessageRequest {
|
||||
model: "claude-3-5-sonnet-20241022".to_string(),
|
||||
messages: vec![InputMessage::user("你好,世界!")],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let stream = client.create_message_stream(request).await?;
|
||||
```
|
||||
@@ -0,0 +1,241 @@
|
||||
use crate::error::ApiError;
|
||||
use crate::prompt_cache::{PromptCache, PromptCacheRecord, PromptCacheStats};
|
||||
use crate::providers::anthropic::{self, AnthropicClient, AuthSource};
|
||||
use crate::providers::openai_compat::{self, OpenAiCompatClient, OpenAiCompatConfig};
|
||||
use crate::providers::{self, ProviderKind};
|
||||
use crate::types::{MessageRequest, MessageResponse, StreamEvent};
|
||||
|
||||
#[allow(clippy::large_enum_variant)]
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum ProviderClient {
|
||||
Anthropic(AnthropicClient),
|
||||
Xai(OpenAiCompatClient),
|
||||
OpenAi(OpenAiCompatClient),
|
||||
}
|
||||
|
||||
impl ProviderClient {
|
||||
pub fn from_model(model: &str) -> Result<Self, ApiError> {
|
||||
Self::from_model_with_anthropic_auth(model, None)
|
||||
}
|
||||
|
||||
pub fn from_model_with_anthropic_auth(
|
||||
model: &str,
|
||||
anthropic_auth: Option<AuthSource>,
|
||||
) -> Result<Self, ApiError> {
|
||||
let resolved_model = providers::resolve_model_alias(model);
|
||||
match providers::detect_provider_kind(&resolved_model) {
|
||||
ProviderKind::Anthropic => Ok(Self::Anthropic(match anthropic_auth {
|
||||
Some(auth) => AnthropicClient::from_auth(auth),
|
||||
None => AnthropicClient::from_env()?,
|
||||
})),
|
||||
ProviderKind::Xai => Ok(Self::Xai(OpenAiCompatClient::from_env(
|
||||
OpenAiCompatConfig::xai(),
|
||||
)?)),
|
||||
ProviderKind::OpenAi => {
|
||||
// DashScope models (qwen-*) also return ProviderKind::OpenAi because they
|
||||
// speak the OpenAI wire format, but they need the DashScope config which
|
||||
// reads DASHSCOPE_API_KEY and points at dashscope.aliyuncs.com.
|
||||
let config = match providers::metadata_for_model(&resolved_model) {
|
||||
Some(meta) if meta.auth_env == "DASHSCOPE_API_KEY" => {
|
||||
OpenAiCompatConfig::dashscope()
|
||||
}
|
||||
_ => OpenAiCompatConfig::openai(),
|
||||
};
|
||||
Ok(Self::OpenAi(OpenAiCompatClient::from_env(config)?))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn provider_kind(&self) -> ProviderKind {
|
||||
match self {
|
||||
Self::Anthropic(_) => ProviderKind::Anthropic,
|
||||
Self::Xai(_) => ProviderKind::Xai,
|
||||
Self::OpenAi(_) => ProviderKind::OpenAi,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_prompt_cache(self, prompt_cache: PromptCache) -> Self {
|
||||
match self {
|
||||
Self::Anthropic(client) => Self::Anthropic(client.with_prompt_cache(prompt_cache)),
|
||||
other => other,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn prompt_cache_stats(&self) -> Option<PromptCacheStats> {
|
||||
match self {
|
||||
Self::Anthropic(client) => client.prompt_cache_stats(),
|
||||
Self::Xai(_) | Self::OpenAi(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn take_last_prompt_cache_record(&self) -> Option<PromptCacheRecord> {
|
||||
match self {
|
||||
Self::Anthropic(client) => client.take_last_prompt_cache_record(),
|
||||
Self::Xai(_) | Self::OpenAi(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn send_message(
|
||||
&self,
|
||||
request: &MessageRequest,
|
||||
) -> Result<MessageResponse, ApiError> {
|
||||
match self {
|
||||
Self::Anthropic(client) => client.send_message(request).await,
|
||||
Self::Xai(client) | Self::OpenAi(client) => client.send_message(request).await,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn stream_message(
|
||||
&self,
|
||||
request: &MessageRequest,
|
||||
) -> Result<MessageStream, ApiError> {
|
||||
match self {
|
||||
Self::Anthropic(client) => client
|
||||
.stream_message(request)
|
||||
.await
|
||||
.map(MessageStream::Anthropic),
|
||||
Self::Xai(client) | Self::OpenAi(client) => client
|
||||
.stream_message(request)
|
||||
.await
|
||||
.map(MessageStream::OpenAiCompat),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum MessageStream {
|
||||
Anthropic(anthropic::MessageStream),
|
||||
OpenAiCompat(openai_compat::MessageStream),
|
||||
}
|
||||
|
||||
impl MessageStream {
|
||||
#[must_use]
|
||||
pub fn request_id(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Anthropic(stream) => stream.request_id(),
|
||||
Self::OpenAiCompat(stream) => stream.request_id(),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn next_event(&mut self) -> Result<Option<StreamEvent>, ApiError> {
|
||||
match self {
|
||||
Self::Anthropic(stream) => stream.next_event().await,
|
||||
Self::OpenAiCompat(stream) => stream.next_event().await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub use anthropic::{
|
||||
oauth_token_is_expired, resolve_saved_oauth_token, resolve_startup_auth_source, OAuthTokenSet,
|
||||
};
|
||||
#[must_use]
|
||||
pub fn read_base_url() -> String {
|
||||
anthropic::read_base_url()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn read_xai_base_url() -> String {
|
||||
openai_compat::read_base_url(OpenAiCompatConfig::xai())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
|
||||
use super::ProviderClient;
|
||||
use crate::providers::{detect_provider_kind, resolve_model_alias, ProviderKind};
|
||||
|
||||
/// Serializes every test in this module that mutates process-wide
|
||||
/// environment variables so concurrent test threads cannot observe
|
||||
/// each other's partially-applied state.
|
||||
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
||||
LOCK.get_or_init(|| Mutex::new(()))
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_existing_and_grok_aliases() {
|
||||
assert_eq!(resolve_model_alias("opus"), "claude-opus-4-6");
|
||||
assert_eq!(resolve_model_alias("grok"), "grok-3");
|
||||
assert_eq!(resolve_model_alias("grok-mini"), "grok-3-mini");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_detection_prefers_model_family() {
|
||||
assert_eq!(detect_provider_kind("grok-3"), ProviderKind::Xai);
|
||||
assert_eq!(
|
||||
detect_provider_kind("claude-sonnet-4-6"),
|
||||
ProviderKind::Anthropic
|
||||
);
|
||||
}
|
||||
|
||||
/// Snapshot-restore guard for a single environment variable. Mirrors
|
||||
/// the pattern used in `providers/mod.rs` tests: captures the original
|
||||
/// value on construction, applies the override, and restores on drop so
|
||||
/// tests leave the process env untouched even when they panic.
|
||||
struct EnvVarGuard {
|
||||
key: &'static str,
|
||||
original: Option<std::ffi::OsString>,
|
||||
}
|
||||
|
||||
impl EnvVarGuard {
|
||||
fn set(key: &'static str, value: Option<&str>) -> Self {
|
||||
let original = std::env::var_os(key);
|
||||
match value {
|
||||
Some(value) => std::env::set_var(key, value),
|
||||
None => std::env::remove_var(key),
|
||||
}
|
||||
Self { key, original }
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for EnvVarGuard {
|
||||
fn drop(&mut self) {
|
||||
match self.original.take() {
|
||||
Some(value) => std::env::set_var(self.key, value),
|
||||
None => std::env::remove_var(self.key),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dashscope_model_uses_dashscope_config_not_openai() {
|
||||
// Regression: qwen-plus was being routed to OpenAiCompatConfig::openai()
|
||||
// which reads OPENAI_API_KEY and points at api.openai.com, when it should
|
||||
// use OpenAiCompatConfig::dashscope() which reads DASHSCOPE_API_KEY and
|
||||
// points at dashscope.aliyuncs.com.
|
||||
let _lock = env_lock();
|
||||
let _dashscope = EnvVarGuard::set("DASHSCOPE_API_KEY", Some("test-dashscope-key"));
|
||||
let _openai = EnvVarGuard::set("OPENAI_API_KEY", None);
|
||||
|
||||
let client = ProviderClient::from_model("qwen-plus");
|
||||
|
||||
// Must succeed (not fail with "missing OPENAI_API_KEY")
|
||||
assert!(
|
||||
client.is_ok(),
|
||||
"qwen-plus with DASHSCOPE_API_KEY set should build successfully, got: {:?}",
|
||||
client.err()
|
||||
);
|
||||
|
||||
// Verify it's the OpenAi variant pointed at the DashScope base URL.
|
||||
match client.unwrap() {
|
||||
ProviderClient::OpenAi(openai_client) => {
|
||||
assert!(
|
||||
openai_client.base_url().contains("dashscope.aliyuncs.com"),
|
||||
"qwen-plus should route to DashScope base URL (contains 'dashscope.aliyuncs.com'), got: {}",
|
||||
openai_client.base_url()
|
||||
);
|
||||
}
|
||||
other => panic!(
|
||||
"Expected ProviderClient::OpenAi for qwen-plus, got: {:?}",
|
||||
other
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,572 @@
|
||||
use std::env::VarError;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::time::Duration;
|
||||
|
||||
const GENERIC_FATAL_WRAPPER_MARKERS: &[&str] = &[
|
||||
"something went wrong while processing your request",
|
||||
"please try again, or use /new to start a fresh session",
|
||||
];
|
||||
|
||||
const CONTEXT_WINDOW_ERROR_MARKERS: &[&str] = &[
|
||||
"maximum context length",
|
||||
"context window",
|
||||
"context length",
|
||||
"too many tokens",
|
||||
"prompt is too long",
|
||||
"input is too long",
|
||||
"request is too large",
|
||||
];
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum ApiError {
|
||||
MissingCredentials {
|
||||
provider: &'static str,
|
||||
env_vars: &'static [&'static str],
|
||||
/// Optional, runtime-computed hint appended to the error Display
|
||||
/// output. Populated when the provider resolver can infer what the
|
||||
/// user probably intended (e.g. an OpenAI key is set but Anthropic
|
||||
/// was selected because no Anthropic credentials exist).
|
||||
hint: Option<String>,
|
||||
},
|
||||
ContextWindowExceeded {
|
||||
model: String,
|
||||
estimated_input_tokens: u32,
|
||||
requested_output_tokens: u32,
|
||||
estimated_total_tokens: u32,
|
||||
context_window_tokens: u32,
|
||||
},
|
||||
ExpiredOAuthToken,
|
||||
Auth(String),
|
||||
InvalidApiKeyEnv(VarError),
|
||||
Http(reqwest::Error),
|
||||
Io(std::io::Error),
|
||||
Json {
|
||||
provider: String,
|
||||
model: String,
|
||||
body_snippet: String,
|
||||
source: serde_json::Error,
|
||||
},
|
||||
Api {
|
||||
status: reqwest::StatusCode,
|
||||
error_type: Option<String>,
|
||||
message: Option<String>,
|
||||
request_id: Option<String>,
|
||||
body: String,
|
||||
retryable: bool,
|
||||
},
|
||||
RetriesExhausted {
|
||||
attempts: u32,
|
||||
last_error: Box<ApiError>,
|
||||
},
|
||||
InvalidSseFrame(&'static str),
|
||||
BackoffOverflow {
|
||||
attempt: u32,
|
||||
base_delay: Duration,
|
||||
},
|
||||
}
|
||||
|
||||
impl ApiError {
|
||||
#[must_use]
|
||||
pub const fn missing_credentials(
|
||||
provider: &'static str,
|
||||
env_vars: &'static [&'static str],
|
||||
) -> Self {
|
||||
Self::MissingCredentials {
|
||||
provider,
|
||||
env_vars,
|
||||
hint: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a `MissingCredentials` error carrying an extra, runtime-computed
|
||||
/// hint string that the Display impl appends after the canonical "missing
|
||||
/// <provider> credentials" message. Used by the provider resolver to
|
||||
/// suggest the likely fix when the user has credentials for a different
|
||||
/// provider already in the environment.
|
||||
#[must_use]
|
||||
pub fn missing_credentials_with_hint(
|
||||
provider: &'static str,
|
||||
env_vars: &'static [&'static str],
|
||||
hint: impl Into<String>,
|
||||
) -> Self {
|
||||
Self::MissingCredentials {
|
||||
provider,
|
||||
env_vars,
|
||||
hint: Some(hint.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a `Self::Json` enriched with the provider name, the model that
|
||||
/// was requested, and the first 200 characters of the raw response body so
|
||||
/// that callers can diagnose deserialization failures without re-running
|
||||
/// the request.
|
||||
#[must_use]
|
||||
pub fn json_deserialize(
|
||||
provider: impl Into<String>,
|
||||
model: impl Into<String>,
|
||||
body: &str,
|
||||
source: serde_json::Error,
|
||||
) -> Self {
|
||||
Self::Json {
|
||||
provider: provider.into(),
|
||||
model: model.into(),
|
||||
body_snippet: truncate_body_snippet(body, 200),
|
||||
source,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_retryable(&self) -> bool {
|
||||
match self {
|
||||
Self::Http(error) => error.is_connect() || error.is_timeout() || error.is_request(),
|
||||
Self::Api { retryable, .. } => *retryable,
|
||||
Self::RetriesExhausted { last_error, .. } => last_error.is_retryable(),
|
||||
Self::MissingCredentials { .. }
|
||||
| Self::ContextWindowExceeded { .. }
|
||||
| Self::ExpiredOAuthToken
|
||||
| Self::Auth(_)
|
||||
| Self::InvalidApiKeyEnv(_)
|
||||
| Self::Io(_)
|
||||
| Self::Json { .. }
|
||||
| Self::InvalidSseFrame(_)
|
||||
| Self::BackoffOverflow { .. } => false,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn request_id(&self) -> Option<&str> {
|
||||
match self {
|
||||
Self::Api { request_id, .. } => request_id.as_deref(),
|
||||
Self::RetriesExhausted { last_error, .. } => last_error.request_id(),
|
||||
Self::MissingCredentials { .. }
|
||||
| Self::ContextWindowExceeded { .. }
|
||||
| Self::ExpiredOAuthToken
|
||||
| Self::Auth(_)
|
||||
| Self::InvalidApiKeyEnv(_)
|
||||
| Self::Http(_)
|
||||
| Self::Io(_)
|
||||
| Self::Json { .. }
|
||||
| Self::InvalidSseFrame(_)
|
||||
| Self::BackoffOverflow { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn safe_failure_class(&self) -> &'static str {
|
||||
match self {
|
||||
Self::RetriesExhausted { .. } if self.is_context_window_failure() => "context_window",
|
||||
Self::RetriesExhausted { .. } if self.is_generic_fatal_wrapper() => {
|
||||
"provider_retry_exhausted"
|
||||
}
|
||||
Self::RetriesExhausted { last_error, .. } => last_error.safe_failure_class(),
|
||||
Self::MissingCredentials { .. } | Self::ExpiredOAuthToken | Self::Auth(_) => {
|
||||
"provider_auth"
|
||||
}
|
||||
Self::Api { status, .. } if matches!(status.as_u16(), 401 | 403) => "provider_auth",
|
||||
Self::ContextWindowExceeded { .. } => "context_window",
|
||||
Self::Api { .. } if self.is_context_window_failure() => "context_window",
|
||||
Self::Api { status, .. } if status.as_u16() == 429 => "provider_rate_limit",
|
||||
Self::Api { .. } if self.is_generic_fatal_wrapper() => "provider_internal",
|
||||
Self::Api { .. } => "provider_error",
|
||||
Self::Http(_) | Self::InvalidSseFrame(_) | Self::BackoffOverflow { .. } => {
|
||||
"provider_transport"
|
||||
}
|
||||
Self::InvalidApiKeyEnv(_) | Self::Io(_) | Self::Json { .. } => "runtime_io",
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_generic_fatal_wrapper(&self) -> bool {
|
||||
match self {
|
||||
Self::Api { message, body, .. } => {
|
||||
message
|
||||
.as_deref()
|
||||
.is_some_and(looks_like_generic_fatal_wrapper)
|
||||
|| looks_like_generic_fatal_wrapper(body)
|
||||
}
|
||||
Self::RetriesExhausted { last_error, .. } => last_error.is_generic_fatal_wrapper(),
|
||||
Self::MissingCredentials { .. }
|
||||
| Self::ContextWindowExceeded { .. }
|
||||
| Self::ExpiredOAuthToken
|
||||
| Self::Auth(_)
|
||||
| Self::InvalidApiKeyEnv(_)
|
||||
| Self::Http(_)
|
||||
| Self::Io(_)
|
||||
| Self::Json { .. }
|
||||
| Self::InvalidSseFrame(_)
|
||||
| Self::BackoffOverflow { .. } => false,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_context_window_failure(&self) -> bool {
|
||||
match self {
|
||||
Self::ContextWindowExceeded { .. } => true,
|
||||
Self::Api {
|
||||
status,
|
||||
message,
|
||||
body,
|
||||
..
|
||||
} => {
|
||||
matches!(status.as_u16(), 400 | 413 | 422)
|
||||
&& (message
|
||||
.as_deref()
|
||||
.is_some_and(looks_like_context_window_error)
|
||||
|| looks_like_context_window_error(body))
|
||||
}
|
||||
Self::RetriesExhausted { last_error, .. } => last_error.is_context_window_failure(),
|
||||
Self::MissingCredentials { .. }
|
||||
| Self::ExpiredOAuthToken
|
||||
| Self::Auth(_)
|
||||
| Self::InvalidApiKeyEnv(_)
|
||||
| Self::Http(_)
|
||||
| Self::Io(_)
|
||||
| Self::Json { .. }
|
||||
| Self::InvalidSseFrame(_)
|
||||
| Self::BackoffOverflow { .. } => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for ApiError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::MissingCredentials {
|
||||
provider,
|
||||
env_vars,
|
||||
hint,
|
||||
} => {
|
||||
write!(
|
||||
f,
|
||||
"missing {provider} credentials; export {} before calling the {provider} API",
|
||||
env_vars.join(" or ")
|
||||
)?;
|
||||
if cfg!(target_os = "windows") {
|
||||
if let Some(primary) = env_vars.first() {
|
||||
write!(
|
||||
f,
|
||||
" (on Windows, environment variables set in PowerShell only persist for the current session; use `setx {primary} <value>` to make it permanent, then open a new terminal, or place a `.env` file containing `{primary}=<value>` in the current working directory)"
|
||||
)?;
|
||||
} else {
|
||||
write!(
|
||||
f,
|
||||
" (on Windows, environment variables set in PowerShell only persist for the current session; use `setx` to make them permanent, then open a new terminal, or place a `.env` file in the current working directory)"
|
||||
)?;
|
||||
}
|
||||
}
|
||||
if let Some(hint) = hint {
|
||||
write!(f, " — hint: {hint}")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
Self::ContextWindowExceeded {
|
||||
model,
|
||||
estimated_input_tokens,
|
||||
requested_output_tokens,
|
||||
estimated_total_tokens,
|
||||
context_window_tokens,
|
||||
} => write!(
|
||||
f,
|
||||
"context_window_blocked for {model}: estimated input {estimated_input_tokens} + requested output {requested_output_tokens} = {estimated_total_tokens} tokens exceeds the {context_window_tokens}-token context window; compact the session or reduce request size before retrying"
|
||||
),
|
||||
Self::ExpiredOAuthToken => {
|
||||
write!(
|
||||
f,
|
||||
"saved OAuth token is expired and no refresh token is available"
|
||||
)
|
||||
}
|
||||
Self::Auth(message) => write!(f, "auth error: {message}"),
|
||||
Self::InvalidApiKeyEnv(error) => {
|
||||
write!(f, "failed to read credential environment variable: {error}")
|
||||
}
|
||||
Self::Http(error) => write!(f, "http error: {error}"),
|
||||
Self::Io(error) => write!(f, "io error: {error}"),
|
||||
Self::Json {
|
||||
provider,
|
||||
model,
|
||||
body_snippet,
|
||||
source,
|
||||
} => write!(
|
||||
f,
|
||||
"failed to parse {provider} response for model {model}: {source}; first 200 chars of body: {body_snippet}"
|
||||
),
|
||||
Self::Api {
|
||||
status,
|
||||
error_type,
|
||||
message,
|
||||
request_id,
|
||||
body,
|
||||
..
|
||||
} => {
|
||||
if let (Some(error_type), Some(message)) = (error_type, message) {
|
||||
write!(f, "api returned {status} ({error_type})")?;
|
||||
if let Some(request_id) = request_id {
|
||||
write!(f, " [trace {request_id}]")?;
|
||||
}
|
||||
write!(f, ": {message}")
|
||||
} else {
|
||||
write!(f, "api returned {status}")?;
|
||||
if let Some(request_id) = request_id {
|
||||
write!(f, " [trace {request_id}]")?;
|
||||
}
|
||||
write!(f, ": {body}")
|
||||
}
|
||||
}
|
||||
Self::RetriesExhausted {
|
||||
attempts,
|
||||
last_error,
|
||||
} => write!(f, "api failed after {attempts} attempts: {last_error}"),
|
||||
Self::InvalidSseFrame(message) => write!(f, "invalid sse frame: {message}"),
|
||||
Self::BackoffOverflow {
|
||||
attempt,
|
||||
base_delay,
|
||||
} => write!(
|
||||
f,
|
||||
"retry backoff overflowed on attempt {attempt} with base delay {base_delay:?}"
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for ApiError {}
|
||||
|
||||
impl From<reqwest::Error> for ApiError {
|
||||
fn from(value: reqwest::Error) -> Self {
|
||||
Self::Http(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<std::io::Error> for ApiError {
|
||||
fn from(value: std::io::Error) -> Self {
|
||||
Self::Io(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for ApiError {
|
||||
fn from(value: serde_json::Error) -> Self {
|
||||
Self::Json {
|
||||
provider: "unknown".to_string(),
|
||||
model: "unknown".to_string(),
|
||||
body_snippet: String::new(),
|
||||
source: value,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<VarError> for ApiError {
|
||||
fn from(value: VarError) -> Self {
|
||||
Self::InvalidApiKeyEnv(value)
|
||||
}
|
||||
}
|
||||
|
||||
fn looks_like_generic_fatal_wrapper(text: &str) -> bool {
|
||||
let lowered = text.to_ascii_lowercase();
|
||||
GENERIC_FATAL_WRAPPER_MARKERS
|
||||
.iter()
|
||||
.any(|marker| lowered.contains(marker))
|
||||
}
|
||||
|
||||
fn looks_like_context_window_error(text: &str) -> bool {
|
||||
let lowered = text.to_ascii_lowercase();
|
||||
CONTEXT_WINDOW_ERROR_MARKERS
|
||||
.iter()
|
||||
.any(|marker| lowered.contains(marker))
|
||||
}
|
||||
|
||||
/// Truncate `body` so the resulting snippet contains at most `max_chars`
|
||||
/// characters (counted by Unicode scalar values, not bytes), preserving the
|
||||
/// leading slice of the body that the caller most often needs to inspect.
|
||||
fn truncate_body_snippet(body: &str, max_chars: usize) -> String {
|
||||
let mut taken_chars = 0;
|
||||
let mut byte_end = 0;
|
||||
for (offset, character) in body.char_indices() {
|
||||
if taken_chars >= max_chars {
|
||||
break;
|
||||
}
|
||||
taken_chars += 1;
|
||||
byte_end = offset + character.len_utf8();
|
||||
}
|
||||
if taken_chars >= max_chars && byte_end < body.len() {
|
||||
format!("{}…", &body[..byte_end])
|
||||
} else {
|
||||
body[..byte_end].to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{truncate_body_snippet, ApiError};
|
||||
|
||||
#[test]
|
||||
fn json_deserialize_error_includes_provider_model_and_truncated_body_snippet() {
|
||||
let raw_body = format!("{}{}", "x".repeat(190), "_TAIL_PAST_200_CHARS_MARKER_");
|
||||
let source = serde_json::from_str::<serde_json::Value>("{not json")
|
||||
.expect_err("invalid json should fail to parse");
|
||||
|
||||
let error = ApiError::json_deserialize("Anthropic", "claude-opus-4-6", &raw_body, source);
|
||||
let rendered = error.to_string();
|
||||
|
||||
assert!(
|
||||
rendered.starts_with("failed to parse Anthropic response for model claude-opus-4-6: "),
|
||||
"rendered error should lead with provider and model: {rendered}"
|
||||
);
|
||||
assert!(
|
||||
rendered.contains("first 200 chars of body: "),
|
||||
"rendered error should label the body snippet: {rendered}"
|
||||
);
|
||||
let snippet = rendered
|
||||
.split("first 200 chars of body: ")
|
||||
.nth(1)
|
||||
.expect("snippet section should be present");
|
||||
assert!(
|
||||
snippet.starts_with(&"x".repeat(190)),
|
||||
"snippet should preserve the leading characters of the body: {snippet}"
|
||||
);
|
||||
assert!(
|
||||
snippet.ends_with('…'),
|
||||
"snippet should signal truncation with an ellipsis: {snippet}"
|
||||
);
|
||||
assert!(
|
||||
!snippet.contains("_TAIL_PAST_200_CHARS_MARKER_"),
|
||||
"snippet should drop characters past the 200-char cap: {snippet}"
|
||||
);
|
||||
assert_eq!(error.safe_failure_class(), "runtime_io");
|
||||
assert_eq!(error.request_id(), None);
|
||||
assert!(!error.is_retryable());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_body_snippet_keeps_short_bodies_intact() {
|
||||
assert_eq!(truncate_body_snippet("hello", 200), "hello");
|
||||
assert_eq!(truncate_body_snippet("", 200), "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_body_snippet_caps_long_bodies_at_max_chars() {
|
||||
let body = "a".repeat(250);
|
||||
let snippet = truncate_body_snippet(&body, 200);
|
||||
assert_eq!(snippet.chars().count(), 201, "200 chars + ellipsis");
|
||||
assert!(snippet.ends_with('…'));
|
||||
assert!(snippet.starts_with(&"a".repeat(200)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_body_snippet_does_not_split_multibyte_characters() {
|
||||
let body = "한글한글한글한글한글한글";
|
||||
let snippet = truncate_body_snippet(body, 4);
|
||||
assert_eq!(snippet, "한글한글…");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_generic_fatal_wrapper_and_classifies_it_as_provider_internal() {
|
||||
let error = ApiError::Api {
|
||||
status: reqwest::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
error_type: Some("api_error".to_string()),
|
||||
message: Some(
|
||||
"Something went wrong while processing your request. Please try again, or use /new to start a fresh session."
|
||||
.to_string(),
|
||||
),
|
||||
request_id: Some("req_jobdori_123".to_string()),
|
||||
body: String::new(),
|
||||
retryable: true,
|
||||
};
|
||||
|
||||
assert!(error.is_generic_fatal_wrapper());
|
||||
assert_eq!(error.safe_failure_class(), "provider_internal");
|
||||
assert_eq!(error.request_id(), Some("req_jobdori_123"));
|
||||
assert!(error.to_string().contains("[trace req_jobdori_123]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn retries_exhausted_preserves_nested_request_id_and_failure_class() {
|
||||
let error = ApiError::RetriesExhausted {
|
||||
attempts: 3,
|
||||
last_error: Box::new(ApiError::Api {
|
||||
status: reqwest::StatusCode::BAD_GATEWAY,
|
||||
error_type: Some("api_error".to_string()),
|
||||
message: Some(
|
||||
"Something went wrong while processing your request. Please try again, or use /new to start a fresh session."
|
||||
.to_string(),
|
||||
),
|
||||
request_id: Some("req_nested_456".to_string()),
|
||||
body: String::new(),
|
||||
retryable: true,
|
||||
}),
|
||||
};
|
||||
|
||||
assert!(error.is_generic_fatal_wrapper());
|
||||
assert_eq!(error.safe_failure_class(), "provider_retry_exhausted");
|
||||
assert_eq!(error.request_id(), Some("req_nested_456"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_provider_context_window_errors() {
|
||||
let error = ApiError::Api {
|
||||
status: reqwest::StatusCode::BAD_REQUEST,
|
||||
error_type: Some("invalid_request_error".to_string()),
|
||||
message: Some(
|
||||
"This model's maximum context length is 200000 tokens, but your request used 230000 tokens."
|
||||
.to_string(),
|
||||
),
|
||||
request_id: Some("req_ctx_123".to_string()),
|
||||
body: String::new(),
|
||||
retryable: false,
|
||||
};
|
||||
|
||||
assert!(error.is_context_window_failure());
|
||||
assert_eq!(error.safe_failure_class(), "context_window");
|
||||
assert_eq!(error.request_id(), Some("req_ctx_123"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_credentials_without_hint_renders_the_canonical_message() {
|
||||
// given
|
||||
let error = ApiError::missing_credentials(
|
||||
"Anthropic",
|
||||
&["ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY"],
|
||||
);
|
||||
|
||||
// when
|
||||
let rendered = error.to_string();
|
||||
|
||||
// then
|
||||
assert!(
|
||||
rendered.starts_with(
|
||||
"missing Anthropic credentials; export ANTHROPIC_AUTH_TOKEN or ANTHROPIC_API_KEY before calling the Anthropic API"
|
||||
),
|
||||
"rendered error should lead with the canonical missing-credential message: {rendered}"
|
||||
);
|
||||
assert!(
|
||||
!rendered.contains(" — hint: "),
|
||||
"no hint should be appended when none is supplied: {rendered}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_credentials_with_hint_appends_the_hint_after_base_message() {
|
||||
// given
|
||||
let error = ApiError::missing_credentials_with_hint(
|
||||
"Anthropic",
|
||||
&["ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY"],
|
||||
"I see OPENAI_API_KEY is set — if you meant to use the OpenAI-compat provider, prefix your model name with `openai/` so prefix routing selects it.",
|
||||
);
|
||||
|
||||
// when
|
||||
let rendered = error.to_string();
|
||||
|
||||
// then
|
||||
assert!(
|
||||
rendered.starts_with("missing Anthropic credentials;"),
|
||||
"hint should be appended, not replace the base message: {rendered}"
|
||||
);
|
||||
let hint_marker = " — hint: I see OPENAI_API_KEY is set — if you meant to use the OpenAI-compat provider, prefix your model name with `openai/` so prefix routing selects it.";
|
||||
assert!(
|
||||
rendered.ends_with(hint_marker),
|
||||
"rendered error should end with the hint: {rendered}"
|
||||
);
|
||||
// Classification semantics are unaffected by the presence of a hint.
|
||||
assert_eq!(error.safe_failure_class(), "provider_auth");
|
||||
assert!(!error.is_retryable());
|
||||
assert_eq!(error.request_id(), None);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,344 @@
|
||||
use crate::error::ApiError;
|
||||
|
||||
const HTTP_PROXY_KEYS: [&str; 2] = ["HTTP_PROXY", "http_proxy"];
|
||||
const HTTPS_PROXY_KEYS: [&str; 2] = ["HTTPS_PROXY", "https_proxy"];
|
||||
const NO_PROXY_KEYS: [&str; 2] = ["NO_PROXY", "no_proxy"];
|
||||
|
||||
/// Snapshot of the proxy-related environment variables that influence the
|
||||
/// outbound HTTP client. Captured up front so callers can inspect, log, and
|
||||
/// test the resolved configuration without re-reading the process environment.
|
||||
///
|
||||
/// When `proxy_url` is set it acts as a single catch-all proxy for both
|
||||
/// HTTP and HTTPS traffic, taking precedence over the per-scheme fields.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct ProxyConfig {
|
||||
pub http_proxy: Option<String>,
|
||||
pub https_proxy: Option<String>,
|
||||
pub no_proxy: Option<String>,
|
||||
/// Optional unified proxy URL that applies to both HTTP and HTTPS.
|
||||
/// When set, this takes precedence over `http_proxy` and `https_proxy`.
|
||||
pub proxy_url: Option<String>,
|
||||
}
|
||||
|
||||
impl ProxyConfig {
|
||||
/// Read proxy settings from the live process environment, honouring both
|
||||
/// the upper- and lower-case spellings used by curl, git, and friends.
|
||||
#[must_use]
|
||||
pub fn from_env() -> Self {
|
||||
Self::from_lookup(|key| std::env::var(key).ok())
|
||||
}
|
||||
|
||||
/// Create a proxy configuration from a single URL that applies to both
|
||||
/// HTTP and HTTPS traffic. This is the config-file alternative to setting
|
||||
/// `HTTP_PROXY` and `HTTPS_PROXY` environment variables separately.
|
||||
#[must_use]
|
||||
pub fn from_proxy_url(url: impl Into<String>) -> Self {
|
||||
Self {
|
||||
proxy_url: Some(url.into()),
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn from_lookup<F>(mut lookup: F) -> Self
|
||||
where
|
||||
F: FnMut(&str) -> Option<String>,
|
||||
{
|
||||
Self {
|
||||
http_proxy: first_non_empty(&HTTP_PROXY_KEYS, &mut lookup),
|
||||
https_proxy: first_non_empty(&HTTPS_PROXY_KEYS, &mut lookup),
|
||||
no_proxy: first_non_empty(&NO_PROXY_KEYS, &mut lookup),
|
||||
proxy_url: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.proxy_url.is_none() && self.http_proxy.is_none() && self.https_proxy.is_none()
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a `reqwest::Client` that honours the standard `HTTP_PROXY`,
|
||||
/// `HTTPS_PROXY`, and `NO_PROXY` environment variables. When no proxy is
|
||||
/// configured the client behaves identically to `reqwest::Client::new()`.
|
||||
pub fn build_http_client() -> Result<reqwest::Client, ApiError> {
|
||||
build_http_client_with(&ProxyConfig::from_env())
|
||||
}
|
||||
|
||||
/// Infallible counterpart to [`build_http_client`] for constructors that
|
||||
/// historically returned `Self` rather than `Result<Self, _>`. When the proxy
|
||||
/// configuration is malformed we fall back to a default client so that
|
||||
/// callers retain the previous behaviour and the failure surfaces on the
|
||||
/// first outbound request instead of at construction time.
|
||||
#[must_use]
|
||||
pub fn build_http_client_or_default() -> reqwest::Client {
|
||||
build_http_client().unwrap_or_else(|_| reqwest::Client::new())
|
||||
}
|
||||
|
||||
/// Build a `reqwest::Client` from an explicit [`ProxyConfig`]. Used by tests
|
||||
/// and by callers that want to override process-level environment lookups.
|
||||
///
|
||||
/// When `config.proxy_url` is set it overrides the per-scheme `http_proxy`
|
||||
/// and `https_proxy` fields and is registered as both an HTTP and HTTPS
|
||||
/// proxy so a single value can route every outbound request.
|
||||
pub fn build_http_client_with(config: &ProxyConfig) -> Result<reqwest::Client, ApiError> {
|
||||
let mut builder = reqwest::Client::builder().no_proxy();
|
||||
|
||||
let no_proxy = config
|
||||
.no_proxy
|
||||
.as_deref()
|
||||
.and_then(reqwest::NoProxy::from_string);
|
||||
|
||||
let (http_proxy_url, https_proxy_url) = match config.proxy_url.as_deref() {
|
||||
Some(unified) => (Some(unified), Some(unified)),
|
||||
None => (config.http_proxy.as_deref(), config.https_proxy.as_deref()),
|
||||
};
|
||||
|
||||
if let Some(url) = https_proxy_url {
|
||||
let mut proxy = reqwest::Proxy::https(url)?;
|
||||
if let Some(filter) = no_proxy.clone() {
|
||||
proxy = proxy.no_proxy(Some(filter));
|
||||
}
|
||||
builder = builder.proxy(proxy);
|
||||
}
|
||||
|
||||
if let Some(url) = http_proxy_url {
|
||||
let mut proxy = reqwest::Proxy::http(url)?;
|
||||
if let Some(filter) = no_proxy.clone() {
|
||||
proxy = proxy.no_proxy(Some(filter));
|
||||
}
|
||||
builder = builder.proxy(proxy);
|
||||
}
|
||||
|
||||
Ok(builder.build()?)
|
||||
}
|
||||
|
||||
fn first_non_empty<F>(keys: &[&str], lookup: &mut F) -> Option<String>
|
||||
where
|
||||
F: FnMut(&str) -> Option<String>,
|
||||
{
|
||||
keys.iter()
|
||||
.find_map(|key| lookup(key).filter(|value| !value.is_empty()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
||||
use super::{build_http_client_with, ProxyConfig};
|
||||
|
||||
fn config_from_map(pairs: &[(&str, &str)]) -> ProxyConfig {
|
||||
let map: HashMap<String, String> = pairs
|
||||
.iter()
|
||||
.map(|(key, value)| ((*key).to_string(), (*value).to_string()))
|
||||
.collect();
|
||||
ProxyConfig::from_lookup(|key| map.get(key).cloned())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_is_empty_when_no_env_vars_are_set() {
|
||||
// given
|
||||
let config = config_from_map(&[]);
|
||||
|
||||
// when
|
||||
let empty = config.is_empty();
|
||||
|
||||
// then
|
||||
assert!(empty);
|
||||
assert_eq!(config, ProxyConfig::default());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_reads_uppercase_http_https_and_no_proxy() {
|
||||
// given
|
||||
let pairs = [
|
||||
("HTTP_PROXY", "http://proxy.internal:3128"),
|
||||
("HTTPS_PROXY", "http://secure.internal:3129"),
|
||||
("NO_PROXY", "localhost,127.0.0.1,.corp"),
|
||||
];
|
||||
|
||||
// when
|
||||
let config = config_from_map(&pairs);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
config.http_proxy.as_deref(),
|
||||
Some("http://proxy.internal:3128")
|
||||
);
|
||||
assert_eq!(
|
||||
config.https_proxy.as_deref(),
|
||||
Some("http://secure.internal:3129")
|
||||
);
|
||||
assert_eq!(
|
||||
config.no_proxy.as_deref(),
|
||||
Some("localhost,127.0.0.1,.corp")
|
||||
);
|
||||
assert!(!config.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_falls_back_to_lowercase_keys() {
|
||||
// given
|
||||
let pairs = [
|
||||
("http_proxy", "http://lower.internal:3128"),
|
||||
("https_proxy", "http://lower-secure.internal:3129"),
|
||||
("no_proxy", ".lower"),
|
||||
];
|
||||
|
||||
// when
|
||||
let config = config_from_map(&pairs);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
config.http_proxy.as_deref(),
|
||||
Some("http://lower.internal:3128")
|
||||
);
|
||||
assert_eq!(
|
||||
config.https_proxy.as_deref(),
|
||||
Some("http://lower-secure.internal:3129")
|
||||
);
|
||||
assert_eq!(config.no_proxy.as_deref(), Some(".lower"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_prefers_uppercase_over_lowercase_when_both_set() {
|
||||
// given
|
||||
let pairs = [
|
||||
("HTTP_PROXY", "http://upper.internal:3128"),
|
||||
("http_proxy", "http://lower.internal:3128"),
|
||||
];
|
||||
|
||||
// when
|
||||
let config = config_from_map(&pairs);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
config.http_proxy.as_deref(),
|
||||
Some("http://upper.internal:3128")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_treats_empty_strings_as_unset() {
|
||||
// given
|
||||
let pairs = [("HTTP_PROXY", ""), ("http_proxy", "")];
|
||||
|
||||
// when
|
||||
let config = config_from_map(&pairs);
|
||||
|
||||
// then
|
||||
assert!(config.http_proxy.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_http_client_succeeds_when_no_proxy_is_configured() {
|
||||
// given
|
||||
let config = ProxyConfig::default();
|
||||
|
||||
// when
|
||||
let result = build_http_client_with(&config);
|
||||
|
||||
// then
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_http_client_succeeds_with_valid_http_and_https_proxies() {
|
||||
// given
|
||||
let config = ProxyConfig {
|
||||
http_proxy: Some("http://proxy.internal:3128".to_string()),
|
||||
https_proxy: Some("http://secure.internal:3129".to_string()),
|
||||
no_proxy: Some("localhost,127.0.0.1".to_string()),
|
||||
proxy_url: None,
|
||||
};
|
||||
|
||||
// when
|
||||
let result = build_http_client_with(&config);
|
||||
|
||||
// then
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_http_client_returns_http_error_for_invalid_proxy_url() {
|
||||
// given
|
||||
let config = ProxyConfig {
|
||||
http_proxy: None,
|
||||
https_proxy: Some("not a url".to_string()),
|
||||
no_proxy: None,
|
||||
proxy_url: None,
|
||||
};
|
||||
|
||||
// when
|
||||
let result = build_http_client_with(&config);
|
||||
|
||||
// then
|
||||
let error = result.expect_err("invalid proxy URL must be reported as a build failure");
|
||||
assert!(
|
||||
matches!(error, crate::error::ApiError::Http(_)),
|
||||
"expected ApiError::Http for invalid proxy URL, got: {error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_proxy_url_sets_unified_field_and_leaves_per_scheme_empty() {
|
||||
// given / when
|
||||
let config = ProxyConfig::from_proxy_url("http://unified.internal:3128");
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
config.proxy_url.as_deref(),
|
||||
Some("http://unified.internal:3128")
|
||||
);
|
||||
assert!(config.http_proxy.is_none());
|
||||
assert!(config.https_proxy.is_none());
|
||||
assert!(!config.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_http_client_succeeds_with_unified_proxy_url() {
|
||||
// given
|
||||
let config = ProxyConfig {
|
||||
proxy_url: Some("http://unified.internal:3128".to_string()),
|
||||
no_proxy: Some("localhost".to_string()),
|
||||
..ProxyConfig::default()
|
||||
};
|
||||
|
||||
// when
|
||||
let result = build_http_client_with(&config);
|
||||
|
||||
// then
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_url_takes_precedence_over_per_scheme_fields() {
|
||||
// given – both per-scheme and unified are set
|
||||
let config = ProxyConfig {
|
||||
http_proxy: Some("http://per-scheme.internal:1111".to_string()),
|
||||
https_proxy: Some("http://per-scheme.internal:2222".to_string()),
|
||||
no_proxy: None,
|
||||
proxy_url: Some("http://unified.internal:3128".to_string()),
|
||||
};
|
||||
|
||||
// when – building succeeds (the unified URL is valid)
|
||||
let result = build_http_client_with(&config);
|
||||
|
||||
// then
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_http_client_returns_error_for_invalid_unified_proxy_url() {
|
||||
// given
|
||||
let config = ProxyConfig::from_proxy_url("not a url");
|
||||
|
||||
// when
|
||||
let result = build_http_client_with(&config);
|
||||
|
||||
// then
|
||||
assert!(
|
||||
matches!(result, Err(crate::error::ApiError::Http(_))),
|
||||
"invalid unified proxy URL should fail: {result:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
mod client;
|
||||
mod error;
|
||||
mod http_client;
|
||||
mod prompt_cache;
|
||||
mod providers;
|
||||
mod sse;
|
||||
mod types;
|
||||
|
||||
pub use client::{
|
||||
oauth_token_is_expired, read_base_url, read_xai_base_url, resolve_saved_oauth_token,
|
||||
resolve_startup_auth_source, MessageStream, OAuthTokenSet, ProviderClient,
|
||||
};
|
||||
pub use error::ApiError;
|
||||
pub use http_client::{
|
||||
build_http_client, build_http_client_or_default, build_http_client_with, ProxyConfig,
|
||||
};
|
||||
pub use prompt_cache::{
|
||||
CacheBreakEvent, PromptCache, PromptCacheConfig, PromptCachePaths, PromptCacheRecord,
|
||||
PromptCacheStats,
|
||||
};
|
||||
pub use providers::anthropic::{AnthropicClient, AnthropicClient as ApiClient, AuthSource};
|
||||
pub use providers::openai_compat::{OpenAiCompatClient, OpenAiCompatConfig};
|
||||
pub use providers::{
|
||||
detect_provider_kind, max_tokens_for_model, max_tokens_for_model_with_override,
|
||||
resolve_model_alias, ProviderKind,
|
||||
};
|
||||
pub use sse::{parse_frame, SseParser};
|
||||
pub use types::{
|
||||
ContentBlockDelta, ContentBlockDeltaEvent, ContentBlockStartEvent, ContentBlockStopEvent,
|
||||
InputContentBlock, InputMessage, MessageDelta, MessageDeltaEvent, MessageRequest,
|
||||
MessageResponse, MessageStartEvent, MessageStopEvent, OutputContentBlock, StreamEvent,
|
||||
ToolChoice, ToolDefinition, ToolResultContentBlock, Usage,
|
||||
};
|
||||
|
||||
pub use telemetry::{
|
||||
AnalyticsEvent, AnthropicRequestProfile, ClientIdentity, JsonlTelemetrySink,
|
||||
MemoryTelemetrySink, SessionTraceRecord, SessionTracer, TelemetryEvent, TelemetrySink,
|
||||
DEFAULT_ANTHROPIC_VERSION,
|
||||
};
|
||||
@@ -0,0 +1,735 @@
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::types::{MessageRequest, MessageResponse, Usage};
|
||||
|
||||
const DEFAULT_COMPLETION_TTL_SECS: u64 = 30;
|
||||
const DEFAULT_PROMPT_TTL_SECS: u64 = 5 * 60;
|
||||
const DEFAULT_BREAK_MIN_DROP: u32 = 2_000;
|
||||
const MAX_SANITIZED_LENGTH: usize = 80;
|
||||
const REQUEST_FINGERPRINT_VERSION: u32 = 1;
|
||||
const REQUEST_FINGERPRINT_PREFIX: &str = "v1";
|
||||
const FNV_OFFSET_BASIS: u64 = 0xcbf2_9ce4_8422_2325;
|
||||
const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptCacheConfig {
|
||||
pub session_id: String,
|
||||
pub completion_ttl: Duration,
|
||||
pub prompt_ttl: Duration,
|
||||
pub cache_break_min_drop: u32,
|
||||
}
|
||||
|
||||
impl PromptCacheConfig {
|
||||
#[must_use]
|
||||
pub fn new(session_id: impl Into<String>) -> Self {
|
||||
Self {
|
||||
session_id: session_id.into(),
|
||||
completion_ttl: Duration::from_secs(DEFAULT_COMPLETION_TTL_SECS),
|
||||
prompt_ttl: Duration::from_secs(DEFAULT_PROMPT_TTL_SECS),
|
||||
cache_break_min_drop: DEFAULT_BREAK_MIN_DROP,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for PromptCacheConfig {
|
||||
fn default() -> Self {
|
||||
Self::new("default")
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct PromptCachePaths {
|
||||
pub root: PathBuf,
|
||||
pub session_dir: PathBuf,
|
||||
pub completion_dir: PathBuf,
|
||||
pub session_state_path: PathBuf,
|
||||
pub stats_path: PathBuf,
|
||||
}
|
||||
|
||||
impl PromptCachePaths {
|
||||
#[must_use]
|
||||
pub fn for_session(session_id: &str) -> Self {
|
||||
let root = base_cache_root();
|
||||
let session_dir = root.join(sanitize_path_segment(session_id));
|
||||
let completion_dir = session_dir.join("completions");
|
||||
Self {
|
||||
root,
|
||||
session_state_path: session_dir.join("session-state.json"),
|
||||
stats_path: session_dir.join("stats.json"),
|
||||
session_dir,
|
||||
completion_dir,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn completion_entry_path(&self, request_hash: &str) -> PathBuf {
|
||||
self.completion_dir.join(format!("{request_hash}.json"))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct PromptCacheStats {
|
||||
pub tracked_requests: u64,
|
||||
pub completion_cache_hits: u64,
|
||||
pub completion_cache_misses: u64,
|
||||
pub completion_cache_writes: u64,
|
||||
pub expected_invalidations: u64,
|
||||
pub unexpected_cache_breaks: u64,
|
||||
pub total_cache_creation_input_tokens: u64,
|
||||
pub total_cache_read_input_tokens: u64,
|
||||
pub last_cache_creation_input_tokens: Option<u32>,
|
||||
pub last_cache_read_input_tokens: Option<u32>,
|
||||
pub last_request_hash: Option<String>,
|
||||
pub last_completion_cache_key: Option<String>,
|
||||
pub last_break_reason: Option<String>,
|
||||
pub last_cache_source: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct CacheBreakEvent {
|
||||
pub unexpected: bool,
|
||||
pub reason: String,
|
||||
pub previous_cache_read_input_tokens: u32,
|
||||
pub current_cache_read_input_tokens: u32,
|
||||
pub token_drop: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PromptCacheRecord {
|
||||
pub cache_break: Option<CacheBreakEvent>,
|
||||
pub stats: PromptCacheStats,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PromptCache {
|
||||
inner: Arc<Mutex<PromptCacheInner>>,
|
||||
}
|
||||
|
||||
impl PromptCache {
|
||||
#[must_use]
|
||||
pub fn new(session_id: impl Into<String>) -> Self {
|
||||
Self::with_config(PromptCacheConfig::new(session_id))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_config(config: PromptCacheConfig) -> Self {
|
||||
let paths = PromptCachePaths::for_session(&config.session_id);
|
||||
let stats = read_json::<PromptCacheStats>(&paths.stats_path).unwrap_or_default();
|
||||
let previous = read_json::<TrackedPromptState>(&paths.session_state_path);
|
||||
Self {
|
||||
inner: Arc::new(Mutex::new(PromptCacheInner {
|
||||
config,
|
||||
paths,
|
||||
stats,
|
||||
previous,
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn paths(&self) -> PromptCachePaths {
|
||||
self.lock().paths.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn stats(&self) -> PromptCacheStats {
|
||||
self.lock().stats.clone()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn lookup_completion(&self, request: &MessageRequest) -> Option<MessageResponse> {
|
||||
let request_hash = request_hash_hex(request);
|
||||
let (paths, ttl) = {
|
||||
let inner = self.lock();
|
||||
(inner.paths.clone(), inner.config.completion_ttl)
|
||||
};
|
||||
let entry_path = paths.completion_entry_path(&request_hash);
|
||||
let entry = read_json::<CompletionCacheEntry>(&entry_path);
|
||||
let Some(entry) = entry else {
|
||||
let mut inner = self.lock();
|
||||
inner.stats.completion_cache_misses += 1;
|
||||
inner.stats.last_completion_cache_key = Some(request_hash);
|
||||
persist_state(&inner);
|
||||
return None;
|
||||
};
|
||||
|
||||
if entry.fingerprint_version != current_fingerprint_version() {
|
||||
let mut inner = self.lock();
|
||||
inner.stats.completion_cache_misses += 1;
|
||||
inner.stats.last_completion_cache_key = Some(request_hash.clone());
|
||||
let _ = fs::remove_file(entry_path);
|
||||
persist_state(&inner);
|
||||
return None;
|
||||
}
|
||||
|
||||
let expired = now_unix_secs().saturating_sub(entry.cached_at_unix_secs) >= ttl.as_secs();
|
||||
let mut inner = self.lock();
|
||||
inner.stats.last_completion_cache_key = Some(request_hash.clone());
|
||||
if expired {
|
||||
inner.stats.completion_cache_misses += 1;
|
||||
let _ = fs::remove_file(entry_path);
|
||||
persist_state(&inner);
|
||||
return None;
|
||||
}
|
||||
|
||||
inner.stats.completion_cache_hits += 1;
|
||||
apply_usage_to_stats(
|
||||
&mut inner.stats,
|
||||
&entry.response.usage,
|
||||
&request_hash,
|
||||
"completion-cache",
|
||||
);
|
||||
inner.previous = Some(TrackedPromptState::from_usage(
|
||||
request,
|
||||
&entry.response.usage,
|
||||
));
|
||||
persist_state(&inner);
|
||||
Some(entry.response)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn record_response(
|
||||
&self,
|
||||
request: &MessageRequest,
|
||||
response: &MessageResponse,
|
||||
) -> PromptCacheRecord {
|
||||
self.record_usage_internal(request, &response.usage, Some(response))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn record_usage(&self, request: &MessageRequest, usage: &Usage) -> PromptCacheRecord {
|
||||
self.record_usage_internal(request, usage, None)
|
||||
}
|
||||
|
||||
fn record_usage_internal(
|
||||
&self,
|
||||
request: &MessageRequest,
|
||||
usage: &Usage,
|
||||
response: Option<&MessageResponse>,
|
||||
) -> PromptCacheRecord {
|
||||
let request_hash = request_hash_hex(request);
|
||||
let mut inner = self.lock();
|
||||
let previous = inner.previous.clone();
|
||||
let current = TrackedPromptState::from_usage(request, usage);
|
||||
let cache_break = detect_cache_break(&inner.config, previous.as_ref(), ¤t);
|
||||
|
||||
inner.stats.tracked_requests += 1;
|
||||
apply_usage_to_stats(&mut inner.stats, usage, &request_hash, "api-response");
|
||||
if let Some(event) = &cache_break {
|
||||
if event.unexpected {
|
||||
inner.stats.unexpected_cache_breaks += 1;
|
||||
} else {
|
||||
inner.stats.expected_invalidations += 1;
|
||||
}
|
||||
inner.stats.last_break_reason = Some(event.reason.clone());
|
||||
}
|
||||
|
||||
inner.previous = Some(current);
|
||||
if let Some(response) = response {
|
||||
write_completion_entry(&inner.paths, &request_hash, response);
|
||||
inner.stats.completion_cache_writes += 1;
|
||||
}
|
||||
persist_state(&inner);
|
||||
|
||||
PromptCacheRecord {
|
||||
cache_break,
|
||||
stats: inner.stats.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
fn lock(&self) -> std::sync::MutexGuard<'_, PromptCacheInner> {
|
||||
self.inner
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct PromptCacheInner {
|
||||
config: PromptCacheConfig,
|
||||
paths: PromptCachePaths,
|
||||
stats: PromptCacheStats,
|
||||
previous: Option<TrackedPromptState>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
struct CompletionCacheEntry {
|
||||
cached_at_unix_secs: u64,
|
||||
#[serde(default = "current_fingerprint_version")]
|
||||
fingerprint_version: u32,
|
||||
response: MessageResponse,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
struct TrackedPromptState {
|
||||
observed_at_unix_secs: u64,
|
||||
#[serde(default = "current_fingerprint_version")]
|
||||
fingerprint_version: u32,
|
||||
model_hash: u64,
|
||||
system_hash: u64,
|
||||
tools_hash: u64,
|
||||
messages_hash: u64,
|
||||
cache_read_input_tokens: u32,
|
||||
}
|
||||
|
||||
impl TrackedPromptState {
|
||||
fn from_usage(request: &MessageRequest, usage: &Usage) -> Self {
|
||||
let hashes = RequestFingerprints::from_request(request);
|
||||
Self {
|
||||
observed_at_unix_secs: now_unix_secs(),
|
||||
fingerprint_version: current_fingerprint_version(),
|
||||
model_hash: hashes.model,
|
||||
system_hash: hashes.system,
|
||||
tools_hash: hashes.tools,
|
||||
messages_hash: hashes.messages,
|
||||
cache_read_input_tokens: usage.cache_read_input_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct RequestFingerprints {
|
||||
model: u64,
|
||||
system: u64,
|
||||
tools: u64,
|
||||
messages: u64,
|
||||
}
|
||||
|
||||
impl RequestFingerprints {
|
||||
fn from_request(request: &MessageRequest) -> Self {
|
||||
Self {
|
||||
model: hash_serializable(&request.model),
|
||||
system: hash_serializable(&request.system),
|
||||
tools: hash_serializable(&request.tools),
|
||||
messages: hash_serializable(&request.messages),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn detect_cache_break(
|
||||
config: &PromptCacheConfig,
|
||||
previous: Option<&TrackedPromptState>,
|
||||
current: &TrackedPromptState,
|
||||
) -> Option<CacheBreakEvent> {
|
||||
let previous = previous?;
|
||||
if previous.fingerprint_version != current.fingerprint_version {
|
||||
return Some(CacheBreakEvent {
|
||||
unexpected: false,
|
||||
reason: format!(
|
||||
"fingerprint version changed (v{} -> v{})",
|
||||
previous.fingerprint_version, current.fingerprint_version
|
||||
),
|
||||
previous_cache_read_input_tokens: previous.cache_read_input_tokens,
|
||||
current_cache_read_input_tokens: current.cache_read_input_tokens,
|
||||
token_drop: previous
|
||||
.cache_read_input_tokens
|
||||
.saturating_sub(current.cache_read_input_tokens),
|
||||
});
|
||||
}
|
||||
let token_drop = previous
|
||||
.cache_read_input_tokens
|
||||
.saturating_sub(current.cache_read_input_tokens);
|
||||
if token_drop < config.cache_break_min_drop {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut reasons = Vec::new();
|
||||
if previous.model_hash != current.model_hash {
|
||||
reasons.push("model changed");
|
||||
}
|
||||
if previous.system_hash != current.system_hash {
|
||||
reasons.push("system prompt changed");
|
||||
}
|
||||
if previous.tools_hash != current.tools_hash {
|
||||
reasons.push("tool definitions changed");
|
||||
}
|
||||
if previous.messages_hash != current.messages_hash {
|
||||
reasons.push("message payload changed");
|
||||
}
|
||||
|
||||
let elapsed = current
|
||||
.observed_at_unix_secs
|
||||
.saturating_sub(previous.observed_at_unix_secs);
|
||||
|
||||
let (unexpected, reason) = if reasons.is_empty() {
|
||||
if elapsed > config.prompt_ttl.as_secs() {
|
||||
(
|
||||
false,
|
||||
format!("possible prompt cache TTL expiry after {elapsed}s"),
|
||||
)
|
||||
} else {
|
||||
(
|
||||
true,
|
||||
"cache read tokens dropped while prompt fingerprint remained stable".to_string(),
|
||||
)
|
||||
}
|
||||
} else {
|
||||
(false, reasons.join(", "))
|
||||
};
|
||||
|
||||
Some(CacheBreakEvent {
|
||||
unexpected,
|
||||
reason,
|
||||
previous_cache_read_input_tokens: previous.cache_read_input_tokens,
|
||||
current_cache_read_input_tokens: current.cache_read_input_tokens,
|
||||
token_drop,
|
||||
})
|
||||
}
|
||||
|
||||
fn apply_usage_to_stats(
|
||||
stats: &mut PromptCacheStats,
|
||||
usage: &Usage,
|
||||
request_hash: &str,
|
||||
source: &str,
|
||||
) {
|
||||
stats.total_cache_creation_input_tokens += u64::from(usage.cache_creation_input_tokens);
|
||||
stats.total_cache_read_input_tokens += u64::from(usage.cache_read_input_tokens);
|
||||
stats.last_cache_creation_input_tokens = Some(usage.cache_creation_input_tokens);
|
||||
stats.last_cache_read_input_tokens = Some(usage.cache_read_input_tokens);
|
||||
stats.last_request_hash = Some(request_hash.to_string());
|
||||
stats.last_cache_source = Some(source.to_string());
|
||||
}
|
||||
|
||||
fn persist_state(inner: &PromptCacheInner) {
|
||||
let _ = ensure_cache_dirs(&inner.paths);
|
||||
let _ = write_json(&inner.paths.stats_path, &inner.stats);
|
||||
if let Some(previous) = &inner.previous {
|
||||
let _ = write_json(&inner.paths.session_state_path, previous);
|
||||
}
|
||||
}
|
||||
|
||||
fn write_completion_entry(
|
||||
paths: &PromptCachePaths,
|
||||
request_hash: &str,
|
||||
response: &MessageResponse,
|
||||
) {
|
||||
let _ = ensure_cache_dirs(paths);
|
||||
let entry = CompletionCacheEntry {
|
||||
cached_at_unix_secs: now_unix_secs(),
|
||||
fingerprint_version: current_fingerprint_version(),
|
||||
response: response.clone(),
|
||||
};
|
||||
let _ = write_json(&paths.completion_entry_path(request_hash), &entry);
|
||||
}
|
||||
|
||||
fn ensure_cache_dirs(paths: &PromptCachePaths) -> std::io::Result<()> {
|
||||
fs::create_dir_all(&paths.completion_dir)
|
||||
}
|
||||
|
||||
fn write_json<T: Serialize>(path: &Path, value: &T) -> std::io::Result<()> {
|
||||
let json = serde_json::to_vec_pretty(value)
|
||||
.map_err(|error| std::io::Error::new(std::io::ErrorKind::InvalidData, error))?;
|
||||
fs::write(path, json)
|
||||
}
|
||||
|
||||
fn read_json<T: for<'de> Deserialize<'de>>(path: &Path) -> Option<T> {
|
||||
let bytes = fs::read(path).ok()?;
|
||||
serde_json::from_slice(&bytes).ok()
|
||||
}
|
||||
|
||||
fn request_hash_hex(request: &MessageRequest) -> String {
|
||||
format!(
|
||||
"{REQUEST_FINGERPRINT_PREFIX}-{:016x}",
|
||||
hash_serializable(request)
|
||||
)
|
||||
}
|
||||
|
||||
fn hash_serializable<T: Serialize>(value: &T) -> u64 {
|
||||
let json = serde_json::to_vec(value).unwrap_or_default();
|
||||
stable_hash_bytes(&json)
|
||||
}
|
||||
|
||||
fn sanitize_path_segment(value: &str) -> String {
|
||||
let sanitized: String = value
|
||||
.chars()
|
||||
.map(|ch| if ch.is_ascii_alphanumeric() { ch } else { '-' })
|
||||
.collect();
|
||||
if sanitized.len() <= MAX_SANITIZED_LENGTH {
|
||||
return sanitized;
|
||||
}
|
||||
let suffix = format!("-{:x}", hash_string(value));
|
||||
format!(
|
||||
"{}{}",
|
||||
&sanitized[..MAX_SANITIZED_LENGTH.saturating_sub(suffix.len())],
|
||||
suffix
|
||||
)
|
||||
}
|
||||
|
||||
fn hash_string(value: &str) -> u64 {
|
||||
stable_hash_bytes(value.as_bytes())
|
||||
}
|
||||
|
||||
fn base_cache_root() -> PathBuf {
|
||||
if let Some(config_home) = std::env::var_os("CLAUDE_CONFIG_HOME") {
|
||||
return PathBuf::from(config_home)
|
||||
.join("cache")
|
||||
.join("prompt-cache");
|
||||
}
|
||||
if let Some(home) = std::env::var_os("HOME") {
|
||||
return PathBuf::from(home)
|
||||
.join(".claude")
|
||||
.join("cache")
|
||||
.join("prompt-cache");
|
||||
}
|
||||
std::env::temp_dir().join("claude-prompt-cache")
|
||||
}
|
||||
|
||||
fn now_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_or(0, |duration| duration.as_secs())
|
||||
}
|
||||
|
||||
const fn current_fingerprint_version() -> u32 {
|
||||
REQUEST_FINGERPRINT_VERSION
|
||||
}
|
||||
|
||||
fn stable_hash_bytes(bytes: &[u8]) -> u64 {
|
||||
let mut hash = FNV_OFFSET_BASIS;
|
||||
for byte in bytes {
|
||||
hash ^= u64::from(*byte);
|
||||
hash = hash.wrapping_mul(FNV_PRIME);
|
||||
}
|
||||
hash
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use super::{
|
||||
detect_cache_break, read_json, request_hash_hex, sanitize_path_segment, PromptCache,
|
||||
PromptCacheConfig, PromptCachePaths, TrackedPromptState, REQUEST_FINGERPRINT_PREFIX,
|
||||
};
|
||||
use crate::types::{InputMessage, MessageRequest, MessageResponse, OutputContentBlock, Usage};
|
||||
|
||||
fn test_env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
||||
LOCK.get_or_init(|| Mutex::new(()))
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn path_builder_sanitizes_session_identifier() {
|
||||
let paths = PromptCachePaths::for_session("session:/with spaces");
|
||||
let session_dir = paths
|
||||
.session_dir
|
||||
.file_name()
|
||||
.and_then(|value| value.to_str())
|
||||
.expect("session dir name");
|
||||
assert_eq!(session_dir, "session--with-spaces");
|
||||
assert!(paths.completion_dir.ends_with("completions"));
|
||||
assert!(paths.stats_path.ends_with("stats.json"));
|
||||
assert!(paths.session_state_path.ends_with("session-state.json"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_fingerprint_drives_unexpected_break_detection() {
|
||||
let request = sample_request("same");
|
||||
let previous = TrackedPromptState::from_usage(
|
||||
&request,
|
||||
&Usage {
|
||||
input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 6_000,
|
||||
output_tokens: 0,
|
||||
},
|
||||
);
|
||||
let current = TrackedPromptState::from_usage(
|
||||
&request,
|
||||
&Usage {
|
||||
input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 1_000,
|
||||
output_tokens: 0,
|
||||
},
|
||||
);
|
||||
let event = detect_cache_break(&PromptCacheConfig::default(), Some(&previous), ¤t)
|
||||
.expect("break should be detected");
|
||||
assert!(event.unexpected);
|
||||
assert!(event.reason.contains("stable"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn changed_prompt_marks_break_as_expected() {
|
||||
let previous_request = sample_request("first");
|
||||
let current_request = sample_request("second");
|
||||
let previous = TrackedPromptState::from_usage(
|
||||
&previous_request,
|
||||
&Usage {
|
||||
input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 6_000,
|
||||
output_tokens: 0,
|
||||
},
|
||||
);
|
||||
let current = TrackedPromptState::from_usage(
|
||||
¤t_request,
|
||||
&Usage {
|
||||
input_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 1_000,
|
||||
output_tokens: 0,
|
||||
},
|
||||
);
|
||||
let event = detect_cache_break(&PromptCacheConfig::default(), Some(&previous), ¤t)
|
||||
.expect("break should be detected");
|
||||
assert!(!event.unexpected);
|
||||
assert!(event.reason.contains("message payload changed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completion_cache_round_trip_persists_recent_response() {
|
||||
let _guard = test_env_lock();
|
||||
let temp_root = std::env::temp_dir().join(format!(
|
||||
"prompt-cache-test-{}-{}",
|
||||
std::process::id(),
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time")
|
||||
.as_nanos()
|
||||
));
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
|
||||
let cache = PromptCache::new("unit-test-session");
|
||||
let request = sample_request("cache me");
|
||||
let response = sample_response(42, 12, "cached");
|
||||
|
||||
assert!(cache.lookup_completion(&request).is_none());
|
||||
let record = cache.record_response(&request, &response);
|
||||
assert!(record.cache_break.is_none());
|
||||
|
||||
let cached = cache
|
||||
.lookup_completion(&request)
|
||||
.expect("cached response should load");
|
||||
assert_eq!(cached.content, response.content);
|
||||
|
||||
let stats = cache.stats();
|
||||
assert_eq!(stats.completion_cache_hits, 1);
|
||||
assert_eq!(stats.completion_cache_misses, 1);
|
||||
assert_eq!(stats.completion_cache_writes, 1);
|
||||
|
||||
let persisted = read_json::<super::PromptCacheStats>(&cache.paths().stats_path)
|
||||
.expect("stats should persist");
|
||||
assert_eq!(persisted.completion_cache_hits, 1);
|
||||
|
||||
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn distinct_requests_do_not_collide_in_completion_cache() {
|
||||
let _guard = test_env_lock();
|
||||
let temp_root = std::env::temp_dir().join(format!(
|
||||
"prompt-cache-distinct-{}-{}",
|
||||
std::process::id(),
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time")
|
||||
.as_nanos()
|
||||
));
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
|
||||
let cache = PromptCache::new("distinct-request-session");
|
||||
let first_request = sample_request("first");
|
||||
let second_request = sample_request("second");
|
||||
|
||||
let response = sample_response(42, 12, "cached");
|
||||
let _ = cache.record_response(&first_request, &response);
|
||||
|
||||
assert!(cache.lookup_completion(&second_request).is_none());
|
||||
|
||||
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expired_completion_entries_are_not_reused() {
|
||||
let _guard = test_env_lock();
|
||||
let temp_root = std::env::temp_dir().join(format!(
|
||||
"prompt-cache-expired-{}-{}",
|
||||
std::process::id(),
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time")
|
||||
.as_nanos()
|
||||
));
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
|
||||
let cache = PromptCache::with_config(PromptCacheConfig {
|
||||
session_id: "expired-session".to_string(),
|
||||
completion_ttl: Duration::ZERO,
|
||||
..PromptCacheConfig::default()
|
||||
});
|
||||
let request = sample_request("expire me");
|
||||
let response = sample_response(7, 3, "stale");
|
||||
|
||||
let _ = cache.record_response(&request, &response);
|
||||
|
||||
assert!(cache.lookup_completion(&request).is_none());
|
||||
let stats = cache.stats();
|
||||
assert_eq!(stats.completion_cache_hits, 0);
|
||||
assert_eq!(stats.completion_cache_misses, 1);
|
||||
|
||||
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitize_path_caps_long_values() {
|
||||
let long_value = "x".repeat(200);
|
||||
let sanitized = sanitize_path_segment(&long_value);
|
||||
assert!(sanitized.len() <= 80);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_hashes_are_versioned_and_stable() {
|
||||
let request = sample_request("stable");
|
||||
let first = request_hash_hex(&request);
|
||||
let second = request_hash_hex(&request);
|
||||
assert_eq!(first, second);
|
||||
assert!(first.starts_with(REQUEST_FINGERPRINT_PREFIX));
|
||||
}
|
||||
|
||||
fn sample_request(text: &str) -> MessageRequest {
|
||||
MessageRequest {
|
||||
model: "claude-3-7-sonnet-latest".to_string(),
|
||||
max_tokens: 64,
|
||||
messages: vec![InputMessage::user_text(text)],
|
||||
system: Some("system".to_string()),
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
stream: false,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_response(
|
||||
cache_read_input_tokens: u32,
|
||||
output_tokens: u32,
|
||||
text: &str,
|
||||
) -> MessageResponse {
|
||||
MessageResponse {
|
||||
id: "msg_test".to_string(),
|
||||
kind: "message".to_string(),
|
||||
role: "assistant".to_string(),
|
||||
content: vec![OutputContentBlock::Text {
|
||||
text: text.to_string(),
|
||||
}],
|
||||
model: "claude-3-7-sonnet-latest".to_string(),
|
||||
stop_reason: Some("end_turn".to_string()),
|
||||
stop_sequence: None,
|
||||
usage: Usage {
|
||||
input_tokens: 10,
|
||||
cache_creation_input_tokens: 5,
|
||||
cache_read_input_tokens,
|
||||
output_tokens,
|
||||
},
|
||||
request_id: Some("req_test".to_string()),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -4,6 +4,8 @@ use crate::types::StreamEvent;
|
||||
#[derive(Debug, Default)]
|
||||
pub struct SseParser {
|
||||
buffer: Vec<u8>,
|
||||
provider: Option<String>,
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
impl SseParser {
|
||||
@@ -12,12 +14,23 @@ impl SseParser {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Attach the provider name and model to this parser so that JSON
|
||||
/// deserialization failures within streamed frames carry enough context
|
||||
/// for callers to understand which upstream produced the unparseable
|
||||
/// payload.
|
||||
#[must_use]
|
||||
pub fn with_context(mut self, provider: impl Into<String>, model: impl Into<String>) -> Self {
|
||||
self.provider = Some(provider.into());
|
||||
self.model = Some(model.into());
|
||||
self
|
||||
}
|
||||
|
||||
pub fn push(&mut self, chunk: &[u8]) -> Result<Vec<StreamEvent>, ApiError> {
|
||||
self.buffer.extend_from_slice(chunk);
|
||||
let mut events = Vec::new();
|
||||
|
||||
while let Some(frame) = self.next_frame() {
|
||||
if let Some(event) = parse_frame(&frame)? {
|
||||
if let Some(event) = self.parse_frame_with_context(&frame)? {
|
||||
events.push(event);
|
||||
}
|
||||
}
|
||||
@@ -31,12 +44,18 @@ impl SseParser {
|
||||
}
|
||||
|
||||
let trailing = std::mem::take(&mut self.buffer);
|
||||
match parse_frame(&String::from_utf8_lossy(&trailing))? {
|
||||
match self.parse_frame_with_context(&String::from_utf8_lossy(&trailing))? {
|
||||
Some(event) => Ok(vec![event]),
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_frame_with_context(&self, frame: &str) -> Result<Option<StreamEvent>, ApiError> {
|
||||
let provider = self.provider.as_deref().unwrap_or("unknown");
|
||||
let model = self.model.as_deref().unwrap_or("unknown");
|
||||
parse_frame_with_provider(frame, provider, model)
|
||||
}
|
||||
|
||||
fn next_frame(&mut self) -> Option<String> {
|
||||
let separator = self
|
||||
.buffer
|
||||
@@ -61,6 +80,14 @@ impl SseParser {
|
||||
}
|
||||
|
||||
pub fn parse_frame(frame: &str) -> Result<Option<StreamEvent>, ApiError> {
|
||||
parse_frame_with_provider(frame, "unknown", "unknown")
|
||||
}
|
||||
|
||||
pub(crate) fn parse_frame_with_provider(
|
||||
frame: &str,
|
||||
provider: &str,
|
||||
model: &str,
|
||||
) -> Result<Option<StreamEvent>, ApiError> {
|
||||
let trimmed = frame.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Ok(None);
|
||||
@@ -97,7 +124,7 @@ pub fn parse_frame(frame: &str) -> Result<Option<StreamEvent>, ApiError> {
|
||||
|
||||
serde_json::from_str::<StreamEvent>(&payload)
|
||||
.map(Some)
|
||||
.map_err(ApiError::from)
|
||||
.map_err(|error| ApiError::json_deserialize(provider, model, &payload, error))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
@@ -216,4 +243,88 @@ mod tests {
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_thinking_content_block_start() {
|
||||
let frame = concat!(
|
||||
"event: content_block_start\n",
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\",\"signature\":null}}\n\n"
|
||||
);
|
||||
|
||||
let event = parse_frame(frame).expect("frame should parse");
|
||||
assert_eq!(
|
||||
event,
|
||||
Some(StreamEvent::ContentBlockStart(
|
||||
crate::types::ContentBlockStartEvent {
|
||||
index: 0,
|
||||
content_block: OutputContentBlock::Thinking {
|
||||
thinking: String::new(),
|
||||
signature: None,
|
||||
},
|
||||
},
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_thinking_related_deltas() {
|
||||
let thinking = concat!(
|
||||
"event: content_block_delta\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"step 1\"}}\n\n"
|
||||
);
|
||||
let signature = concat!(
|
||||
"event: content_block_delta\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"sig_123\"}}\n\n"
|
||||
);
|
||||
|
||||
let thinking_event = parse_frame(thinking).expect("thinking delta should parse");
|
||||
let signature_event = parse_frame(signature).expect("signature delta should parse");
|
||||
|
||||
assert_eq!(
|
||||
thinking_event,
|
||||
Some(StreamEvent::ContentBlockDelta(
|
||||
crate::types::ContentBlockDeltaEvent {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::ThinkingDelta {
|
||||
thinking: "step 1".to_string(),
|
||||
},
|
||||
}
|
||||
))
|
||||
);
|
||||
assert_eq!(
|
||||
signature_event,
|
||||
Some(StreamEvent::ContentBlockDelta(
|
||||
crate::types::ContentBlockDeltaEvent {
|
||||
index: 0,
|
||||
delta: ContentBlockDelta::SignatureDelta {
|
||||
signature: "sig_123".to_string(),
|
||||
},
|
||||
}
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_message_delta_frame_with_empty_usage_when_parsed_then_usage_defaults_to_zero() {
|
||||
// given
|
||||
let frame = concat!(
|
||||
"event: message_delta\n",
|
||||
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{}}\n\n"
|
||||
);
|
||||
|
||||
// when
|
||||
let event = parse_frame(frame).expect("frame should parse");
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
event,
|
||||
Some(StreamEvent::MessageDelta(crate::types::MessageDeltaEvent {
|
||||
delta: MessageDelta {
|
||||
stop_reason: Some("end_turn".to_string()),
|
||||
stop_sequence: None,
|
||||
},
|
||||
usage: Usage::default(),
|
||||
}))
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,8 @@
|
||||
use runtime::{pricing_for_model, TokenUsage, UsageCostEstimate};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
|
||||
pub struct MessageRequest {
|
||||
pub model: String,
|
||||
pub max_tokens: u32,
|
||||
@@ -12,10 +13,24 @@ pub struct MessageRequest {
|
||||
pub tools: Option<Vec<ToolDefinition>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tool_choice: Option<ToolChoice>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub thinking: Option<ThinkingConfig>,
|
||||
#[serde(default, skip_serializing_if = "std::ops::Not::not")]
|
||||
pub stream: bool,
|
||||
/// OpenAI-compatible tuning parameters. Optional — omitted from payload when None.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub temperature: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub top_p: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub frequency_penalty: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub presence_penalty: Option<f64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub stop: Option<Vec<String>>,
|
||||
/// Reasoning effort level for OpenAI-compatible reasoning models (e.g. `o4-mini`).
|
||||
/// Accepted values: `"low"`, `"medium"`, `"high"`. Omitted when `None`.
|
||||
/// Silently ignored by backends that do not support it.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub reasoning_effort: Option<String>,
|
||||
}
|
||||
|
||||
impl MessageRequest {
|
||||
@@ -26,23 +41,6 @@ impl MessageRequest {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ThinkingConfig {
|
||||
#[serde(rename = "type")]
|
||||
pub kind: String,
|
||||
pub budget_tokens: u32,
|
||||
}
|
||||
|
||||
impl ThinkingConfig {
|
||||
#[must_use]
|
||||
pub fn enabled(budget_tokens: u32) -> Self {
|
||||
Self {
|
||||
kind: "enabled".to_string(),
|
||||
budget_tokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct InputMessage {
|
||||
pub role: String,
|
||||
@@ -83,9 +81,6 @@ pub enum InputContentBlock {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
Image {
|
||||
source: ImageSource,
|
||||
},
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
@@ -99,14 +94,6 @@ pub enum InputContentBlock {
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ImageSource {
|
||||
#[serde(rename = "type")]
|
||||
pub kind: String,
|
||||
pub media_type: String,
|
||||
pub data: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ToolResultContentBlock {
|
||||
@@ -142,6 +129,7 @@ pub struct MessageResponse {
|
||||
pub stop_reason: Option<String>,
|
||||
#[serde(default)]
|
||||
pub stop_sequence: Option<String>,
|
||||
#[serde(default)]
|
||||
pub usage: Usage,
|
||||
#[serde(default)]
|
||||
pub request_id: Option<String>,
|
||||
@@ -160,32 +148,60 @@ pub enum OutputContentBlock {
|
||||
Text {
|
||||
text: String,
|
||||
},
|
||||
Thinking {
|
||||
thinking: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
signature: Option<String>,
|
||||
},
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
input: Value,
|
||||
},
|
||||
Thinking {
|
||||
#[serde(default)]
|
||||
thinking: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
signature: Option<String>,
|
||||
},
|
||||
RedactedThinking {
|
||||
data: Value,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct Usage {
|
||||
#[serde(default)]
|
||||
pub input_tokens: u32,
|
||||
#[serde(default)]
|
||||
pub cache_creation_input_tokens: u32,
|
||||
#[serde(default)]
|
||||
pub cache_read_input_tokens: u32,
|
||||
#[serde(default)]
|
||||
pub output_tokens: u32,
|
||||
}
|
||||
|
||||
impl Usage {
|
||||
#[must_use]
|
||||
pub const fn total_tokens(&self) -> u32 {
|
||||
self.input_tokens + self.output_tokens
|
||||
self.input_tokens
|
||||
+ self.output_tokens
|
||||
+ self.cache_creation_input_tokens
|
||||
+ self.cache_read_input_tokens
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub const fn token_usage(&self) -> TokenUsage {
|
||||
TokenUsage {
|
||||
input_tokens: self.input_tokens,
|
||||
output_tokens: self.output_tokens,
|
||||
cache_creation_input_tokens: self.cache_creation_input_tokens,
|
||||
cache_read_input_tokens: self.cache_read_input_tokens,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn estimated_cost_usd(&self, model: &str) -> UsageCostEstimate {
|
||||
let usage = self.token_usage();
|
||||
pricing_for_model(model).map_or_else(
|
||||
|| usage.estimate_cost_usd(),
|
||||
|pricing| usage.estimate_cost_usd_with_pricing(pricing),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -197,6 +213,7 @@ pub struct MessageStartEvent {
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct MessageDeltaEvent {
|
||||
pub delta: MessageDelta,
|
||||
#[serde(default)]
|
||||
pub usage: Usage,
|
||||
}
|
||||
|
||||
@@ -224,9 +241,9 @@ pub struct ContentBlockDeltaEvent {
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ContentBlockDelta {
|
||||
TextDelta { text: String },
|
||||
InputJsonDelta { partial_json: String },
|
||||
ThinkingDelta { thinking: String },
|
||||
SignatureDelta { signature: String },
|
||||
InputJsonDelta { partial_json: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
@@ -247,3 +264,47 @@ pub enum StreamEvent {
|
||||
ContentBlockStop(ContentBlockStopEvent),
|
||||
MessageStop(MessageStopEvent),
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use runtime::format_usd;
|
||||
|
||||
use super::{MessageResponse, Usage};
|
||||
|
||||
#[test]
|
||||
fn usage_total_tokens_includes_cache_tokens() {
|
||||
let usage = Usage {
|
||||
input_tokens: 10,
|
||||
cache_creation_input_tokens: 2,
|
||||
cache_read_input_tokens: 3,
|
||||
output_tokens: 4,
|
||||
};
|
||||
|
||||
assert_eq!(usage.total_tokens(), 19);
|
||||
assert_eq!(usage.token_usage().total_tokens(), 19);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn message_response_estimates_cost_from_model_usage() {
|
||||
let response = MessageResponse {
|
||||
id: "msg_cost".to_string(),
|
||||
kind: "message".to_string(),
|
||||
role: "assistant".to_string(),
|
||||
content: Vec::new(),
|
||||
model: "claude-sonnet-4-20250514".to_string(),
|
||||
stop_reason: Some("end_turn".to_string()),
|
||||
stop_sequence: None,
|
||||
usage: Usage {
|
||||
input_tokens: 1_000_000,
|
||||
cache_creation_input_tokens: 100_000,
|
||||
cache_read_input_tokens: 200_000,
|
||||
output_tokens: 500_000,
|
||||
},
|
||||
request_id: None,
|
||||
};
|
||||
|
||||
let cost = response.usage.estimated_cost_usd(&response.model);
|
||||
assert_eq!(format_usd(cost.total_cost_usd()), "$54.6750");
|
||||
assert_eq!(response.total_tokens(), 1_800_000);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,928 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Mutex as StdMutex, OnceLock};
|
||||
use std::time::Duration;
|
||||
|
||||
use api::{
|
||||
AnthropicClient, ApiClient, ApiError, AuthSource, ContentBlockDelta, ContentBlockDeltaEvent,
|
||||
ContentBlockStartEvent, InputContentBlock, InputMessage, MessageDeltaEvent, MessageRequest,
|
||||
OutputContentBlock, PromptCache, PromptCacheConfig, ProviderClient, StreamEvent, ToolChoice,
|
||||
ToolDefinition,
|
||||
};
|
||||
use serde_json::json;
|
||||
use telemetry::{ClientIdentity, MemoryTelemetrySink, SessionTracer, TelemetryEvent};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
static LOCK: OnceLock<StdMutex<()>> = OnceLock::new();
|
||||
LOCK.get_or_init(|| StdMutex::new(()))
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_message_posts_json_and_parses_response() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let body = concat!(
|
||||
"{",
|
||||
"\"id\":\"msg_test\",",
|
||||
"\"type\":\"message\",",
|
||||
"\"role\":\"assistant\",",
|
||||
"\"content\":[{\"type\":\"text\",\"text\":\"Hello from Claude\"}],",
|
||||
"\"model\":\"claude-3-7-sonnet-latest\",",
|
||||
"\"stop_reason\":\"end_turn\",",
|
||||
"\"stop_sequence\":null,",
|
||||
"\"usage\":{\"input_tokens\":12,\"output_tokens\":4},",
|
||||
"\"request_id\":\"req_body_123\"",
|
||||
"}"
|
||||
);
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response("200 OK", "application/json", body)],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = ApiClient::new("test-key")
|
||||
.with_auth_token(Some("proxy-token".to_string()))
|
||||
.with_base_url(server.base_url());
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.id, "msg_test");
|
||||
assert_eq!(response.total_tokens(), 16);
|
||||
assert_eq!(response.request_id.as_deref(), Some("req_body_123"));
|
||||
assert_eq!(response.usage.cache_creation_input_tokens, 0);
|
||||
assert_eq!(response.usage.cache_read_input_tokens, 0);
|
||||
assert_eq!(
|
||||
response.content,
|
||||
vec![OutputContentBlock::Text {
|
||||
text: "Hello from Claude".to_string(),
|
||||
}]
|
||||
);
|
||||
|
||||
let captured = state.lock().await;
|
||||
let request = captured.first().expect("server should capture request");
|
||||
assert_eq!(request.method, "POST");
|
||||
assert_eq!(request.path, "/v1/messages");
|
||||
assert_eq!(
|
||||
request.headers.get("x-api-key").map(String::as_str),
|
||||
Some("test-key")
|
||||
);
|
||||
assert_eq!(
|
||||
request.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer proxy-token")
|
||||
);
|
||||
assert_eq!(
|
||||
request.headers.get("anthropic-version").map(String::as_str),
|
||||
Some("2023-06-01")
|
||||
);
|
||||
assert_eq!(
|
||||
request.headers.get("user-agent").map(String::as_str),
|
||||
Some("claude-code/0.1.0")
|
||||
);
|
||||
assert_eq!(
|
||||
request.headers.get("anthropic-beta").map(String::as_str),
|
||||
Some("claude-code-20250219,prompt-caching-scope-2026-01-05")
|
||||
);
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_str(&request.body).expect("request body should be json");
|
||||
assert_eq!(
|
||||
body.get("model").and_then(serde_json::Value::as_str),
|
||||
Some("claude-3-7-sonnet-latest")
|
||||
);
|
||||
assert!(body.get("stream").is_none());
|
||||
assert_eq!(body["tools"][0]["name"], json!("get_weather"));
|
||||
assert_eq!(body["tool_choice"]["type"], json!("auto"));
|
||||
assert!(
|
||||
body.get("betas").is_none(),
|
||||
"betas must travel via the anthropic-beta header, not the request body"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_message_blocks_oversized_requests_before_the_http_call() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response("200 OK", "application/json", "{}")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = AnthropicClient::new("test-key").with_base_url(server.base_url());
|
||||
let error = client
|
||||
.send_message(&MessageRequest {
|
||||
model: "claude-sonnet-4-6".to_string(),
|
||||
max_tokens: 64_000,
|
||||
messages: vec![InputMessage {
|
||||
role: "user".to_string(),
|
||||
content: vec![InputContentBlock::Text {
|
||||
text: "x".repeat(600_000),
|
||||
}],
|
||||
}],
|
||||
system: Some("Keep the answer short.".to_string()),
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
stream: false,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.expect_err("oversized request should fail local context-window preflight");
|
||||
|
||||
assert!(matches!(error, ApiError::ContextWindowExceeded { .. }));
|
||||
assert!(
|
||||
state.lock().await.is_empty(),
|
||||
"preflight failure should avoid any upstream HTTP request"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_message_applies_request_profile_and_records_telemetry() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response_with_headers(
|
||||
"200 OK",
|
||||
"application/json",
|
||||
concat!(
|
||||
"{",
|
||||
"\"id\":\"msg_profile\",",
|
||||
"\"type\":\"message\",",
|
||||
"\"role\":\"assistant\",",
|
||||
"\"content\":[{\"type\":\"text\",\"text\":\"ok\"}],",
|
||||
"\"model\":\"claude-3-7-sonnet-latest\",",
|
||||
"\"stop_reason\":\"end_turn\",",
|
||||
"\"stop_sequence\":null,",
|
||||
"\"usage\":{\"input_tokens\":1,\"cache_creation_input_tokens\":2,\"cache_read_input_tokens\":3,\"output_tokens\":1}",
|
||||
"}"
|
||||
),
|
||||
&[("request-id", "req_profile_123")],
|
||||
)],
|
||||
)
|
||||
.await;
|
||||
let sink = Arc::new(MemoryTelemetrySink::default());
|
||||
|
||||
let client = AnthropicClient::new("test-key")
|
||||
.with_base_url(server.base_url())
|
||||
.with_client_identity(ClientIdentity::new("claude-code", "9.9.9").with_runtime("rust-cli"))
|
||||
.with_beta("tools-2026-04-01")
|
||||
.with_extra_body_param("metadata", json!({"source": "clawd-code"}))
|
||||
.with_session_tracer(SessionTracer::new("session-telemetry", sink.clone()));
|
||||
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.request_id.as_deref(), Some("req_profile_123"));
|
||||
|
||||
let captured = state.lock().await;
|
||||
let request = captured.first().expect("server should capture request");
|
||||
assert_eq!(
|
||||
request.headers.get("anthropic-beta").map(String::as_str),
|
||||
Some("claude-code-20250219,prompt-caching-scope-2026-01-05,tools-2026-04-01")
|
||||
);
|
||||
assert_eq!(
|
||||
request.headers.get("user-agent").map(String::as_str),
|
||||
Some("claude-code/9.9.9")
|
||||
);
|
||||
let body: serde_json::Value =
|
||||
serde_json::from_str(&request.body).expect("request body should be json");
|
||||
assert_eq!(body["metadata"]["source"], json!("clawd-code"));
|
||||
assert!(
|
||||
body.get("betas").is_none(),
|
||||
"betas must travel via the anthropic-beta header, not the request body"
|
||||
);
|
||||
|
||||
let events = sink.events();
|
||||
assert_eq!(events.len(), 6);
|
||||
assert!(matches!(
|
||||
&events[0],
|
||||
TelemetryEvent::HttpRequestStarted {
|
||||
session_id,
|
||||
attempt: 1,
|
||||
method,
|
||||
path,
|
||||
..
|
||||
} if session_id == "session-telemetry" && method == "POST" && path == "/v1/messages"
|
||||
));
|
||||
assert!(matches!(
|
||||
&events[1],
|
||||
TelemetryEvent::SessionTrace(trace) if trace.name == "http_request_started"
|
||||
));
|
||||
assert!(matches!(
|
||||
&events[2],
|
||||
TelemetryEvent::HttpRequestSucceeded {
|
||||
request_id,
|
||||
status: 200,
|
||||
..
|
||||
} if request_id.as_deref() == Some("req_profile_123")
|
||||
));
|
||||
assert!(matches!(
|
||||
&events[3],
|
||||
TelemetryEvent::SessionTrace(trace) if trace.name == "http_request_succeeded"
|
||||
));
|
||||
assert!(matches!(
|
||||
&events[4],
|
||||
TelemetryEvent::Analytics(event)
|
||||
if event.namespace == "api"
|
||||
&& event.action == "message_usage"
|
||||
&& event.properties.get("request_id") == Some(&json!("req_profile_123"))
|
||||
&& event.properties.get("total_tokens") == Some(&json!(7))
|
||||
&& event.properties.get("estimated_cost_usd") == Some(&json!("$0.0001"))
|
||||
));
|
||||
assert!(matches!(
|
||||
&events[5],
|
||||
TelemetryEvent::SessionTrace(trace) if trace.name == "analytics"
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_message_parses_prompt_cache_token_usage_from_response() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let body = concat!(
|
||||
"{",
|
||||
"\"id\":\"msg_cache_tokens\",",
|
||||
"\"type\":\"message\",",
|
||||
"\"role\":\"assistant\",",
|
||||
"\"content\":[{\"type\":\"text\",\"text\":\"Cache tokens\"}],",
|
||||
"\"model\":\"claude-3-7-sonnet-latest\",",
|
||||
"\"stop_reason\":\"end_turn\",",
|
||||
"\"stop_sequence\":null,",
|
||||
"\"usage\":{\"input_tokens\":12,\"cache_creation_input_tokens\":321,\"cache_read_input_tokens\":654,\"output_tokens\":4}",
|
||||
"}"
|
||||
);
|
||||
let server = spawn_server(
|
||||
state,
|
||||
vec![http_response("200 OK", "application/json", body)],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = AnthropicClient::new("test-key").with_base_url(server.base_url());
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.usage.input_tokens, 12);
|
||||
assert_eq!(response.usage.cache_creation_input_tokens, 321);
|
||||
assert_eq!(response.usage.cache_read_input_tokens, 654);
|
||||
assert_eq!(response.usage.output_tokens, 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn given_empty_usage_object_when_send_message_parses_response_then_usage_defaults_to_zero() {
|
||||
// given
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let body = concat!(
|
||||
"{",
|
||||
"\"id\":\"msg_empty_usage\",",
|
||||
"\"type\":\"message\",",
|
||||
"\"role\":\"assistant\",",
|
||||
"\"content\":[{\"type\":\"text\",\"text\":\"Hello from Claude\"}],",
|
||||
"\"model\":\"claude-3-7-sonnet-latest\",",
|
||||
"\"stop_reason\":\"end_turn\",",
|
||||
"\"stop_sequence\":null,",
|
||||
"\"usage\":{}",
|
||||
"}"
|
||||
);
|
||||
let server = spawn_server(
|
||||
state,
|
||||
vec![http_response("200 OK", "application/json", body)],
|
||||
)
|
||||
.await;
|
||||
let client = AnthropicClient::new("test-key").with_base_url(server.base_url());
|
||||
|
||||
// when
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("response with empty usage object should still parse");
|
||||
|
||||
// then
|
||||
assert_eq!(response.id, "msg_empty_usage");
|
||||
assert_eq!(response.total_tokens(), 0);
|
||||
assert_eq!(response.usage.input_tokens, 0);
|
||||
assert_eq!(response.usage.cache_creation_input_tokens, 0);
|
||||
assert_eq!(response.usage.cache_read_input_tokens, 0);
|
||||
assert_eq!(response.usage.output_tokens, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn stream_message_parses_sse_events_with_tool_use() {
|
||||
let _guard = env_lock();
|
||||
let temp_root = std::env::temp_dir().join(format!(
|
||||
"api-stream-cache-{}-{}",
|
||||
std::process::id(),
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("time")
|
||||
.as_nanos()
|
||||
));
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let sse = concat!(
|
||||
"event: message_start\n",
|
||||
"data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_stream\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":8,\"cache_creation_input_tokens\":13,\"cache_read_input_tokens\":21,\"output_tokens\":0}}}\n\n",
|
||||
"event: content_block_start\n",
|
||||
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"toolu_123\",\"name\":\"get_weather\",\"input\":{}}}\n\n",
|
||||
"event: content_block_delta\n",
|
||||
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"city\\\":\\\"Paris\\\"}\"}}\n\n",
|
||||
"event: content_block_stop\n",
|
||||
"data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
|
||||
"event: message_delta\n",
|
||||
"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\",\"stop_sequence\":null},\"usage\":{\"input_tokens\":8,\"cache_creation_input_tokens\":34,\"cache_read_input_tokens\":55,\"output_tokens\":1}}\n\n",
|
||||
"event: message_stop\n",
|
||||
"data: {\"type\":\"message_stop\"}\n\n",
|
||||
"data: [DONE]\n\n"
|
||||
);
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response_with_headers(
|
||||
"200 OK",
|
||||
"text/event-stream",
|
||||
sse,
|
||||
&[("request-id", "req_stream_456")],
|
||||
)],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = ApiClient::new("test-key")
|
||||
.with_auth_token(Some("proxy-token".to_string()))
|
||||
.with_base_url(server.base_url())
|
||||
.with_prompt_cache(PromptCache::new("stream-session"));
|
||||
let mut stream = client
|
||||
.stream_message(&sample_request(false))
|
||||
.await
|
||||
.expect("stream should start");
|
||||
|
||||
assert_eq!(stream.request_id(), Some("req_stream_456"));
|
||||
|
||||
let mut events = Vec::new();
|
||||
while let Some(event) = stream
|
||||
.next_event()
|
||||
.await
|
||||
.expect("stream event should parse")
|
||||
{
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
assert_eq!(events.len(), 6);
|
||||
assert!(matches!(events[0], StreamEvent::MessageStart(_)));
|
||||
assert!(matches!(
|
||||
events[1],
|
||||
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
|
||||
content_block: OutputContentBlock::ToolUse { .. },
|
||||
..
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
events[2],
|
||||
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
|
||||
delta: ContentBlockDelta::InputJsonDelta { .. },
|
||||
..
|
||||
})
|
||||
));
|
||||
assert!(matches!(events[3], StreamEvent::ContentBlockStop(_)));
|
||||
assert!(matches!(
|
||||
events[4],
|
||||
StreamEvent::MessageDelta(MessageDeltaEvent { .. })
|
||||
));
|
||||
assert!(matches!(events[5], StreamEvent::MessageStop(_)));
|
||||
|
||||
match &events[1] {
|
||||
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
|
||||
content_block: OutputContentBlock::ToolUse { name, input, .. },
|
||||
..
|
||||
}) => {
|
||||
assert_eq!(name, "get_weather");
|
||||
assert_eq!(input, &json!({}));
|
||||
}
|
||||
other => panic!("expected tool_use block, got {other:?}"),
|
||||
}
|
||||
|
||||
let captured = state.lock().await;
|
||||
let request = captured.first().expect("server should capture request");
|
||||
assert!(request.body.contains("\"stream\":true"));
|
||||
|
||||
let cache_stats = client
|
||||
.prompt_cache_stats()
|
||||
.expect("prompt cache stats should exist");
|
||||
assert_eq!(cache_stats.tracked_requests, 1);
|
||||
assert_eq!(cache_stats.last_cache_creation_input_tokens, Some(34));
|
||||
assert_eq!(cache_stats.last_cache_read_input_tokens, Some(55));
|
||||
assert_eq!(
|
||||
cache_stats.last_cache_source.as_deref(),
|
||||
Some("api-response")
|
||||
);
|
||||
|
||||
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn retries_retryable_failures_before_succeeding() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![
|
||||
http_response(
|
||||
"429 Too Many Requests",
|
||||
"application/json",
|
||||
"{\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down\"}}",
|
||||
),
|
||||
http_response(
|
||||
"200 OK",
|
||||
"application/json",
|
||||
"{\"id\":\"msg_retry\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Recovered\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}",
|
||||
),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = ApiClient::new("test-key")
|
||||
.with_base_url(server.base_url())
|
||||
.with_retry_policy(2, Duration::from_millis(1), Duration::from_millis(2));
|
||||
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("retry should eventually succeed");
|
||||
|
||||
assert_eq!(response.total_tokens(), 5);
|
||||
assert_eq!(state.lock().await.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_client_dispatches_anthropic_requests() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response(
|
||||
"200 OK",
|
||||
"application/json",
|
||||
"{\"id\":\"msg_provider\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Dispatched\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}",
|
||||
)],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = ProviderClient::from_model_with_anthropic_auth(
|
||||
"claude-sonnet-4-6",
|
||||
Some(AuthSource::ApiKey("test-key".to_string())),
|
||||
)
|
||||
.expect("anthropic provider client should be constructed");
|
||||
let client = match client {
|
||||
ProviderClient::Anthropic(client) => {
|
||||
ProviderClient::Anthropic(client.with_base_url(server.base_url()))
|
||||
}
|
||||
other => panic!("expected anthropic provider, got {other:?}"),
|
||||
};
|
||||
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("provider-dispatched request should succeed");
|
||||
|
||||
assert_eq!(response.total_tokens(), 5);
|
||||
|
||||
let captured = state.lock().await;
|
||||
let request = captured.first().expect("server should capture request");
|
||||
assert_eq!(request.path, "/v1/messages");
|
||||
assert_eq!(
|
||||
request.headers.get("x-api-key").map(String::as_str),
|
||||
Some("test-key")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn surfaces_retry_exhaustion_for_persistent_retryable_errors() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![
|
||||
http_response(
|
||||
"503 Service Unavailable",
|
||||
"application/json",
|
||||
"{\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"busy\"}}",
|
||||
),
|
||||
http_response(
|
||||
"503 Service Unavailable",
|
||||
"application/json",
|
||||
"{\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"still busy\"}}",
|
||||
),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = ApiClient::new("test-key")
|
||||
.with_base_url(server.base_url())
|
||||
.with_retry_policy(1, Duration::from_millis(1), Duration::from_millis(2));
|
||||
|
||||
let error = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect_err("persistent 503 should fail");
|
||||
|
||||
match error {
|
||||
ApiError::RetriesExhausted {
|
||||
attempts,
|
||||
last_error,
|
||||
} => {
|
||||
assert_eq!(attempts, 2);
|
||||
assert!(matches!(
|
||||
*last_error,
|
||||
ApiError::Api {
|
||||
status: reqwest::StatusCode::SERVICE_UNAVAILABLE,
|
||||
retryable: true,
|
||||
..
|
||||
}
|
||||
));
|
||||
}
|
||||
other => panic!("expected retries exhausted, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn retries_multiple_retryable_failures_with_exponential_backoff_and_jitter() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![
|
||||
http_response(
|
||||
"429 Too Many Requests",
|
||||
"application/json",
|
||||
"{\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down\"}}",
|
||||
),
|
||||
http_response(
|
||||
"500 Internal Server Error",
|
||||
"application/json",
|
||||
"{\"type\":\"error\",\"error\":{\"type\":\"api_error\",\"message\":\"boom\"}}",
|
||||
),
|
||||
http_response(
|
||||
"503 Service Unavailable",
|
||||
"application/json",
|
||||
"{\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"busy\"}}",
|
||||
),
|
||||
http_response(
|
||||
"429 Too Many Requests",
|
||||
"application/json",
|
||||
"{\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"message\":\"slow down again\"}}",
|
||||
),
|
||||
http_response(
|
||||
"503 Service Unavailable",
|
||||
"application/json",
|
||||
"{\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"still busy\"}}",
|
||||
),
|
||||
http_response(
|
||||
"200 OK",
|
||||
"application/json",
|
||||
"{\"id\":\"msg_exp_retry\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Recovered after 5\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}",
|
||||
),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = ApiClient::new("test-key")
|
||||
.with_base_url(server.base_url())
|
||||
.with_retry_policy(8, Duration::from_millis(1), Duration::from_millis(4));
|
||||
let started_at = std::time::Instant::now();
|
||||
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("8-retry policy should absorb 5 retryable failures");
|
||||
|
||||
let elapsed = started_at.elapsed();
|
||||
assert_eq!(response.total_tokens(), 5);
|
||||
assert_eq!(
|
||||
state.lock().await.len(),
|
||||
6,
|
||||
"client should issue 1 original + 5 retry requests before the 200"
|
||||
);
|
||||
// Jittered sleeps are bounded by 2 * max_backoff per retry (base + jitter),
|
||||
// so 5 sleeps fit comfortably below this upper bound with generous slack.
|
||||
assert!(
|
||||
elapsed < Duration::from_secs(5),
|
||||
"retries should complete promptly, took {elapsed:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn send_message_reuses_recent_completion_cache_entries() {
|
||||
let _guard = env_lock();
|
||||
let temp_root = std::env::temp_dir().join(format!(
|
||||
"api-prompt-cache-{}-{}",
|
||||
std::process::id(),
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("time")
|
||||
.as_nanos()
|
||||
));
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
|
||||
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response(
|
||||
"200 OK",
|
||||
"application/json",
|
||||
"{\"id\":\"msg_cached\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Cached once\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"cache_creation_input_tokens\":5,\"cache_read_input_tokens\":4000,\"output_tokens\":2}}",
|
||||
)],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = AnthropicClient::new("test-key")
|
||||
.with_base_url(server.base_url())
|
||||
.with_prompt_cache(PromptCache::new("integration-session"));
|
||||
|
||||
let first = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("first request should succeed");
|
||||
let second = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("second request should reuse cache");
|
||||
|
||||
assert_eq!(first.content, second.content);
|
||||
assert_eq!(state.lock().await.len(), 1);
|
||||
|
||||
let cache_stats = client
|
||||
.prompt_cache_stats()
|
||||
.expect("prompt cache stats should exist");
|
||||
assert_eq!(cache_stats.completion_cache_hits, 1);
|
||||
assert_eq!(cache_stats.completion_cache_misses, 1);
|
||||
assert_eq!(cache_stats.completion_cache_writes, 1);
|
||||
|
||||
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn send_message_tracks_unexpected_prompt_cache_breaks() {
|
||||
let _guard = env_lock();
|
||||
let temp_root = std::env::temp_dir().join(format!(
|
||||
"api-prompt-break-{}-{}",
|
||||
std::process::id(),
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("time")
|
||||
.as_nanos()
|
||||
));
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", &temp_root);
|
||||
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state,
|
||||
vec![
|
||||
http_response(
|
||||
"200 OK",
|
||||
"application/json",
|
||||
"{\"id\":\"msg_one\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"One\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"cache_creation_input_tokens\":5,\"cache_read_input_tokens\":6000,\"output_tokens\":2}}",
|
||||
),
|
||||
http_response(
|
||||
"200 OK",
|
||||
"application/json",
|
||||
"{\"id\":\"msg_two\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[{\"type\":\"text\",\"text\":\"Two\"}],\"model\":\"claude-3-7-sonnet-latest\",\"stop_reason\":\"end_turn\",\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"cache_creation_input_tokens\":0,\"cache_read_input_tokens\":1000,\"output_tokens\":2}}",
|
||||
),
|
||||
],
|
||||
)
|
||||
.await;
|
||||
|
||||
let request = sample_request(false);
|
||||
let client = AnthropicClient::new("test-key")
|
||||
.with_base_url(server.base_url())
|
||||
.with_prompt_cache(PromptCache::with_config(PromptCacheConfig {
|
||||
session_id: "break-session".to_string(),
|
||||
completion_ttl: Duration::from_secs(0),
|
||||
..PromptCacheConfig::default()
|
||||
}));
|
||||
|
||||
client
|
||||
.send_message(&request)
|
||||
.await
|
||||
.expect("first response should succeed");
|
||||
client
|
||||
.send_message(&request)
|
||||
.await
|
||||
.expect("second response should succeed");
|
||||
|
||||
let cache_stats = client
|
||||
.prompt_cache_stats()
|
||||
.expect("prompt cache stats should exist");
|
||||
assert_eq!(cache_stats.unexpected_cache_breaks, 1);
|
||||
assert_eq!(
|
||||
cache_stats.last_break_reason.as_deref(),
|
||||
Some("cache read tokens dropped while prompt fingerprint remained stable")
|
||||
);
|
||||
|
||||
std::fs::remove_dir_all(temp_root).expect("cleanup temp root");
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[ignore = "requires ANTHROPIC_API_KEY and network access"]
|
||||
async fn live_stream_smoke_test() {
|
||||
let client = ApiClient::from_env().expect("ANTHROPIC_API_KEY must be set");
|
||||
let mut stream = client
|
||||
.stream_message(&MessageRequest {
|
||||
model: std::env::var("ANTHROPIC_MODEL")
|
||||
.unwrap_or_else(|_| "claude-3-7-sonnet-latest".to_string()),
|
||||
max_tokens: 32,
|
||||
messages: vec![InputMessage::user_text(
|
||||
"Reply with exactly: hello from rust",
|
||||
)],
|
||||
system: None,
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
stream: false,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.expect("live stream should start");
|
||||
|
||||
while let Some(_event) = stream
|
||||
.next_event()
|
||||
.await
|
||||
.expect("live stream should yield events")
|
||||
{}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct CapturedRequest {
|
||||
method: String,
|
||||
path: String,
|
||||
headers: HashMap<String, String>,
|
||||
body: String,
|
||||
}
|
||||
|
||||
struct TestServer {
|
||||
base_url: String,
|
||||
join_handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl TestServer {
|
||||
fn base_url(&self) -> String {
|
||||
self.base_url.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TestServer {
|
||||
fn drop(&mut self) {
|
||||
self.join_handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
async fn spawn_server(
|
||||
state: Arc<Mutex<Vec<CapturedRequest>>>,
|
||||
responses: Vec<String>,
|
||||
) -> TestServer {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let address = listener
|
||||
.local_addr()
|
||||
.expect("listener should have local addr");
|
||||
let join_handle = tokio::spawn(async move {
|
||||
for response in responses {
|
||||
let (mut socket, _) = listener.accept().await.expect("server should accept");
|
||||
let mut buffer = Vec::new();
|
||||
let mut header_end = None;
|
||||
|
||||
loop {
|
||||
let mut chunk = [0_u8; 1024];
|
||||
let read = socket
|
||||
.read(&mut chunk)
|
||||
.await
|
||||
.expect("request read should succeed");
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
buffer.extend_from_slice(&chunk[..read]);
|
||||
if let Some(position) = find_header_end(&buffer) {
|
||||
header_end = Some(position);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let header_end = header_end.expect("request should include headers");
|
||||
let (header_bytes, remaining) = buffer.split_at(header_end);
|
||||
let header_text =
|
||||
String::from_utf8(header_bytes.to_vec()).expect("headers should be utf8");
|
||||
let mut lines = header_text.split("\r\n");
|
||||
let request_line = lines.next().expect("request line should exist");
|
||||
let mut parts = request_line.split_whitespace();
|
||||
let method = parts.next().expect("method should exist").to_string();
|
||||
let path = parts.next().expect("path should exist").to_string();
|
||||
let mut headers = HashMap::new();
|
||||
let mut content_length = 0_usize;
|
||||
for line in lines {
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let (name, value) = line.split_once(':').expect("header should have colon");
|
||||
let value = value.trim().to_string();
|
||||
if name.eq_ignore_ascii_case("content-length") {
|
||||
content_length = value.parse().expect("content length should parse");
|
||||
}
|
||||
headers.insert(name.to_ascii_lowercase(), value);
|
||||
}
|
||||
|
||||
let mut body = remaining[4..].to_vec();
|
||||
while body.len() < content_length {
|
||||
let mut chunk = vec![0_u8; content_length - body.len()];
|
||||
let read = socket
|
||||
.read(&mut chunk)
|
||||
.await
|
||||
.expect("body read should succeed");
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
body.extend_from_slice(&chunk[..read]);
|
||||
}
|
||||
|
||||
state.lock().await.push(CapturedRequest {
|
||||
method,
|
||||
path,
|
||||
headers,
|
||||
body: String::from_utf8(body).expect("body should be utf8"),
|
||||
});
|
||||
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("response write should succeed");
|
||||
}
|
||||
});
|
||||
|
||||
TestServer {
|
||||
base_url: format!("http://{address}"),
|
||||
join_handle,
|
||||
}
|
||||
}
|
||||
|
||||
fn find_header_end(bytes: &[u8]) -> Option<usize> {
|
||||
bytes.windows(4).position(|window| window == b"\r\n\r\n")
|
||||
}
|
||||
|
||||
fn http_response(status: &str, content_type: &str, body: &str) -> String {
|
||||
http_response_with_headers(status, content_type, body, &[])
|
||||
}
|
||||
|
||||
fn http_response_with_headers(
|
||||
status: &str,
|
||||
content_type: &str,
|
||||
body: &str,
|
||||
headers: &[(&str, &str)],
|
||||
) -> String {
|
||||
let mut extra_headers = String::new();
|
||||
for (name, value) in headers {
|
||||
use std::fmt::Write as _;
|
||||
write!(&mut extra_headers, "{name}: {value}\r\n").expect("header write should succeed");
|
||||
}
|
||||
format!(
|
||||
"HTTP/1.1 {status}\r\ncontent-type: {content_type}\r\n{extra_headers}content-length: {}\r\nconnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_request(stream: bool) -> MessageRequest {
|
||||
MessageRequest {
|
||||
model: "claude-3-7-sonnet-latest".to_string(),
|
||||
max_tokens: 64,
|
||||
messages: vec![InputMessage {
|
||||
role: "user".to_string(),
|
||||
content: vec![
|
||||
InputContentBlock::Text {
|
||||
text: "Say hello".to_string(),
|
||||
},
|
||||
InputContentBlock::ToolResult {
|
||||
tool_use_id: "toolu_prev".to_string(),
|
||||
content: vec![api::ToolResultContentBlock::Json {
|
||||
value: json!({"forecast": "sunny"}),
|
||||
}],
|
||||
is_error: false,
|
||||
},
|
||||
],
|
||||
}],
|
||||
system: Some("Use tools when needed".to_string()),
|
||||
tools: Some(vec![ToolDefinition {
|
||||
name: "get_weather".to_string(),
|
||||
description: Some("Fetches the weather".to_string()),
|
||||
input_schema: json!({
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"]
|
||||
}),
|
||||
}]),
|
||||
tool_choice: Some(ToolChoice::Auto),
|
||||
stream,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,531 @@
|
||||
use std::collections::HashMap;
|
||||
use std::ffi::OsString;
|
||||
use std::sync::Arc;
|
||||
use std::sync::{Mutex as StdMutex, OnceLock};
|
||||
|
||||
use api::{
|
||||
ApiError, ContentBlockDelta, ContentBlockDeltaEvent, ContentBlockStartEvent,
|
||||
ContentBlockStopEvent, InputContentBlock, InputMessage, MessageDeltaEvent, MessageRequest,
|
||||
OpenAiCompatClient, OpenAiCompatConfig, OutputContentBlock, ProviderClient, StreamEvent,
|
||||
ToolChoice, ToolDefinition,
|
||||
};
|
||||
use serde_json::json;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_message_uses_openai_compatible_endpoint_and_auth() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let body = concat!(
|
||||
"{",
|
||||
"\"id\":\"chatcmpl_test\",",
|
||||
"\"model\":\"grok-3\",",
|
||||
"\"choices\":[{",
|
||||
"\"message\":{\"role\":\"assistant\",\"content\":\"Hello from Grok\",\"tool_calls\":[]},",
|
||||
"\"finish_reason\":\"stop\"",
|
||||
"}],",
|
||||
"\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":5}",
|
||||
"}"
|
||||
);
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response("200 OK", "application/json", body)],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = OpenAiCompatClient::new("xai-test-key", OpenAiCompatConfig::xai())
|
||||
.with_base_url(server.base_url());
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.model, "grok-3");
|
||||
assert_eq!(response.total_tokens(), 16);
|
||||
assert_eq!(
|
||||
response.content,
|
||||
vec![OutputContentBlock::Text {
|
||||
text: "Hello from Grok".to_string(),
|
||||
}]
|
||||
);
|
||||
|
||||
let captured = state.lock().await;
|
||||
let request = captured.first().expect("server should capture request");
|
||||
assert_eq!(request.path, "/chat/completions");
|
||||
assert_eq!(
|
||||
request.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer xai-test-key")
|
||||
);
|
||||
let body: serde_json::Value = serde_json::from_str(&request.body).expect("json body");
|
||||
assert_eq!(body["model"], json!("grok-3"));
|
||||
assert_eq!(body["messages"][0]["role"], json!("system"));
|
||||
assert_eq!(body["tools"][0]["type"], json!("function"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_message_blocks_oversized_xai_requests_before_the_http_call() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response("200 OK", "application/json", "{}")],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = OpenAiCompatClient::new("xai-test-key", OpenAiCompatConfig::xai())
|
||||
.with_base_url(server.base_url());
|
||||
let error = client
|
||||
.send_message(&MessageRequest {
|
||||
model: "grok-3".to_string(),
|
||||
max_tokens: 64_000,
|
||||
messages: vec![InputMessage {
|
||||
role: "user".to_string(),
|
||||
content: vec![InputContentBlock::Text {
|
||||
text: "x".repeat(300_000),
|
||||
}],
|
||||
}],
|
||||
system: Some("Keep the answer short.".to_string()),
|
||||
tools: None,
|
||||
tool_choice: None,
|
||||
stream: false,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.expect_err("oversized request should fail local context-window preflight");
|
||||
|
||||
assert!(matches!(error, ApiError::ContextWindowExceeded { .. }));
|
||||
assert!(
|
||||
state.lock().await.is_empty(),
|
||||
"preflight failure should avoid any upstream HTTP request"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_message_accepts_full_chat_completions_endpoint_override() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let body = concat!(
|
||||
"{",
|
||||
"\"id\":\"chatcmpl_full_endpoint\",",
|
||||
"\"model\":\"grok-3\",",
|
||||
"\"choices\":[{",
|
||||
"\"message\":{\"role\":\"assistant\",\"content\":\"Endpoint override works\",\"tool_calls\":[]},",
|
||||
"\"finish_reason\":\"stop\"",
|
||||
"}],",
|
||||
"\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":3}",
|
||||
"}"
|
||||
);
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response("200 OK", "application/json", body)],
|
||||
)
|
||||
.await;
|
||||
|
||||
let endpoint_url = format!("{}/chat/completions", server.base_url());
|
||||
let client = OpenAiCompatClient::new("xai-test-key", OpenAiCompatConfig::xai())
|
||||
.with_base_url(endpoint_url);
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
assert_eq!(response.total_tokens(), 10);
|
||||
|
||||
let captured = state.lock().await;
|
||||
let request = captured.first().expect("server should capture request");
|
||||
assert_eq!(request.path, "/chat/completions");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stream_message_normalizes_text_and_multiple_tool_calls() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let sse = concat!(
|
||||
"data: {\"id\":\"chatcmpl_stream\",\"model\":\"grok-3\",\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n",
|
||||
"data: {\"id\":\"chatcmpl_stream\",\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"weather\",\"arguments\":\"{\\\"city\\\":\\\"Paris\\\"}\"}},{\"index\":1,\"id\":\"call_2\",\"function\":{\"name\":\"clock\",\"arguments\":\"{\\\"zone\\\":\\\"UTC\\\"}\"}}]}}]}\n\n",
|
||||
"data: {\"id\":\"chatcmpl_stream\",\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}]}\n\n",
|
||||
"data: [DONE]\n\n"
|
||||
);
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response_with_headers(
|
||||
"200 OK",
|
||||
"text/event-stream",
|
||||
sse,
|
||||
&[("x-request-id", "req_grok_stream")],
|
||||
)],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = OpenAiCompatClient::new("xai-test-key", OpenAiCompatConfig::xai())
|
||||
.with_base_url(server.base_url());
|
||||
let mut stream = client
|
||||
.stream_message(&sample_request(false))
|
||||
.await
|
||||
.expect("stream should start");
|
||||
|
||||
assert_eq!(stream.request_id(), Some("req_grok_stream"));
|
||||
|
||||
let mut events = Vec::new();
|
||||
while let Some(event) = stream.next_event().await.expect("event should parse") {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
assert!(matches!(events[0], StreamEvent::MessageStart(_)));
|
||||
assert!(matches!(
|
||||
events[1],
|
||||
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
|
||||
content_block: OutputContentBlock::Text { .. },
|
||||
..
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
events[2],
|
||||
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
|
||||
delta: ContentBlockDelta::TextDelta { .. },
|
||||
..
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
events[3],
|
||||
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
|
||||
index: 1,
|
||||
content_block: OutputContentBlock::ToolUse { .. },
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
events[4],
|
||||
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
|
||||
index: 1,
|
||||
delta: ContentBlockDelta::InputJsonDelta { .. },
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
events[5],
|
||||
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
|
||||
index: 2,
|
||||
content_block: OutputContentBlock::ToolUse { .. },
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
events[6],
|
||||
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
|
||||
index: 2,
|
||||
delta: ContentBlockDelta::InputJsonDelta { .. },
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
events[7],
|
||||
StreamEvent::ContentBlockStop(ContentBlockStopEvent { index: 1 })
|
||||
));
|
||||
assert!(matches!(
|
||||
events[8],
|
||||
StreamEvent::ContentBlockStop(ContentBlockStopEvent { index: 2 })
|
||||
));
|
||||
assert!(matches!(
|
||||
events[9],
|
||||
StreamEvent::ContentBlockStop(ContentBlockStopEvent { index: 0 })
|
||||
));
|
||||
assert!(matches!(events[10], StreamEvent::MessageDelta(_)));
|
||||
assert!(matches!(events[11], StreamEvent::MessageStop(_)));
|
||||
|
||||
let captured = state.lock().await;
|
||||
let request = captured.first().expect("captured request");
|
||||
assert_eq!(request.path, "/chat/completions");
|
||||
assert!(request.body.contains("\"stream\":true"));
|
||||
}
|
||||
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
#[tokio::test]
|
||||
async fn openai_streaming_requests_opt_into_usage_chunks() {
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let sse = concat!(
|
||||
"data: {\"id\":\"chatcmpl_openai_stream\",\"model\":\"gpt-5\",\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n",
|
||||
"data: {\"id\":\"chatcmpl_openai_stream\",\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
|
||||
"data: {\"id\":\"chatcmpl_openai_stream\",\"choices\":[],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":4}}\n\n",
|
||||
"data: [DONE]\n\n"
|
||||
);
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response_with_headers(
|
||||
"200 OK",
|
||||
"text/event-stream",
|
||||
sse,
|
||||
&[("x-request-id", "req_openai_stream")],
|
||||
)],
|
||||
)
|
||||
.await;
|
||||
|
||||
let client = OpenAiCompatClient::new("openai-test-key", OpenAiCompatConfig::openai())
|
||||
.with_base_url(server.base_url());
|
||||
let mut stream = client
|
||||
.stream_message(&sample_request(false))
|
||||
.await
|
||||
.expect("stream should start");
|
||||
|
||||
assert_eq!(stream.request_id(), Some("req_openai_stream"));
|
||||
|
||||
let mut events = Vec::new();
|
||||
while let Some(event) = stream.next_event().await.expect("event should parse") {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
assert!(matches!(events[0], StreamEvent::MessageStart(_)));
|
||||
assert!(matches!(
|
||||
events[1],
|
||||
StreamEvent::ContentBlockStart(ContentBlockStartEvent {
|
||||
content_block: OutputContentBlock::Text { .. },
|
||||
..
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
events[2],
|
||||
StreamEvent::ContentBlockDelta(ContentBlockDeltaEvent {
|
||||
delta: ContentBlockDelta::TextDelta { .. },
|
||||
..
|
||||
})
|
||||
));
|
||||
assert!(matches!(
|
||||
events[3],
|
||||
StreamEvent::ContentBlockStop(ContentBlockStopEvent { index: 0 })
|
||||
));
|
||||
assert!(matches!(
|
||||
events[4],
|
||||
StreamEvent::MessageDelta(MessageDeltaEvent { .. })
|
||||
));
|
||||
assert!(matches!(events[5], StreamEvent::MessageStop(_)));
|
||||
|
||||
match &events[4] {
|
||||
StreamEvent::MessageDelta(MessageDeltaEvent { usage, .. }) => {
|
||||
assert_eq!(usage.input_tokens, 9);
|
||||
assert_eq!(usage.output_tokens, 4);
|
||||
}
|
||||
other => panic!("expected message delta, got {other:?}"),
|
||||
}
|
||||
|
||||
let captured = state.lock().await;
|
||||
let request = captured.first().expect("captured request");
|
||||
assert_eq!(request.path, "/chat/completions");
|
||||
let body: serde_json::Value = serde_json::from_str(&request.body).expect("json body");
|
||||
assert_eq!(body["stream"], json!(true));
|
||||
assert_eq!(body["stream_options"], json!({"include_usage": true}));
|
||||
}
|
||||
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
#[tokio::test]
|
||||
async fn provider_client_dispatches_xai_requests_from_env() {
|
||||
let _lock = env_lock();
|
||||
let _api_key = ScopedEnvVar::set("XAI_API_KEY", "xai-test-key");
|
||||
|
||||
let state = Arc::new(Mutex::new(Vec::<CapturedRequest>::new()));
|
||||
let server = spawn_server(
|
||||
state.clone(),
|
||||
vec![http_response(
|
||||
"200 OK",
|
||||
"application/json",
|
||||
"{\"id\":\"chatcmpl_provider\",\"model\":\"grok-3\",\"choices\":[{\"message\":{\"role\":\"assistant\",\"content\":\"Through provider client\",\"tool_calls\":[]},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":4}}",
|
||||
)],
|
||||
)
|
||||
.await;
|
||||
let _base_url = ScopedEnvVar::set("XAI_BASE_URL", server.base_url());
|
||||
|
||||
let client =
|
||||
ProviderClient::from_model("grok").expect("xAI provider client should be constructed");
|
||||
assert!(matches!(client, ProviderClient::Xai(_)));
|
||||
|
||||
let response = client
|
||||
.send_message(&sample_request(false))
|
||||
.await
|
||||
.expect("provider-dispatched request should succeed");
|
||||
|
||||
assert_eq!(response.total_tokens(), 13);
|
||||
|
||||
let captured = state.lock().await;
|
||||
let request = captured.first().expect("captured request");
|
||||
assert_eq!(request.path, "/chat/completions");
|
||||
assert_eq!(
|
||||
request.headers.get("authorization").map(String::as_str),
|
||||
Some("Bearer xai-test-key")
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct CapturedRequest {
|
||||
path: String,
|
||||
headers: HashMap<String, String>,
|
||||
body: String,
|
||||
}
|
||||
|
||||
struct TestServer {
|
||||
base_url: String,
|
||||
join_handle: tokio::task::JoinHandle<()>,
|
||||
}
|
||||
|
||||
impl TestServer {
|
||||
fn base_url(&self) -> String {
|
||||
self.base_url.clone()
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for TestServer {
|
||||
fn drop(&mut self) {
|
||||
self.join_handle.abort();
|
||||
}
|
||||
}
|
||||
|
||||
async fn spawn_server(
|
||||
state: Arc<Mutex<Vec<CapturedRequest>>>,
|
||||
responses: Vec<String>,
|
||||
) -> TestServer {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let address = listener.local_addr().expect("listener addr");
|
||||
let join_handle = tokio::spawn(async move {
|
||||
for response in responses {
|
||||
let (mut socket, _) = listener.accept().await.expect("accept");
|
||||
let mut buffer = Vec::new();
|
||||
let mut header_end = None;
|
||||
loop {
|
||||
let mut chunk = [0_u8; 1024];
|
||||
let read = socket.read(&mut chunk).await.expect("read request");
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
buffer.extend_from_slice(&chunk[..read]);
|
||||
if let Some(position) = find_header_end(&buffer) {
|
||||
header_end = Some(position);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let header_end = header_end.expect("headers should exist");
|
||||
let (header_bytes, remaining) = buffer.split_at(header_end);
|
||||
let header_text = String::from_utf8(header_bytes.to_vec()).expect("utf8 headers");
|
||||
let mut lines = header_text.split("\r\n");
|
||||
let request_line = lines.next().expect("request line");
|
||||
let path = request_line
|
||||
.split_whitespace()
|
||||
.nth(1)
|
||||
.expect("path")
|
||||
.to_string();
|
||||
let mut headers = HashMap::new();
|
||||
let mut content_length = 0_usize;
|
||||
for line in lines {
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let (name, value) = line.split_once(':').expect("header");
|
||||
let value = value.trim().to_string();
|
||||
if name.eq_ignore_ascii_case("content-length") {
|
||||
content_length = value.parse().expect("content length");
|
||||
}
|
||||
headers.insert(name.to_ascii_lowercase(), value);
|
||||
}
|
||||
|
||||
let mut body = remaining[4..].to_vec();
|
||||
while body.len() < content_length {
|
||||
let mut chunk = vec![0_u8; content_length - body.len()];
|
||||
let read = socket.read(&mut chunk).await.expect("read body");
|
||||
if read == 0 {
|
||||
break;
|
||||
}
|
||||
body.extend_from_slice(&chunk[..read]);
|
||||
}
|
||||
|
||||
state.lock().await.push(CapturedRequest {
|
||||
path,
|
||||
headers,
|
||||
body: String::from_utf8(body).expect("utf8 body"),
|
||||
});
|
||||
|
||||
socket
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("write response");
|
||||
}
|
||||
});
|
||||
|
||||
TestServer {
|
||||
base_url: format!("http://{address}"),
|
||||
join_handle,
|
||||
}
|
||||
}
|
||||
|
||||
fn find_header_end(bytes: &[u8]) -> Option<usize> {
|
||||
bytes.windows(4).position(|window| window == b"\r\n\r\n")
|
||||
}
|
||||
|
||||
fn http_response(status: &str, content_type: &str, body: &str) -> String {
|
||||
http_response_with_headers(status, content_type, body, &[])
|
||||
}
|
||||
|
||||
fn http_response_with_headers(
|
||||
status: &str,
|
||||
content_type: &str,
|
||||
body: &str,
|
||||
headers: &[(&str, &str)],
|
||||
) -> String {
|
||||
let mut extra_headers = String::new();
|
||||
for (name, value) in headers {
|
||||
use std::fmt::Write as _;
|
||||
write!(&mut extra_headers, "{name}: {value}\r\n").expect("header write");
|
||||
}
|
||||
format!(
|
||||
"HTTP/1.1 {status}\r\ncontent-type: {content_type}\r\n{extra_headers}content-length: {}\r\nconnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
)
|
||||
}
|
||||
|
||||
fn sample_request(stream: bool) -> MessageRequest {
|
||||
MessageRequest {
|
||||
model: "grok-3".to_string(),
|
||||
max_tokens: 64,
|
||||
messages: vec![InputMessage {
|
||||
role: "user".to_string(),
|
||||
content: vec![InputContentBlock::Text {
|
||||
text: "Say hello".to_string(),
|
||||
}],
|
||||
}],
|
||||
system: Some("Use tools when needed".to_string()),
|
||||
tools: Some(vec![ToolDefinition {
|
||||
name: "weather".to_string(),
|
||||
description: Some("Fetches weather".to_string()),
|
||||
input_schema: json!({
|
||||
"type": "object",
|
||||
"properties": {"city": {"type": "string"}},
|
||||
"required": ["city"]
|
||||
}),
|
||||
}]),
|
||||
tool_choice: Some(ToolChoice::Auto),
|
||||
stream,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
static LOCK: OnceLock<StdMutex<()>> = OnceLock::new();
|
||||
LOCK.get_or_init(|| StdMutex::new(()))
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
struct ScopedEnvVar {
|
||||
key: &'static str,
|
||||
previous: Option<OsString>,
|
||||
}
|
||||
|
||||
impl ScopedEnvVar {
|
||||
fn set(key: &'static str, value: impl AsRef<std::ffi::OsStr>) -> Self {
|
||||
let previous = std::env::var_os(key);
|
||||
std::env::set_var(key, value);
|
||||
Self { key, previous }
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ScopedEnvVar {
|
||||
fn drop(&mut self) {
|
||||
match &self.previous {
|
||||
Some(value) => std::env::set_var(self.key, value),
|
||||
None => std::env::remove_var(self.key),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
use std::ffi::OsString;
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
|
||||
use api::{read_xai_base_url, ApiError, AuthSource, ProviderClient, ProviderKind};
|
||||
|
||||
#[test]
|
||||
fn provider_client_routes_grok_aliases_through_xai() {
|
||||
let _lock = env_lock();
|
||||
let _xai_api_key = EnvVarGuard::set("XAI_API_KEY", Some("xai-test-key"));
|
||||
|
||||
let client = ProviderClient::from_model("grok-mini").expect("grok alias should resolve");
|
||||
|
||||
assert_eq!(client.provider_kind(), ProviderKind::Xai);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_client_reports_missing_xai_credentials_for_grok_models() {
|
||||
let _lock = env_lock();
|
||||
let _xai_api_key = EnvVarGuard::set("XAI_API_KEY", None);
|
||||
|
||||
let error = ProviderClient::from_model("grok-3")
|
||||
.expect_err("grok requests without XAI_API_KEY should fail fast");
|
||||
|
||||
match error {
|
||||
ApiError::MissingCredentials {
|
||||
provider, env_vars, ..
|
||||
} => {
|
||||
assert_eq!(provider, "xAI");
|
||||
assert_eq!(env_vars, &["XAI_API_KEY"]);
|
||||
}
|
||||
other => panic!("expected missing xAI credentials, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_client_uses_explicit_anthropic_auth_without_env_lookup() {
|
||||
let _lock = env_lock();
|
||||
let _anthropic_api_key = EnvVarGuard::set("ANTHROPIC_API_KEY", None);
|
||||
let _anthropic_auth_token = EnvVarGuard::set("ANTHROPIC_AUTH_TOKEN", None);
|
||||
|
||||
let client = ProviderClient::from_model_with_anthropic_auth(
|
||||
"claude-sonnet-4-6",
|
||||
Some(AuthSource::ApiKey("anthropic-test-key".to_string())),
|
||||
)
|
||||
.expect("explicit anthropic auth should avoid env lookup");
|
||||
|
||||
assert_eq!(client.provider_kind(), ProviderKind::Anthropic);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_xai_base_url_prefers_env_override() {
|
||||
let _lock = env_lock();
|
||||
let _xai_base_url = EnvVarGuard::set("XAI_BASE_URL", Some("https://example.xai.test/v1"));
|
||||
|
||||
assert_eq!(read_xai_base_url(), "https://example.xai.test/v1");
|
||||
}
|
||||
|
||||
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
||||
LOCK.get_or_init(|| Mutex::new(()))
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
struct EnvVarGuard {
|
||||
key: &'static str,
|
||||
original: Option<OsString>,
|
||||
}
|
||||
|
||||
impl EnvVarGuard {
|
||||
fn set(key: &'static str, value: Option<&str>) -> Self {
|
||||
let original = std::env::var_os(key);
|
||||
match value {
|
||||
Some(value) => std::env::set_var(key, value),
|
||||
None => std::env::remove_var(key),
|
||||
}
|
||||
Self { key, original }
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for EnvVarGuard {
|
||||
fn drop(&mut self) {
|
||||
match &self.original {
|
||||
Some(value) => std::env::set_var(self.key, value),
|
||||
None => std::env::remove_var(self.key),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,173 @@
|
||||
use std::ffi::OsString;
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
|
||||
use api::{build_http_client_with, ProxyConfig};
|
||||
|
||||
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
|
||||
LOCK.get_or_init(|| Mutex::new(()))
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
|
||||
struct EnvVarGuard {
|
||||
key: &'static str,
|
||||
original: Option<OsString>,
|
||||
}
|
||||
|
||||
impl EnvVarGuard {
|
||||
fn set(key: &'static str, value: Option<&str>) -> Self {
|
||||
let original = std::env::var_os(key);
|
||||
match value {
|
||||
Some(value) => std::env::set_var(key, value),
|
||||
None => std::env::remove_var(key),
|
||||
}
|
||||
Self { key, original }
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for EnvVarGuard {
|
||||
fn drop(&mut self) {
|
||||
match &self.original {
|
||||
Some(value) => std::env::set_var(self.key, value),
|
||||
None => std::env::remove_var(self.key),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_from_env_reads_uppercase_proxy_vars() {
|
||||
// given
|
||||
let _lock = env_lock();
|
||||
let _http = EnvVarGuard::set("HTTP_PROXY", Some("http://proxy.corp:3128"));
|
||||
let _https = EnvVarGuard::set("HTTPS_PROXY", Some("http://secure.corp:3129"));
|
||||
let _no = EnvVarGuard::set("NO_PROXY", Some("localhost,127.0.0.1"));
|
||||
let _http_lower = EnvVarGuard::set("http_proxy", None);
|
||||
let _https_lower = EnvVarGuard::set("https_proxy", None);
|
||||
let _no_lower = EnvVarGuard::set("no_proxy", None);
|
||||
|
||||
// when
|
||||
let config = ProxyConfig::from_env();
|
||||
|
||||
// then
|
||||
assert_eq!(config.http_proxy.as_deref(), Some("http://proxy.corp:3128"));
|
||||
assert_eq!(
|
||||
config.https_proxy.as_deref(),
|
||||
Some("http://secure.corp:3129")
|
||||
);
|
||||
assert_eq!(config.no_proxy.as_deref(), Some("localhost,127.0.0.1"));
|
||||
assert!(config.proxy_url.is_none());
|
||||
assert!(!config.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_from_env_reads_lowercase_proxy_vars() {
|
||||
// given
|
||||
let _lock = env_lock();
|
||||
let _http = EnvVarGuard::set("HTTP_PROXY", None);
|
||||
let _https = EnvVarGuard::set("HTTPS_PROXY", None);
|
||||
let _no = EnvVarGuard::set("NO_PROXY", None);
|
||||
let _http_lower = EnvVarGuard::set("http_proxy", Some("http://lower.corp:3128"));
|
||||
let _https_lower = EnvVarGuard::set("https_proxy", Some("http://lower-secure.corp:3129"));
|
||||
let _no_lower = EnvVarGuard::set("no_proxy", Some(".internal"));
|
||||
|
||||
// when
|
||||
let config = ProxyConfig::from_env();
|
||||
|
||||
// then
|
||||
assert_eq!(config.http_proxy.as_deref(), Some("http://lower.corp:3128"));
|
||||
assert_eq!(
|
||||
config.https_proxy.as_deref(),
|
||||
Some("http://lower-secure.corp:3129")
|
||||
);
|
||||
assert_eq!(config.no_proxy.as_deref(), Some(".internal"));
|
||||
assert!(!config.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_from_env_is_empty_when_no_vars_set() {
|
||||
// given
|
||||
let _lock = env_lock();
|
||||
let _http = EnvVarGuard::set("HTTP_PROXY", None);
|
||||
let _https = EnvVarGuard::set("HTTPS_PROXY", None);
|
||||
let _no = EnvVarGuard::set("NO_PROXY", None);
|
||||
let _http_lower = EnvVarGuard::set("http_proxy", None);
|
||||
let _https_lower = EnvVarGuard::set("https_proxy", None);
|
||||
let _no_lower = EnvVarGuard::set("no_proxy", None);
|
||||
|
||||
// when
|
||||
let config = ProxyConfig::from_env();
|
||||
|
||||
// then
|
||||
assert!(config.is_empty());
|
||||
assert!(config.http_proxy.is_none());
|
||||
assert!(config.https_proxy.is_none());
|
||||
assert!(config.no_proxy.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_from_env_treats_empty_values_as_unset() {
|
||||
// given
|
||||
let _lock = env_lock();
|
||||
let _http = EnvVarGuard::set("HTTP_PROXY", Some(""));
|
||||
let _https = EnvVarGuard::set("HTTPS_PROXY", Some(""));
|
||||
let _http_lower = EnvVarGuard::set("http_proxy", Some(""));
|
||||
let _https_lower = EnvVarGuard::set("https_proxy", Some(""));
|
||||
let _no = EnvVarGuard::set("NO_PROXY", Some(""));
|
||||
let _no_lower = EnvVarGuard::set("no_proxy", Some(""));
|
||||
|
||||
// when
|
||||
let config = ProxyConfig::from_env();
|
||||
|
||||
// then
|
||||
assert!(config.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_client_with_env_proxy_config_succeeds() {
|
||||
// given
|
||||
let _lock = env_lock();
|
||||
let _http = EnvVarGuard::set("HTTP_PROXY", Some("http://proxy.corp:3128"));
|
||||
let _https = EnvVarGuard::set("HTTPS_PROXY", Some("http://secure.corp:3129"));
|
||||
let _no = EnvVarGuard::set("NO_PROXY", Some("localhost"));
|
||||
let _http_lower = EnvVarGuard::set("http_proxy", None);
|
||||
let _https_lower = EnvVarGuard::set("https_proxy", None);
|
||||
let _no_lower = EnvVarGuard::set("no_proxy", None);
|
||||
let config = ProxyConfig::from_env();
|
||||
|
||||
// when
|
||||
let result = build_http_client_with(&config);
|
||||
|
||||
// then
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_client_with_proxy_url_config_succeeds() {
|
||||
// given
|
||||
let config = ProxyConfig::from_proxy_url("http://unified.corp:3128");
|
||||
|
||||
// when
|
||||
let result = build_http_client_with(&config);
|
||||
|
||||
// then
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn proxy_config_from_env_prefers_uppercase_over_lowercase() {
|
||||
// given
|
||||
let _lock = env_lock();
|
||||
let _http_upper = EnvVarGuard::set("HTTP_PROXY", Some("http://upper.corp:3128"));
|
||||
let _http_lower = EnvVarGuard::set("http_proxy", Some("http://lower.corp:3128"));
|
||||
let _https = EnvVarGuard::set("HTTPS_PROXY", None);
|
||||
let _https_lower = EnvVarGuard::set("https_proxy", None);
|
||||
let _no = EnvVarGuard::set("NO_PROXY", None);
|
||||
let _no_lower = EnvVarGuard::set("no_proxy", None);
|
||||
|
||||
// when
|
||||
let config = ProxyConfig::from_env();
|
||||
|
||||
// then
|
||||
assert_eq!(config.http_proxy.as_deref(), Some("http://upper.corp:3128"));
|
||||
}
|
||||
@@ -1,24 +1,28 @@
|
||||
[package]
|
||||
name = "rusty-claude-cli"
|
||||
name = "claw-cli"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
publish.workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "claw"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
api = { path = "../api" }
|
||||
commands = { path = "../commands" }
|
||||
compat-harness = { path = "../compat-harness" }
|
||||
crossterm = "0.28"
|
||||
pulldown-cmark = "0.13"
|
||||
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
|
||||
rustyline = "15"
|
||||
runtime = { path = "../runtime" }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
sha2 = "0.10"
|
||||
plugins = { path = "../plugins" }
|
||||
serde_json.workspace = true
|
||||
syntect = "5"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "time"] }
|
||||
tools = { path = "../tools" }
|
||||
dotenvy = "0.15"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,58 @@
|
||||
# Claw CLI 模块 (claw-cli)
|
||||
|
||||
本模块实现了 Claw 应用程序的主要命令行界面 (CLI)。它提供了交互式的 REPL 环境和非交互式的命令执行功能。
|
||||
|
||||
## 概览
|
||||
|
||||
`claw-cli` 模块是整个项目的“胶水”,负责编排 `runtime`、`api`、`tools` 和 `plugins` 模块之间的交互。它捕获用户输入,管理应用程序状态,并以用户友好的格式渲染 AI 响应。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **交互式 REPL**:一个功能齐全的 Read-Eval-Print Loop,用于与 AI 对话,支持:
|
||||
- 通过 `rustyline` 实现多行输入和命令历史。
|
||||
- 实时的 Markdown 流式显示和代码语法高亮。
|
||||
- 为耗时的工具调用显示动态的加载动画 (Spinner) 和进度指示器。
|
||||
- **子命令**:
|
||||
- `prompt`:运行单次 Prompt 并退出(单次模式)。
|
||||
- `login`/`logout`:处理与 Claw 平台的 OAuth 身份验证。
|
||||
- `init`:初始化新项目/仓库以配合 Claw 使用。
|
||||
- `resume`:恢复并继续之前的对话会话。
|
||||
- `agents`/`skills`:管理并发现可用的智能体和技能。
|
||||
- **权限管理**:对工具执行权限的细粒度控制:
|
||||
- `read-only`:安全的分析模式。
|
||||
- `workspace-write`:允许在当前工作区内进行修改。
|
||||
- `danger-full-access`:高级任务的无限制访问。
|
||||
- **OAuth 流程**:集成了本地 HTTP 服务器,无缝处理用户机器上的 OAuth 回调重定向。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`main.rs`**: 主要入口点。负责解析命令行参数、初始化环境,并调度到适当的操作(REPL 或子命令)。
|
||||
- **`render/`**: 包含 `TerminalRenderer` 和 Markdown 流渲染逻辑。使用 `syntect` 进行语法高亮,并使用 `crossterm` 进行终端操作。
|
||||
- **`input/`**: 处理用户输入捕获,包括对多行 Prompt 和斜杠命令 (Slash Command) 的特殊处理。
|
||||
- **`init.rs`**: 处理项目级初始化和仓库设置。
|
||||
|
||||
### 交互循环流程
|
||||
|
||||
1. CLI 初始化 `ConversationRuntime` 并加载项目上下文。
|
||||
2. 使用 `rustyline` 进入捕获用户输入的循环。
|
||||
3. 检查用户输入是否为“斜杠命令”(如 `/compact`、`/model`)。
|
||||
4. 普通 Prompt 通过 `runtime` 发送给 AI。
|
||||
5. AI 事件(文本增量、工具调用)逐步渲染到终端。
|
||||
6. 会话定期保存,以便将来可以恢复。
|
||||
|
||||
## 使用方法
|
||||
|
||||
主二进制程序名为 `claw`。
|
||||
|
||||
```bash
|
||||
# 启动交互式 REPL
|
||||
claw
|
||||
|
||||
# 运行单次 Prompt
|
||||
claw prompt "解释一下这个项目的架构"
|
||||
|
||||
# 登录服务
|
||||
claw login
|
||||
```
|
||||
@@ -112,7 +112,7 @@ impl CliApp {
|
||||
|
||||
pub fn run_repl(&mut self) -> io::Result<()> {
|
||||
let mut editor = LineEditor::new("› ", Vec::new());
|
||||
println!("Rusty Claude CLI interactive mode");
|
||||
println!("Claw Code interactive mode");
|
||||
println!("Type /help for commands. Shift+Enter or Ctrl+J inserts a newline.");
|
||||
|
||||
loop {
|
||||
@@ -162,6 +162,10 @@ impl CliApp {
|
||||
writeln!(out, "Unknown slash command: /{name}")?;
|
||||
Ok(CommandResult::Continue)
|
||||
}
|
||||
_ => {
|
||||
writeln!(out, "Slash command unavailable in this mode")?;
|
||||
Ok(CommandResult::Continue)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -172,7 +176,7 @@ impl CliApp {
|
||||
SlashCommand::Help => "/help",
|
||||
SlashCommand::Status => "/status",
|
||||
SlashCommand::Compact => "/compact",
|
||||
SlashCommand::Unknown(_) => continue,
|
||||
_ => continue,
|
||||
};
|
||||
writeln!(out, " {name:<9} {}", handler.summary)?;
|
||||
}
|
||||
@@ -385,14 +389,14 @@ mod tests {
|
||||
#[test]
|
||||
fn session_state_tracks_config_values() {
|
||||
let config = SessionConfig {
|
||||
model: "claude".into(),
|
||||
permission_mode: PermissionMode::WorkspaceWrite,
|
||||
model: "sonnet".into(),
|
||||
permission_mode: PermissionMode::DangerFullAccess,
|
||||
config: Some(PathBuf::from("settings.toml")),
|
||||
output_format: OutputFormat::Text,
|
||||
};
|
||||
|
||||
assert_eq!(config.model, "claude");
|
||||
assert_eq!(config.permission_mode, PermissionMode::WorkspaceWrite);
|
||||
assert_eq!(config.model, "sonnet");
|
||||
assert_eq!(config.permission_mode, PermissionMode::DangerFullAccess);
|
||||
assert_eq!(config.config, Some(PathBuf::from("settings.toml")));
|
||||
}
|
||||
}
|
||||
@@ -3,16 +3,12 @@ use std::path::PathBuf;
|
||||
use clap::{Parser, Subcommand, ValueEnum};
|
||||
|
||||
#[derive(Debug, Clone, Parser, PartialEq, Eq)]
|
||||
#[command(
|
||||
name = "rusty-claude-cli",
|
||||
version,
|
||||
about = "Rust Claude CLI prototype"
|
||||
)]
|
||||
#[command(name = "claw-cli", version, about = "Claw Code CLI")]
|
||||
pub struct Cli {
|
||||
#[arg(long, default_value = "claude-3-7-sonnet")]
|
||||
#[arg(long, default_value = "claude-opus-4-6")]
|
||||
pub model: String,
|
||||
|
||||
#[arg(long, value_enum, default_value_t = PermissionMode::WorkspaceWrite)]
|
||||
#[arg(long, value_enum, default_value_t = PermissionMode::DangerFullAccess)]
|
||||
pub permission_mode: PermissionMode,
|
||||
|
||||
#[arg(long)]
|
||||
@@ -62,9 +58,9 @@ mod tests {
|
||||
#[test]
|
||||
fn parses_requested_flags() {
|
||||
let cli = Cli::parse_from([
|
||||
"rusty-claude-cli",
|
||||
"claw-cli",
|
||||
"--model",
|
||||
"claude-3-5-haiku",
|
||||
"claude-haiku-4-5-20251213",
|
||||
"--permission-mode",
|
||||
"read-only",
|
||||
"--config",
|
||||
@@ -76,7 +72,7 @@ mod tests {
|
||||
"world",
|
||||
]);
|
||||
|
||||
assert_eq!(cli.model, "claude-3-5-haiku");
|
||||
assert_eq!(cli.model, "claude-haiku-4-5-20251213");
|
||||
assert_eq!(cli.permission_mode, PermissionMode::ReadOnly);
|
||||
assert_eq!(
|
||||
cli.config.as_deref(),
|
||||
@@ -93,10 +89,16 @@ mod tests {
|
||||
|
||||
#[test]
|
||||
fn parses_login_and_logout_commands() {
|
||||
let login = Cli::parse_from(["rusty-claude-cli", "login"]);
|
||||
let login = Cli::parse_from(["claw-cli", "login"]);
|
||||
assert_eq!(login.command, Some(Command::Login));
|
||||
|
||||
let logout = Cli::parse_from(["rusty-claude-cli", "logout"]);
|
||||
let logout = Cli::parse_from(["claw-cli", "logout"]);
|
||||
assert_eq!(logout.command, Some(Command::Logout));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defaults_to_danger_full_access_permission_mode() {
|
||||
let cli = Cli::parse_from(["claw-cli"]);
|
||||
assert_eq!(cli.permission_mode, PermissionMode::DangerFullAccess);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,432 @@
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
const STARTER_CLAW_JSON: &str = concat!(
|
||||
"{\n",
|
||||
" \"permissions\": {\n",
|
||||
" \"defaultMode\": \"dontAsk\"\n",
|
||||
" }\n",
|
||||
"}\n",
|
||||
);
|
||||
const GITIGNORE_COMMENT: &str = "# Claw Code local artifacts";
|
||||
const GITIGNORE_ENTRIES: [&str; 2] = [".claw/settings.local.json", ".claw/sessions/"];
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum InitStatus {
|
||||
Created,
|
||||
Updated,
|
||||
Skipped,
|
||||
}
|
||||
|
||||
impl InitStatus {
|
||||
#[must_use]
|
||||
pub(crate) fn label(self) -> &'static str {
|
||||
match self {
|
||||
Self::Created => "created",
|
||||
Self::Updated => "updated",
|
||||
Self::Skipped => "skipped (already exists)",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct InitArtifact {
|
||||
pub(crate) name: &'static str,
|
||||
pub(crate) status: InitStatus,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct InitReport {
|
||||
pub(crate) project_root: PathBuf,
|
||||
pub(crate) artifacts: Vec<InitArtifact>,
|
||||
}
|
||||
|
||||
impl InitReport {
|
||||
#[must_use]
|
||||
pub(crate) fn render(&self) -> String {
|
||||
let mut lines = vec![
|
||||
"Init".to_string(),
|
||||
format!(" Project {}", self.project_root.display()),
|
||||
];
|
||||
for artifact in &self.artifacts {
|
||||
lines.push(format!(
|
||||
" {:<16} {}",
|
||||
artifact.name,
|
||||
artifact.status.label()
|
||||
));
|
||||
}
|
||||
lines.push(" Next step Review and tailor the generated guidance".to_string());
|
||||
lines.join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
#[allow(clippy::struct_excessive_bools)]
|
||||
struct RepoDetection {
|
||||
rust_workspace: bool,
|
||||
rust_root: bool,
|
||||
python: bool,
|
||||
package_json: bool,
|
||||
typescript: bool,
|
||||
nextjs: bool,
|
||||
react: bool,
|
||||
vite: bool,
|
||||
nest: bool,
|
||||
src_dir: bool,
|
||||
tests_dir: bool,
|
||||
rust_dir: bool,
|
||||
}
|
||||
|
||||
pub(crate) fn initialize_repo(cwd: &Path) -> Result<InitReport, Box<dyn std::error::Error>> {
|
||||
let mut artifacts = Vec::new();
|
||||
|
||||
let claw_dir = cwd.join(".claw");
|
||||
artifacts.push(InitArtifact {
|
||||
name: ".claw/",
|
||||
status: ensure_dir(&claw_dir)?,
|
||||
});
|
||||
|
||||
let claw_json = cwd.join(".claw.json");
|
||||
artifacts.push(InitArtifact {
|
||||
name: ".claw.json",
|
||||
status: write_file_if_missing(&claw_json, STARTER_CLAW_JSON)?,
|
||||
});
|
||||
|
||||
let gitignore = cwd.join(".gitignore");
|
||||
artifacts.push(InitArtifact {
|
||||
name: ".gitignore",
|
||||
status: ensure_gitignore_entries(&gitignore)?,
|
||||
});
|
||||
|
||||
let claw_md = cwd.join("CLAW.md");
|
||||
let content = render_init_claw_md(cwd);
|
||||
artifacts.push(InitArtifact {
|
||||
name: "CLAW.md",
|
||||
status: write_file_if_missing(&claw_md, &content)?,
|
||||
});
|
||||
|
||||
Ok(InitReport {
|
||||
project_root: cwd.to_path_buf(),
|
||||
artifacts,
|
||||
})
|
||||
}
|
||||
|
||||
fn ensure_dir(path: &Path) -> Result<InitStatus, std::io::Error> {
|
||||
if path.is_dir() {
|
||||
return Ok(InitStatus::Skipped);
|
||||
}
|
||||
fs::create_dir_all(path)?;
|
||||
Ok(InitStatus::Created)
|
||||
}
|
||||
|
||||
fn write_file_if_missing(path: &Path, content: &str) -> Result<InitStatus, std::io::Error> {
|
||||
if path.exists() {
|
||||
return Ok(InitStatus::Skipped);
|
||||
}
|
||||
fs::write(path, content)?;
|
||||
Ok(InitStatus::Created)
|
||||
}
|
||||
|
||||
fn ensure_gitignore_entries(path: &Path) -> Result<InitStatus, std::io::Error> {
|
||||
if !path.exists() {
|
||||
let mut lines = vec![GITIGNORE_COMMENT.to_string()];
|
||||
lines.extend(GITIGNORE_ENTRIES.iter().map(|entry| (*entry).to_string()));
|
||||
fs::write(path, format!("{}\n", lines.join("\n")))?;
|
||||
return Ok(InitStatus::Created);
|
||||
}
|
||||
|
||||
let existing = fs::read_to_string(path)?;
|
||||
let mut lines = existing.lines().map(ToOwned::to_owned).collect::<Vec<_>>();
|
||||
let mut changed = false;
|
||||
|
||||
if !lines.iter().any(|line| line == GITIGNORE_COMMENT) {
|
||||
lines.push(GITIGNORE_COMMENT.to_string());
|
||||
changed = true;
|
||||
}
|
||||
|
||||
for entry in GITIGNORE_ENTRIES {
|
||||
if !lines.iter().any(|line| line == entry) {
|
||||
lines.push(entry.to_string());
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
|
||||
if !changed {
|
||||
return Ok(InitStatus::Skipped);
|
||||
}
|
||||
|
||||
fs::write(path, format!("{}\n", lines.join("\n")))?;
|
||||
Ok(InitStatus::Updated)
|
||||
}
|
||||
|
||||
pub(crate) fn render_init_claw_md(cwd: &Path) -> String {
|
||||
let detection = detect_repo(cwd);
|
||||
let mut lines = vec![
|
||||
"# CLAW.md".to_string(),
|
||||
String::new(),
|
||||
"This file provides guidance to Claw Code (clawcode.dev) when working with code in this repository.".to_string(),
|
||||
String::new(),
|
||||
];
|
||||
|
||||
let detected_languages = detected_languages(&detection);
|
||||
let detected_frameworks = detected_frameworks(&detection);
|
||||
lines.push("## Detected stack".to_string());
|
||||
if detected_languages.is_empty() {
|
||||
lines.push("- No specific language markers were detected yet; document the primary language and verification commands once the project structure settles.".to_string());
|
||||
} else {
|
||||
lines.push(format!("- Languages: {}.", detected_languages.join(", ")));
|
||||
}
|
||||
if detected_frameworks.is_empty() {
|
||||
lines.push("- Frameworks: none detected from the supported starter markers.".to_string());
|
||||
} else {
|
||||
lines.push(format!(
|
||||
"- Frameworks/tooling markers: {}.",
|
||||
detected_frameworks.join(", ")
|
||||
));
|
||||
}
|
||||
lines.push(String::new());
|
||||
|
||||
let verification_lines = verification_lines(cwd, &detection);
|
||||
if !verification_lines.is_empty() {
|
||||
lines.push("## Verification".to_string());
|
||||
lines.extend(verification_lines);
|
||||
lines.push(String::new());
|
||||
}
|
||||
|
||||
let structure_lines = repository_shape_lines(&detection);
|
||||
if !structure_lines.is_empty() {
|
||||
lines.push("## Repository shape".to_string());
|
||||
lines.extend(structure_lines);
|
||||
lines.push(String::new());
|
||||
}
|
||||
|
||||
let framework_lines = framework_notes(&detection);
|
||||
if !framework_lines.is_empty() {
|
||||
lines.push("## Framework notes".to_string());
|
||||
lines.extend(framework_lines);
|
||||
lines.push(String::new());
|
||||
}
|
||||
|
||||
lines.push("## Working agreement".to_string());
|
||||
lines.push("- Prefer small, reviewable changes and keep generated bootstrap files aligned with actual repo workflows.".to_string());
|
||||
lines.push("- Keep shared defaults in `.claw.json`; reserve `.claw/settings.local.json` for machine-local overrides.".to_string());
|
||||
lines.push("- Do not overwrite existing `CLAW.md` content automatically; update it intentionally when repo workflows change.".to_string());
|
||||
lines.push(String::new());
|
||||
|
||||
lines.join("\n")
|
||||
}
|
||||
|
||||
fn detect_repo(cwd: &Path) -> RepoDetection {
|
||||
let package_json_contents = fs::read_to_string(cwd.join("package.json"))
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase();
|
||||
RepoDetection {
|
||||
rust_workspace: cwd.join("rust").join("Cargo.toml").is_file(),
|
||||
rust_root: cwd.join("Cargo.toml").is_file(),
|
||||
python: cwd.join("pyproject.toml").is_file()
|
||||
|| cwd.join("requirements.txt").is_file()
|
||||
|| cwd.join("setup.py").is_file(),
|
||||
package_json: cwd.join("package.json").is_file(),
|
||||
typescript: cwd.join("tsconfig.json").is_file()
|
||||
|| package_json_contents.contains("typescript"),
|
||||
nextjs: package_json_contents.contains("\"next\""),
|
||||
react: package_json_contents.contains("\"react\""),
|
||||
vite: package_json_contents.contains("\"vite\""),
|
||||
nest: package_json_contents.contains("@nestjs"),
|
||||
src_dir: cwd.join("src").is_dir(),
|
||||
tests_dir: cwd.join("tests").is_dir(),
|
||||
rust_dir: cwd.join("rust").is_dir(),
|
||||
}
|
||||
}
|
||||
|
||||
fn detected_languages(detection: &RepoDetection) -> Vec<&'static str> {
|
||||
let mut languages = Vec::new();
|
||||
if detection.rust_workspace || detection.rust_root {
|
||||
languages.push("Rust");
|
||||
}
|
||||
if detection.python {
|
||||
languages.push("Python");
|
||||
}
|
||||
if detection.typescript {
|
||||
languages.push("TypeScript");
|
||||
} else if detection.package_json {
|
||||
languages.push("JavaScript/Node.js");
|
||||
}
|
||||
languages
|
||||
}
|
||||
|
||||
fn detected_frameworks(detection: &RepoDetection) -> Vec<&'static str> {
|
||||
let mut frameworks = Vec::new();
|
||||
if detection.nextjs {
|
||||
frameworks.push("Next.js");
|
||||
}
|
||||
if detection.react {
|
||||
frameworks.push("React");
|
||||
}
|
||||
if detection.vite {
|
||||
frameworks.push("Vite");
|
||||
}
|
||||
if detection.nest {
|
||||
frameworks.push("NestJS");
|
||||
}
|
||||
frameworks
|
||||
}
|
||||
|
||||
fn verification_lines(cwd: &Path, detection: &RepoDetection) -> Vec<String> {
|
||||
let mut lines = Vec::new();
|
||||
if detection.rust_workspace {
|
||||
lines.push("- Run Rust verification from `rust/`: `cargo fmt`, `cargo clippy --workspace --all-targets -- -D warnings`, `cargo test --workspace`".to_string());
|
||||
} else if detection.rust_root {
|
||||
lines.push("- Run Rust verification from the repo root: `cargo fmt`, `cargo clippy --workspace --all-targets -- -D warnings`, `cargo test --workspace`".to_string());
|
||||
}
|
||||
if detection.python {
|
||||
if cwd.join("pyproject.toml").is_file() {
|
||||
lines.push("- Run the Python project checks declared in `pyproject.toml` (for example: `pytest`, `ruff check`, and `mypy` when configured).".to_string());
|
||||
} else {
|
||||
lines.push(
|
||||
"- Run the repo's Python test/lint commands before shipping changes.".to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
if detection.package_json {
|
||||
lines.push("- Run the JavaScript/TypeScript checks from `package.json` before shipping changes (`npm test`, `npm run lint`, `npm run build`, or the repo equivalent).".to_string());
|
||||
}
|
||||
if detection.tests_dir && detection.src_dir {
|
||||
lines.push("- `src/` and `tests/` are both present; update both surfaces together when behavior changes.".to_string());
|
||||
}
|
||||
lines
|
||||
}
|
||||
|
||||
fn repository_shape_lines(detection: &RepoDetection) -> Vec<String> {
|
||||
let mut lines = Vec::new();
|
||||
if detection.rust_dir {
|
||||
lines.push(
|
||||
"- `rust/` contains the Rust workspace and active CLI/runtime implementation."
|
||||
.to_string(),
|
||||
);
|
||||
}
|
||||
if detection.src_dir {
|
||||
lines.push("- `src/` contains source files that should stay consistent with generated guidance and tests.".to_string());
|
||||
}
|
||||
if detection.tests_dir {
|
||||
lines.push("- `tests/` contains validation surfaces that should be reviewed alongside code changes.".to_string());
|
||||
}
|
||||
lines
|
||||
}
|
||||
|
||||
fn framework_notes(detection: &RepoDetection) -> Vec<String> {
|
||||
let mut lines = Vec::new();
|
||||
if detection.nextjs {
|
||||
lines.push("- Next.js detected: preserve routing/data-fetching conventions and verify production builds after changing app structure.".to_string());
|
||||
}
|
||||
if detection.react && !detection.nextjs {
|
||||
lines.push("- React detected: keep component behavior covered with focused tests and avoid unnecessary prop/API churn.".to_string());
|
||||
}
|
||||
if detection.vite {
|
||||
lines.push("- Vite detected: validate the production bundle after changing build-sensitive configuration or imports.".to_string());
|
||||
}
|
||||
if detection.nest {
|
||||
lines.push("- NestJS detected: keep module/provider boundaries explicit and verify controller/service wiring after refactors.".to_string());
|
||||
}
|
||||
lines
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{initialize_repo, render_init_claw_md};
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_dir() -> std::path::PathBuf {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!("claw-init-{nanos}"))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn initialize_repo_creates_expected_files_and_gitignore_entries() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(root.join("rust")).expect("create rust dir");
|
||||
fs::write(root.join("rust").join("Cargo.toml"), "[workspace]\n").expect("write cargo");
|
||||
|
||||
let report = initialize_repo(&root).expect("init should succeed");
|
||||
let rendered = report.render();
|
||||
assert!(rendered.contains(".claw/ created"));
|
||||
assert!(rendered.contains(".claw.json created"));
|
||||
assert!(rendered.contains(".gitignore created"));
|
||||
assert!(rendered.contains("CLAW.md created"));
|
||||
assert!(root.join(".claw").is_dir());
|
||||
assert!(root.join(".claw.json").is_file());
|
||||
assert!(root.join("CLAW.md").is_file());
|
||||
assert_eq!(
|
||||
fs::read_to_string(root.join(".claw.json")).expect("read claw json"),
|
||||
concat!(
|
||||
"{\n",
|
||||
" \"permissions\": {\n",
|
||||
" \"defaultMode\": \"dontAsk\"\n",
|
||||
" }\n",
|
||||
"}\n",
|
||||
)
|
||||
);
|
||||
let gitignore = fs::read_to_string(root.join(".gitignore")).expect("read gitignore");
|
||||
assert!(gitignore.contains(".claw/settings.local.json"));
|
||||
assert!(gitignore.contains(".claw/sessions/"));
|
||||
let claw_md = fs::read_to_string(root.join("CLAW.md")).expect("read claw md");
|
||||
assert!(claw_md.contains("Languages: Rust."));
|
||||
assert!(claw_md.contains("cargo clippy --workspace --all-targets -- -D warnings"));
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn initialize_repo_is_idempotent_and_preserves_existing_files() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create root");
|
||||
fs::write(root.join("CLAW.md"), "custom guidance\n").expect("write existing claw md");
|
||||
fs::write(root.join(".gitignore"), ".claw/settings.local.json\n").expect("write gitignore");
|
||||
|
||||
let first = initialize_repo(&root).expect("first init should succeed");
|
||||
assert!(first
|
||||
.render()
|
||||
.contains("CLAW.md skipped (already exists)"));
|
||||
let second = initialize_repo(&root).expect("second init should succeed");
|
||||
let second_rendered = second.render();
|
||||
assert!(second_rendered.contains(".claw/ skipped (already exists)"));
|
||||
assert!(second_rendered.contains(".claw.json skipped (already exists)"));
|
||||
assert!(second_rendered.contains(".gitignore skipped (already exists)"));
|
||||
assert!(second_rendered.contains("CLAW.md skipped (already exists)"));
|
||||
assert_eq!(
|
||||
fs::read_to_string(root.join("CLAW.md")).expect("read existing claw md"),
|
||||
"custom guidance\n"
|
||||
);
|
||||
let gitignore = fs::read_to_string(root.join(".gitignore")).expect("read gitignore");
|
||||
assert_eq!(gitignore.matches(".claw/settings.local.json").count(), 1);
|
||||
assert_eq!(gitignore.matches(".claw/sessions/").count(), 1);
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_init_template_mentions_detected_python_and_nextjs_markers() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create root");
|
||||
fs::write(root.join("pyproject.toml"), "[project]\nname = \"demo\"\n")
|
||||
.expect("write pyproject");
|
||||
fs::write(
|
||||
root.join("package.json"),
|
||||
r#"{"dependencies":{"next":"14.0.0","react":"18.0.0"},"devDependencies":{"typescript":"5.0.0"}}"#,
|
||||
)
|
||||
.expect("write package json");
|
||||
|
||||
let rendered = render_init_claw_md(Path::new(&root));
|
||||
assert!(rendered.contains("Languages: Python, TypeScript."));
|
||||
assert!(rendered.contains("Frameworks/tooling markers: Next.js, React."));
|
||||
assert!(rendered.contains("pyproject.toml"));
|
||||
assert!(rendered.contains("Next.js detected"));
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,797 @@
|
||||
use std::fmt::Write as FmtWrite;
|
||||
use std::io::{self, Write};
|
||||
|
||||
use crossterm::cursor::{MoveToColumn, RestorePosition, SavePosition};
|
||||
use crossterm::style::{Color, Print, ResetColor, SetForegroundColor, Stylize};
|
||||
use crossterm::terminal::{Clear, ClearType};
|
||||
use crossterm::{execute, queue};
|
||||
use pulldown_cmark::{CodeBlockKind, Event, Options, Parser, Tag, TagEnd};
|
||||
use syntect::easy::HighlightLines;
|
||||
use syntect::highlighting::{Theme, ThemeSet};
|
||||
use syntect::parsing::SyntaxSet;
|
||||
use syntect::util::{as_24_bit_terminal_escaped, LinesWithEndings};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct ColorTheme {
|
||||
heading: Color,
|
||||
emphasis: Color,
|
||||
strong: Color,
|
||||
inline_code: Color,
|
||||
link: Color,
|
||||
quote: Color,
|
||||
table_border: Color,
|
||||
code_block_border: Color,
|
||||
spinner_active: Color,
|
||||
spinner_done: Color,
|
||||
spinner_failed: Color,
|
||||
}
|
||||
|
||||
impl Default for ColorTheme {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
heading: Color::Cyan,
|
||||
emphasis: Color::Magenta,
|
||||
strong: Color::Yellow,
|
||||
inline_code: Color::Green,
|
||||
link: Color::Blue,
|
||||
quote: Color::DarkGrey,
|
||||
table_border: Color::DarkCyan,
|
||||
code_block_border: Color::DarkGrey,
|
||||
spinner_active: Color::Blue,
|
||||
spinner_done: Color::Green,
|
||||
spinner_failed: Color::Red,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
pub struct Spinner {
|
||||
frame_index: usize,
|
||||
}
|
||||
|
||||
impl Spinner {
|
||||
const FRAMES: [&str; 10] = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
|
||||
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn tick(
|
||||
&mut self,
|
||||
label: &str,
|
||||
theme: &ColorTheme,
|
||||
out: &mut impl Write,
|
||||
) -> io::Result<()> {
|
||||
let frame = Self::FRAMES[self.frame_index % Self::FRAMES.len()];
|
||||
self.frame_index += 1;
|
||||
queue!(
|
||||
out,
|
||||
SavePosition,
|
||||
MoveToColumn(0),
|
||||
Clear(ClearType::CurrentLine),
|
||||
SetForegroundColor(theme.spinner_active),
|
||||
Print(format!("{frame} {label}")),
|
||||
ResetColor,
|
||||
RestorePosition
|
||||
)?;
|
||||
out.flush()
|
||||
}
|
||||
|
||||
pub fn finish(
|
||||
&mut self,
|
||||
label: &str,
|
||||
theme: &ColorTheme,
|
||||
out: &mut impl Write,
|
||||
) -> io::Result<()> {
|
||||
self.frame_index = 0;
|
||||
execute!(
|
||||
out,
|
||||
MoveToColumn(0),
|
||||
Clear(ClearType::CurrentLine),
|
||||
SetForegroundColor(theme.spinner_done),
|
||||
Print(format!("✔ {label}\n")),
|
||||
ResetColor
|
||||
)?;
|
||||
out.flush()
|
||||
}
|
||||
|
||||
pub fn fail(
|
||||
&mut self,
|
||||
label: &str,
|
||||
theme: &ColorTheme,
|
||||
out: &mut impl Write,
|
||||
) -> io::Result<()> {
|
||||
self.frame_index = 0;
|
||||
execute!(
|
||||
out,
|
||||
MoveToColumn(0),
|
||||
Clear(ClearType::CurrentLine),
|
||||
SetForegroundColor(theme.spinner_failed),
|
||||
Print(format!("✘ {label}\n")),
|
||||
ResetColor
|
||||
)?;
|
||||
out.flush()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum ListKind {
|
||||
Unordered,
|
||||
Ordered { next_index: u64 },
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
struct TableState {
|
||||
headers: Vec<String>,
|
||||
rows: Vec<Vec<String>>,
|
||||
current_row: Vec<String>,
|
||||
current_cell: String,
|
||||
in_head: bool,
|
||||
}
|
||||
|
||||
impl TableState {
|
||||
fn push_cell(&mut self) {
|
||||
let cell = self.current_cell.trim().to_string();
|
||||
self.current_row.push(cell);
|
||||
self.current_cell.clear();
|
||||
}
|
||||
|
||||
fn finish_row(&mut self) {
|
||||
if self.current_row.is_empty() {
|
||||
return;
|
||||
}
|
||||
let row = std::mem::take(&mut self.current_row);
|
||||
if self.in_head {
|
||||
self.headers = row;
|
||||
} else {
|
||||
self.rows.push(row);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
struct RenderState {
|
||||
emphasis: usize,
|
||||
strong: usize,
|
||||
heading_level: Option<u8>,
|
||||
quote: usize,
|
||||
list_stack: Vec<ListKind>,
|
||||
link_stack: Vec<LinkState>,
|
||||
table: Option<TableState>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct LinkState {
|
||||
destination: String,
|
||||
text: String,
|
||||
}
|
||||
|
||||
impl RenderState {
|
||||
fn style_text(&self, text: &str, theme: &ColorTheme) -> String {
|
||||
let mut style = text.stylize();
|
||||
|
||||
if matches!(self.heading_level, Some(1 | 2)) || self.strong > 0 {
|
||||
style = style.bold();
|
||||
}
|
||||
if self.emphasis > 0 {
|
||||
style = style.italic();
|
||||
}
|
||||
|
||||
if let Some(level) = self.heading_level {
|
||||
style = match level {
|
||||
1 => style.with(theme.heading),
|
||||
2 => style.white(),
|
||||
3 => style.with(Color::Blue),
|
||||
_ => style.with(Color::Grey),
|
||||
};
|
||||
} else if self.strong > 0 {
|
||||
style = style.with(theme.strong);
|
||||
} else if self.emphasis > 0 {
|
||||
style = style.with(theme.emphasis);
|
||||
}
|
||||
|
||||
if self.quote > 0 {
|
||||
style = style.with(theme.quote);
|
||||
}
|
||||
|
||||
format!("{style}")
|
||||
}
|
||||
|
||||
fn append_raw(&mut self, output: &mut String, text: &str) {
|
||||
if let Some(link) = self.link_stack.last_mut() {
|
||||
link.text.push_str(text);
|
||||
} else if let Some(table) = self.table.as_mut() {
|
||||
table.current_cell.push_str(text);
|
||||
} else {
|
||||
output.push_str(text);
|
||||
}
|
||||
}
|
||||
|
||||
fn append_styled(&mut self, output: &mut String, text: &str, theme: &ColorTheme) {
|
||||
let styled = self.style_text(text, theme);
|
||||
self.append_raw(output, &styled);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct TerminalRenderer {
|
||||
syntax_set: SyntaxSet,
|
||||
syntax_theme: Theme,
|
||||
color_theme: ColorTheme,
|
||||
}
|
||||
|
||||
impl Default for TerminalRenderer {
|
||||
fn default() -> Self {
|
||||
let syntax_set = SyntaxSet::load_defaults_newlines();
|
||||
let syntax_theme = ThemeSet::load_defaults()
|
||||
.themes
|
||||
.remove("base16-ocean.dark")
|
||||
.unwrap_or_default();
|
||||
Self {
|
||||
syntax_set,
|
||||
syntax_theme,
|
||||
color_theme: ColorTheme::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TerminalRenderer {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn color_theme(&self) -> &ColorTheme {
|
||||
&self.color_theme
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn render_markdown(&self, markdown: &str) -> String {
|
||||
let mut output = String::new();
|
||||
let mut state = RenderState::default();
|
||||
let mut code_language = String::new();
|
||||
let mut code_buffer = String::new();
|
||||
let mut in_code_block = false;
|
||||
|
||||
for event in Parser::new_ext(markdown, Options::all()) {
|
||||
self.render_event(
|
||||
event,
|
||||
&mut state,
|
||||
&mut output,
|
||||
&mut code_buffer,
|
||||
&mut code_language,
|
||||
&mut in_code_block,
|
||||
);
|
||||
}
|
||||
|
||||
output.trim_end().to_string()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn markdown_to_ansi(&self, markdown: &str) -> String {
|
||||
self.render_markdown(markdown)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_lines)]
|
||||
fn render_event(
|
||||
&self,
|
||||
event: Event<'_>,
|
||||
state: &mut RenderState,
|
||||
output: &mut String,
|
||||
code_buffer: &mut String,
|
||||
code_language: &mut String,
|
||||
in_code_block: &mut bool,
|
||||
) {
|
||||
match event {
|
||||
Event::Start(Tag::Heading { level, .. }) => {
|
||||
self.start_heading(state, level as u8, output);
|
||||
}
|
||||
Event::End(TagEnd::Paragraph) => output.push_str("\n\n"),
|
||||
Event::Start(Tag::BlockQuote(..)) => self.start_quote(state, output),
|
||||
Event::End(TagEnd::BlockQuote(..)) => {
|
||||
state.quote = state.quote.saturating_sub(1);
|
||||
output.push('\n');
|
||||
}
|
||||
Event::End(TagEnd::Heading(..)) => {
|
||||
state.heading_level = None;
|
||||
output.push_str("\n\n");
|
||||
}
|
||||
Event::End(TagEnd::Item) | Event::SoftBreak | Event::HardBreak => {
|
||||
state.append_raw(output, "\n");
|
||||
}
|
||||
Event::Start(Tag::List(first_item)) => {
|
||||
let kind = match first_item {
|
||||
Some(index) => ListKind::Ordered { next_index: index },
|
||||
None => ListKind::Unordered,
|
||||
};
|
||||
state.list_stack.push(kind);
|
||||
}
|
||||
Event::End(TagEnd::List(..)) => {
|
||||
state.list_stack.pop();
|
||||
output.push('\n');
|
||||
}
|
||||
Event::Start(Tag::Item) => Self::start_item(state, output),
|
||||
Event::Start(Tag::CodeBlock(kind)) => {
|
||||
*in_code_block = true;
|
||||
*code_language = match kind {
|
||||
CodeBlockKind::Indented => String::from("text"),
|
||||
CodeBlockKind::Fenced(lang) => lang.to_string(),
|
||||
};
|
||||
code_buffer.clear();
|
||||
self.start_code_block(code_language, output);
|
||||
}
|
||||
Event::End(TagEnd::CodeBlock) => {
|
||||
self.finish_code_block(code_buffer, code_language, output);
|
||||
*in_code_block = false;
|
||||
code_language.clear();
|
||||
code_buffer.clear();
|
||||
}
|
||||
Event::Start(Tag::Emphasis) => state.emphasis += 1,
|
||||
Event::End(TagEnd::Emphasis) => state.emphasis = state.emphasis.saturating_sub(1),
|
||||
Event::Start(Tag::Strong) => state.strong += 1,
|
||||
Event::End(TagEnd::Strong) => state.strong = state.strong.saturating_sub(1),
|
||||
Event::Code(code) => {
|
||||
let rendered =
|
||||
format!("{}", format!("`{code}`").with(self.color_theme.inline_code));
|
||||
state.append_raw(output, &rendered);
|
||||
}
|
||||
Event::Rule => output.push_str("---\n"),
|
||||
Event::Text(text) => {
|
||||
self.push_text(text.as_ref(), state, output, code_buffer, *in_code_block);
|
||||
}
|
||||
Event::Html(html) | Event::InlineHtml(html) => {
|
||||
state.append_raw(output, &html);
|
||||
}
|
||||
Event::FootnoteReference(reference) => {
|
||||
state.append_raw(output, &format!("[{reference}]"));
|
||||
}
|
||||
Event::TaskListMarker(done) => {
|
||||
state.append_raw(output, if done { "[x] " } else { "[ ] " });
|
||||
}
|
||||
Event::InlineMath(math) | Event::DisplayMath(math) => {
|
||||
state.append_raw(output, &math);
|
||||
}
|
||||
Event::Start(Tag::Link { dest_url, .. }) => {
|
||||
state.link_stack.push(LinkState {
|
||||
destination: dest_url.to_string(),
|
||||
text: String::new(),
|
||||
});
|
||||
}
|
||||
Event::End(TagEnd::Link) => {
|
||||
if let Some(link) = state.link_stack.pop() {
|
||||
let label = if link.text.is_empty() {
|
||||
link.destination.clone()
|
||||
} else {
|
||||
link.text
|
||||
};
|
||||
let rendered = format!(
|
||||
"{}",
|
||||
format!("[{label}]({})", link.destination)
|
||||
.underlined()
|
||||
.with(self.color_theme.link)
|
||||
);
|
||||
state.append_raw(output, &rendered);
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::Image { dest_url, .. }) => {
|
||||
let rendered = format!(
|
||||
"{}",
|
||||
format!("[image:{dest_url}]").with(self.color_theme.link)
|
||||
);
|
||||
state.append_raw(output, &rendered);
|
||||
}
|
||||
Event::Start(Tag::Table(..)) => state.table = Some(TableState::default()),
|
||||
Event::End(TagEnd::Table) => {
|
||||
if let Some(table) = state.table.take() {
|
||||
output.push_str(&self.render_table(&table));
|
||||
output.push_str("\n\n");
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::TableHead) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.in_head = true;
|
||||
}
|
||||
}
|
||||
Event::End(TagEnd::TableHead) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.finish_row();
|
||||
table.in_head = false;
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::TableRow) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.current_row.clear();
|
||||
table.current_cell.clear();
|
||||
}
|
||||
}
|
||||
Event::End(TagEnd::TableRow) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.finish_row();
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::TableCell) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.current_cell.clear();
|
||||
}
|
||||
}
|
||||
Event::End(TagEnd::TableCell) => {
|
||||
if let Some(table) = state.table.as_mut() {
|
||||
table.push_cell();
|
||||
}
|
||||
}
|
||||
Event::Start(Tag::Paragraph | Tag::MetadataBlock(..) | _)
|
||||
| Event::End(TagEnd::Image | TagEnd::MetadataBlock(..) | _) => {}
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::unused_self)]
|
||||
fn start_heading(&self, state: &mut RenderState, level: u8, output: &mut String) {
|
||||
state.heading_level = Some(level);
|
||||
if !output.is_empty() {
|
||||
output.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
fn start_quote(&self, state: &mut RenderState, output: &mut String) {
|
||||
state.quote += 1;
|
||||
let _ = write!(output, "{}", "│ ".with(self.color_theme.quote));
|
||||
}
|
||||
|
||||
fn start_item(state: &mut RenderState, output: &mut String) {
|
||||
let depth = state.list_stack.len().saturating_sub(1);
|
||||
output.push_str(&" ".repeat(depth));
|
||||
|
||||
let marker = match state.list_stack.last_mut() {
|
||||
Some(ListKind::Ordered { next_index }) => {
|
||||
let value = *next_index;
|
||||
*next_index += 1;
|
||||
format!("{value}. ")
|
||||
}
|
||||
_ => "• ".to_string(),
|
||||
};
|
||||
output.push_str(&marker);
|
||||
}
|
||||
|
||||
fn start_code_block(&self, code_language: &str, output: &mut String) {
|
||||
let label = if code_language.is_empty() {
|
||||
"code".to_string()
|
||||
} else {
|
||||
code_language.to_string()
|
||||
};
|
||||
let _ = writeln!(
|
||||
output,
|
||||
"{}",
|
||||
format!("╭─ {label}")
|
||||
.bold()
|
||||
.with(self.color_theme.code_block_border)
|
||||
);
|
||||
}
|
||||
|
||||
fn finish_code_block(&self, code_buffer: &str, code_language: &str, output: &mut String) {
|
||||
output.push_str(&self.highlight_code(code_buffer, code_language));
|
||||
let _ = write!(
|
||||
output,
|
||||
"{}",
|
||||
"╰─".bold().with(self.color_theme.code_block_border)
|
||||
);
|
||||
output.push_str("\n\n");
|
||||
}
|
||||
|
||||
fn push_text(
|
||||
&self,
|
||||
text: &str,
|
||||
state: &mut RenderState,
|
||||
output: &mut String,
|
||||
code_buffer: &mut String,
|
||||
in_code_block: bool,
|
||||
) {
|
||||
if in_code_block {
|
||||
code_buffer.push_str(text);
|
||||
} else {
|
||||
state.append_styled(output, text, &self.color_theme);
|
||||
}
|
||||
}
|
||||
|
||||
fn render_table(&self, table: &TableState) -> String {
|
||||
let mut rows = Vec::new();
|
||||
if !table.headers.is_empty() {
|
||||
rows.push(table.headers.clone());
|
||||
}
|
||||
rows.extend(table.rows.iter().cloned());
|
||||
|
||||
if rows.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
|
||||
let column_count = rows.iter().map(Vec::len).max().unwrap_or(0);
|
||||
let widths = (0..column_count)
|
||||
.map(|column| {
|
||||
rows.iter()
|
||||
.filter_map(|row| row.get(column))
|
||||
.map(|cell| visible_width(cell))
|
||||
.max()
|
||||
.unwrap_or(0)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let border = format!("{}", "│".with(self.color_theme.table_border));
|
||||
let separator = widths
|
||||
.iter()
|
||||
.map(|width| "─".repeat(*width + 2))
|
||||
.collect::<Vec<_>>()
|
||||
.join(&format!("{}", "┼".with(self.color_theme.table_border)));
|
||||
let separator = format!("{border}{separator}{border}");
|
||||
|
||||
let mut output = String::new();
|
||||
if !table.headers.is_empty() {
|
||||
output.push_str(&self.render_table_row(&table.headers, &widths, true));
|
||||
output.push('\n');
|
||||
output.push_str(&separator);
|
||||
if !table.rows.is_empty() {
|
||||
output.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
for (index, row) in table.rows.iter().enumerate() {
|
||||
output.push_str(&self.render_table_row(row, &widths, false));
|
||||
if index + 1 < table.rows.len() {
|
||||
output.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
fn render_table_row(&self, row: &[String], widths: &[usize], is_header: bool) -> String {
|
||||
let border = format!("{}", "│".with(self.color_theme.table_border));
|
||||
let mut line = String::new();
|
||||
line.push_str(&border);
|
||||
|
||||
for (index, width) in widths.iter().enumerate() {
|
||||
let cell = row.get(index).map_or("", String::as_str);
|
||||
line.push(' ');
|
||||
if is_header {
|
||||
let _ = write!(line, "{}", cell.bold().with(self.color_theme.heading));
|
||||
} else {
|
||||
line.push_str(cell);
|
||||
}
|
||||
let padding = width.saturating_sub(visible_width(cell));
|
||||
line.push_str(&" ".repeat(padding + 1));
|
||||
line.push_str(&border);
|
||||
}
|
||||
|
||||
line
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn highlight_code(&self, code: &str, language: &str) -> String {
|
||||
let syntax = self
|
||||
.syntax_set
|
||||
.find_syntax_by_token(language)
|
||||
.unwrap_or_else(|| self.syntax_set.find_syntax_plain_text());
|
||||
let mut syntax_highlighter = HighlightLines::new(syntax, &self.syntax_theme);
|
||||
let mut colored_output = String::new();
|
||||
|
||||
for line in LinesWithEndings::from(code) {
|
||||
match syntax_highlighter.highlight_line(line, &self.syntax_set) {
|
||||
Ok(ranges) => {
|
||||
let escaped = as_24_bit_terminal_escaped(&ranges[..], false);
|
||||
colored_output.push_str(&apply_code_block_background(&escaped));
|
||||
}
|
||||
Err(_) => colored_output.push_str(&apply_code_block_background(line)),
|
||||
}
|
||||
}
|
||||
|
||||
colored_output
|
||||
}
|
||||
|
||||
pub fn stream_markdown(&self, markdown: &str, out: &mut impl Write) -> io::Result<()> {
|
||||
let rendered_markdown = self.markdown_to_ansi(markdown);
|
||||
write!(out, "{rendered_markdown}")?;
|
||||
if !rendered_markdown.ends_with('\n') {
|
||||
writeln!(out)?;
|
||||
}
|
||||
out.flush()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Clone, PartialEq, Eq)]
|
||||
pub struct MarkdownStreamState {
|
||||
pending: String,
|
||||
}
|
||||
|
||||
impl MarkdownStreamState {
|
||||
#[must_use]
|
||||
pub fn push(&mut self, renderer: &TerminalRenderer, delta: &str) -> Option<String> {
|
||||
self.pending.push_str(delta);
|
||||
let split = find_stream_safe_boundary(&self.pending)?;
|
||||
let ready = self.pending[..split].to_string();
|
||||
self.pending.drain(..split);
|
||||
Some(renderer.markdown_to_ansi(&ready))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn flush(&mut self, renderer: &TerminalRenderer) -> Option<String> {
|
||||
if self.pending.trim().is_empty() {
|
||||
self.pending.clear();
|
||||
None
|
||||
} else {
|
||||
let pending = std::mem::take(&mut self.pending);
|
||||
Some(renderer.markdown_to_ansi(&pending))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_code_block_background(line: &str) -> String {
|
||||
let trimmed = line.trim_end_matches('\n');
|
||||
let trailing_newline = if trimmed.len() == line.len() {
|
||||
""
|
||||
} else {
|
||||
"\n"
|
||||
};
|
||||
let with_background = trimmed.replace("\u{1b}[0m", "\u{1b}[0;48;5;236m");
|
||||
format!("\u{1b}[48;5;236m{with_background}\u{1b}[0m{trailing_newline}")
|
||||
}
|
||||
|
||||
fn find_stream_safe_boundary(markdown: &str) -> Option<usize> {
|
||||
let mut in_fence = false;
|
||||
let mut last_boundary = None;
|
||||
|
||||
for (offset, line) in markdown.split_inclusive('\n').scan(0usize, |cursor, line| {
|
||||
let start = *cursor;
|
||||
*cursor += line.len();
|
||||
Some((start, line))
|
||||
}) {
|
||||
let trimmed = line.trim_start();
|
||||
if trimmed.starts_with("```") || trimmed.starts_with("~~~") {
|
||||
in_fence = !in_fence;
|
||||
if !in_fence {
|
||||
last_boundary = Some(offset + line.len());
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if in_fence {
|
||||
continue;
|
||||
}
|
||||
|
||||
if trimmed.is_empty() {
|
||||
last_boundary = Some(offset + line.len());
|
||||
}
|
||||
}
|
||||
|
||||
last_boundary
|
||||
}
|
||||
|
||||
fn visible_width(input: &str) -> usize {
|
||||
strip_ansi(input).chars().count()
|
||||
}
|
||||
|
||||
fn strip_ansi(input: &str) -> String {
|
||||
let mut output = String::new();
|
||||
let mut chars = input.chars().peekable();
|
||||
|
||||
while let Some(ch) = chars.next() {
|
||||
if ch == '\u{1b}' {
|
||||
if chars.peek() == Some(&'[') {
|
||||
chars.next();
|
||||
for next in chars.by_ref() {
|
||||
if next.is_ascii_alphabetic() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
output.push(ch);
|
||||
}
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{strip_ansi, MarkdownStreamState, Spinner, TerminalRenderer};
|
||||
|
||||
#[test]
|
||||
fn renders_markdown_with_styling_and_lists() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output = terminal_renderer
|
||||
.render_markdown("# Heading\n\nThis is **bold** and *italic*.\n\n- item\n\n`code`");
|
||||
|
||||
assert!(markdown_output.contains("Heading"));
|
||||
assert!(markdown_output.contains("• item"));
|
||||
assert!(markdown_output.contains("code"));
|
||||
assert!(markdown_output.contains('\u{1b}'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_links_as_colored_markdown_labels() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output =
|
||||
terminal_renderer.render_markdown("See [Claw](https://example.com/docs) now.");
|
||||
let plain_text = strip_ansi(&markdown_output);
|
||||
|
||||
assert!(plain_text.contains("[Claw](https://example.com/docs)"));
|
||||
assert!(markdown_output.contains('\u{1b}'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn highlights_fenced_code_blocks() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output =
|
||||
terminal_renderer.markdown_to_ansi("```rust\nfn hi() { println!(\"hi\"); }\n```");
|
||||
let plain_text = strip_ansi(&markdown_output);
|
||||
|
||||
assert!(plain_text.contains("╭─ rust"));
|
||||
assert!(plain_text.contains("fn hi"));
|
||||
assert!(markdown_output.contains('\u{1b}'));
|
||||
assert!(markdown_output.contains("[48;5;236m"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_ordered_and_nested_lists() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output =
|
||||
terminal_renderer.render_markdown("1. first\n2. second\n - nested\n - child");
|
||||
let plain_text = strip_ansi(&markdown_output);
|
||||
|
||||
assert!(plain_text.contains("1. first"));
|
||||
assert!(plain_text.contains("2. second"));
|
||||
assert!(plain_text.contains(" • nested"));
|
||||
assert!(plain_text.contains(" • child"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn renders_tables_with_alignment() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let markdown_output = terminal_renderer
|
||||
.render_markdown("| Name | Value |\n| ---- | ----- |\n| alpha | 1 |\n| beta | 22 |");
|
||||
let plain_text = strip_ansi(&markdown_output);
|
||||
let lines = plain_text.lines().collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(lines[0], "│ Name │ Value │");
|
||||
assert_eq!(lines[1], "│───────┼───────│");
|
||||
assert_eq!(lines[2], "│ alpha │ 1 │");
|
||||
assert_eq!(lines[3], "│ beta │ 22 │");
|
||||
assert!(markdown_output.contains('\u{1b}'));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streaming_state_waits_for_complete_blocks() {
|
||||
let renderer = TerminalRenderer::new();
|
||||
let mut state = MarkdownStreamState::default();
|
||||
|
||||
assert_eq!(state.push(&renderer, "# Heading"), None);
|
||||
let flushed = state
|
||||
.push(&renderer, "\n\nParagraph\n\n")
|
||||
.expect("completed block");
|
||||
let plain_text = strip_ansi(&flushed);
|
||||
assert!(plain_text.contains("Heading"));
|
||||
assert!(plain_text.contains("Paragraph"));
|
||||
|
||||
assert_eq!(state.push(&renderer, "```rust\nfn main() {}\n"), None);
|
||||
let code = state
|
||||
.push(&renderer, "```\n")
|
||||
.expect("closed code fence flushes");
|
||||
assert!(strip_ansi(&code).contains("fn main()"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn spinner_advances_frames() {
|
||||
let terminal_renderer = TerminalRenderer::new();
|
||||
let mut spinner = Spinner::new();
|
||||
let mut out = Vec::new();
|
||||
spinner
|
||||
.tick("Working", terminal_renderer.color_theme(), &mut out)
|
||||
.expect("tick succeeds");
|
||||
spinner
|
||||
.tick("Working", terminal_renderer.color_theme(), &mut out)
|
||||
.expect("tick succeeds");
|
||||
|
||||
let output = String::from_utf8_lossy(&out);
|
||||
assert!(output.contains("Working"));
|
||||
}
|
||||
}
|
||||
@@ -9,4 +9,6 @@ publish.workspace = true
|
||||
workspace = true
|
||||
|
||||
[dependencies]
|
||||
plugins = { path = "../plugins" }
|
||||
runtime = { path = "../runtime" }
|
||||
serde_json.workspace = true
|
||||
@@ -0,0 +1,55 @@
|
||||
# 命令模块 (commands)
|
||||
|
||||
本模块负责定义和管理 Claw 交互界面中使用的“斜杠命令”(Slash Commands),并提供相关的解析和执行逻辑。
|
||||
|
||||
## 概览
|
||||
|
||||
`commands` 模块的主要职责包括:
|
||||
- 定义所有可用的斜杠命令及其元数据(别名、说明、类别等)。
|
||||
- 提供命令注册表 (`CommandRegistry`),用于在 CLI 中发现和分发命令。
|
||||
- 实现复杂的管理命令,如插件管理 (`/plugins`)、智能体查看 (`/agents`) 和技能查看 (`/skills`)。
|
||||
- 提供命令建议功能,支持基于编辑距离 (Levenshtein distance) 的模糊匹配。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **斜杠命令规范 (SlashCommandSpec)**:每个命令都包含详尽的元数据,包括所属类别(核心、工作区、会话、Git、自动化)以及是否支持在恢复会话时执行。
|
||||
- **命令分类**:
|
||||
- **核心 (Core)**:`/help`, `/status`, `/model`, `/permissions`, `/cost` 等。
|
||||
- **工作区 (Workspace)**:`/config`, `/memory`, `/diff`, `/teleport` 等。
|
||||
- **会话 (Session)**:`/clear`, `/resume`, `/export`, `/session` 等。
|
||||
- **Git 交互**:`/branch`, `/commit`, `/pr`, `/issue` 等。
|
||||
- **自动化 (Automation)**:`/plugins`, `/agents`, `/skills`, `/ultraplan` 等。
|
||||
- **模糊匹配与建议**:当用户输入错误的命令时,系统会自动推荐最接近的合法命令。
|
||||
- **插件集成**:`/plugins` 命令允许用户动态安装、启用、禁用或卸载插件,并能通知运行时重新加载环境。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`lib.rs`**: 包含了绝大部分逻辑。
|
||||
- **`SlashCommand` 枚举**: 定义了所有命令的强类型表示。
|
||||
- **`SlashCommandSpec` 结构体**: 存储命令的静态配置信息。
|
||||
- **`handle_plugins_slash_command`**: 处理复杂的插件管理工作流。
|
||||
- **`suggest_slash_commands`**: 实现基于 Levenshtein 距离的建议算法。
|
||||
|
||||
### 工作流程
|
||||
|
||||
1. 用户在 REPL 中输入以 `/` 开头的字符串。
|
||||
2. `claw-cli` 调用 `SlashCommand::parse` 进行解析。
|
||||
3. 解析后的命令被分发到相应的处理器。
|
||||
4. 处理结果(通常包含要显示给用户的消息,以及可选的会话更新或运行时重新加载请求)返回给 CLI。
|
||||
|
||||
## 使用示例 (内部)
|
||||
|
||||
```rust
|
||||
use commands::{SlashCommand, suggest_slash_commands};
|
||||
|
||||
// 解析命令
|
||||
if let Some(cmd) = SlashCommand::parse("/model sonnet") {
|
||||
// 处理模型切换逻辑
|
||||
}
|
||||
|
||||
// 获取建议
|
||||
let suggestions = suggest_slash_commands("hpel", 3);
|
||||
// 返回 ["/help"]
|
||||
```
|
||||
@@ -0,0 +1,48 @@
|
||||
# 兼容性测试套件模块 (compat-harness)
|
||||
|
||||
本模块提供了一套工具,专门用于分析和提取上游引用实现(如原始的 `claude-code` TypeScript 源码)中的元数据,以确保 Rust 版本的实现与其保持功能兼容。
|
||||
|
||||
## 概览
|
||||
|
||||
`compat-harness` 的主要职责是:
|
||||
- 定位上游仓库的源码路径。
|
||||
- 从 TypeScript 源码文件中提取命令 (`commands`)、工具 (`tools`) 和启动阶段 (`bootstrap phases`) 的定义。
|
||||
- 自动生成功能清单 (`ExtractedManifest`),供运行时或测试使用,以验证 Rust 版本的覆盖率。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **上游路径解析 (UpstreamPaths)**:能够自动识别多种常见的上游仓库目录结构,并支持通过环境变量 `CLAUDE_CODE_UPSTREAM` 进行覆盖。
|
||||
- **静态代码分析**:通过解析 TypeScript 源码,识别特定的代码模式(如 `export const INTERNAL_ONLY_COMMANDS` 或基于 `feature()` 的功能开关)。
|
||||
- **清单提取 (Manifest Extraction)**:
|
||||
- **命令提取**:识别内置命令、仅限内部使用的命令以及受功能开关控制的命令。
|
||||
- **工具提取**:识别基础工具和条件加载的工具。
|
||||
- **启动计划提取**:分析 CLI 入口文件,重建启动时的各个阶段(如 `FastPathVersion`, `MainRuntime` 等)。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`lib.rs`**: 包含了核心的提取逻辑。
|
||||
- **`UpstreamPaths` 结构体**: 封装了寻找 `commands.ts`、`tools.ts` 和 `cli.tsx` 的逻辑。
|
||||
- **`extract_commands` & `extract_tools`**: 使用字符串解析技术,识别 TypeScript 的 `import` 和赋值操作,提取符号名称。
|
||||
- **`extract_bootstrap_plan`**: 搜索特定的标志性字符串(如 `--version` 或 `daemon-worker`),从而推断出上游程序的启动流程。
|
||||
|
||||
### 工作流程
|
||||
|
||||
1. 模块根据预设路径或环境变量寻找上游 `claude-code` 仓库。
|
||||
2. 读取关键的 `.ts` 或 `.tsx` 文件内容。
|
||||
3. 执行正则表达式风格的行解析,提取出所有定义的命令和工具名称。
|
||||
4. 将提取结果组织成 `ExtractedManifest` 对象。
|
||||
|
||||
## 使用示例 (内部测试)
|
||||
|
||||
```rust
|
||||
use compat_harness::{UpstreamPaths, extract_manifest};
|
||||
|
||||
// 指定工作区目录,自动寻找上游路径
|
||||
let paths = UpstreamPaths::from_workspace_dir("path/to/workspace");
|
||||
// 提取功能清单
|
||||
if let Ok(manifest) = extract_manifest(&paths) {
|
||||
println!("上游发现 {} 个工具", manifest.tools.entries().len());
|
||||
}
|
||||
```
|
||||
@@ -70,16 +70,12 @@ fn upstream_repo_candidates(primary_repo_root: &Path) -> Vec<PathBuf> {
|
||||
}
|
||||
|
||||
for ancestor in primary_repo_root.ancestors().take(4) {
|
||||
candidates.push(ancestor.join("claude-code"));
|
||||
candidates.push(ancestor.join("claw-code"));
|
||||
candidates.push(ancestor.join("clawd-code"));
|
||||
}
|
||||
|
||||
candidates.push(
|
||||
primary_repo_root
|
||||
.join("reference-source")
|
||||
.join("claude-code"),
|
||||
);
|
||||
candidates.push(primary_repo_root.join("vendor").join("claude-code"));
|
||||
candidates.push(primary_repo_root.join("reference-source").join("claw-code"));
|
||||
candidates.push(primary_repo_root.join("vendor").join("claw-code"));
|
||||
|
||||
let mut deduped = Vec::new();
|
||||
for candidate in candidates {
|
||||
@@ -1,18 +1,16 @@
|
||||
[package]
|
||||
name = "runtime"
|
||||
name = "lsp"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
publish.workspace = true
|
||||
|
||||
[dependencies]
|
||||
sha2 = "0.10"
|
||||
glob = "0.3"
|
||||
regex = "1"
|
||||
lsp-types.workspace = true
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
tokio = { version = "1", features = ["io-util", "macros", "process", "rt", "rt-multi-thread", "time"] }
|
||||
walkdir = "2"
|
||||
serde_json.workspace = true
|
||||
tokio = { version = "1", features = ["io-util", "macros", "process", "rt", "rt-multi-thread", "sync", "time"] }
|
||||
url = "2"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,56 @@
|
||||
# LSP 模块 (lsp)
|
||||
|
||||
本模块实现了语言服务协议 (Language Server Protocol, LSP) 的客户端功能,允许系统通过集成的编程语言服务器获取代码的语义信息、错误诊断和符号导航。
|
||||
|
||||
## 概览
|
||||
|
||||
`lsp` 模块的主要职责是:
|
||||
- 管理多个 LSP 服务器的生命周期(启动、初始化、关闭)。
|
||||
- 与服务器进行异步 JSON-RPC 通信。
|
||||
- 提供跨语言的代码智能功能,如:
|
||||
- **转到定义 (Go to Definition)**
|
||||
- **查找引用 (Find References)**
|
||||
- **工作区诊断 (Workspace Diagnostics)**
|
||||
- 为 AI 提示词 (Prompt) 提供上下文增强,将代码中的实时错误和符号关系反馈给 LLM。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **LspManager**: 核心管理类,负责协调不同语言的服务器配置和文档状态。
|
||||
- **上下文增强 (Context Enrichment)**:定义了 `LspContextEnrichment` 结构,能够将复杂的 LSP 响应(如诊断信息和定义)转换为易于 AI 理解的 Markdown 格式。
|
||||
- **多服务器支持**:支持根据文件扩展名将请求路由到不同的语言服务器(如 `rust-analyzer`, `pyright` 等)。
|
||||
- **同步机制**:处理文档的 `didOpen`、`didChange` 和 `didSave` 消息,确保服务器拥有最新的代码视图。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`manager.rs`**: 实现了 `LspManager`。它维护一个服务器池,并提供高层 API 来执行跨服务器的请求。
|
||||
- **`client.rs`**: 实现底层的 LSP 客户端逻辑,处理基于 `tokio` 的异步 I/O 和 JSON-RPC 消息的分帧与解析。
|
||||
- **`types.rs`**: 定义了本模块使用的专用数据类型,并对 `lsp-types` 库中的类型进行了简化和包装,以便于内部使用。
|
||||
- **`error.rs`**: 定义了 LSP 相关的错误处理。
|
||||
|
||||
### 工作流程
|
||||
|
||||
1. 系统根据配置初始化 `LspManager`。
|
||||
2. 当打开一个文件时,`LspManager` 启动相应的服务器并发送 `initialize` 请求。
|
||||
3. `LspManager` 跟踪文档的打开状态,并在内容变化时同步到服务器。
|
||||
4. 当需要对某个符号进行分析时,调用 `go_to_definition` 等方法,模块负责发送请求并解析返回的 `Location`。
|
||||
5. 诊断信息异步通过 `textDocument/publishDiagnostics` 通知到达,模块会缓存这些信息供后续查询。
|
||||
|
||||
## 使用示例 (内部)
|
||||
|
||||
```rust
|
||||
use lsp::{LspManager, LspServerConfig};
|
||||
|
||||
// 配置并初始化管理器
|
||||
let configs = vec![LspServerConfig {
|
||||
name: "rust-analyzer".to_string(),
|
||||
command: "rust-analyzer".to_string(),
|
||||
..Default::default()
|
||||
}];
|
||||
let manager = LspManager::new(configs)?;
|
||||
|
||||
// 获取某个位置的上下文增强信息
|
||||
let enrichment = manager.context_enrichment(&file_path, position).await?;
|
||||
println!("{}", enrichment.render_prompt_section());
|
||||
```
|
||||
@@ -0,0 +1,465 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicI64, Ordering};
|
||||
|
||||
use lsp_types::{
|
||||
Diagnostic, GotoDefinitionResponse, Location, LocationLink, Position, PublishDiagnosticsParams,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter};
|
||||
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
|
||||
use tokio::sync::{oneshot, Mutex};
|
||||
|
||||
use crate::error::LspError;
|
||||
use crate::types::{LspServerConfig, SymbolLocation};
|
||||
|
||||
type PendingRequestMap = BTreeMap<i64, oneshot::Sender<Result<Value, LspError>>>;
|
||||
|
||||
pub(crate) struct LspClient {
|
||||
config: LspServerConfig,
|
||||
writer: Mutex<BufWriter<ChildStdin>>,
|
||||
child: Mutex<Child>,
|
||||
pending_requests: Arc<Mutex<PendingRequestMap>>,
|
||||
diagnostics: Arc<Mutex<BTreeMap<String, Vec<Diagnostic>>>>,
|
||||
open_documents: Mutex<BTreeMap<PathBuf, i32>>,
|
||||
next_request_id: AtomicI64,
|
||||
}
|
||||
|
||||
impl LspClient {
|
||||
pub(crate) async fn connect(config: LspServerConfig) -> Result<Self, LspError> {
|
||||
let mut command = Command::new(&config.command);
|
||||
command
|
||||
.args(&config.args)
|
||||
.current_dir(&config.workspace_root)
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.envs(config.env.clone());
|
||||
|
||||
let mut child = command.spawn()?;
|
||||
let stdin = child
|
||||
.stdin
|
||||
.take()
|
||||
.ok_or_else(|| LspError::Protocol("missing LSP stdin pipe".to_string()))?;
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| LspError::Protocol("missing LSP stdout pipe".to_string()))?;
|
||||
let stderr = child.stderr.take();
|
||||
|
||||
let client = Self {
|
||||
config,
|
||||
writer: Mutex::new(BufWriter::new(stdin)),
|
||||
child: Mutex::new(child),
|
||||
pending_requests: Arc::new(Mutex::new(BTreeMap::new())),
|
||||
diagnostics: Arc::new(Mutex::new(BTreeMap::new())),
|
||||
open_documents: Mutex::new(BTreeMap::new()),
|
||||
next_request_id: AtomicI64::new(1),
|
||||
};
|
||||
|
||||
client.spawn_reader(stdout);
|
||||
if let Some(stderr) = stderr {
|
||||
Self::spawn_stderr_drain(stderr);
|
||||
}
|
||||
client.initialize().await?;
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
pub(crate) async fn ensure_document_open(&self, path: &Path) -> Result<(), LspError> {
|
||||
if self.is_document_open(path).await {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let contents = std::fs::read_to_string(path)?;
|
||||
self.open_document(path, &contents).await
|
||||
}
|
||||
|
||||
pub(crate) async fn open_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
|
||||
let uri = file_url(path)?;
|
||||
let language_id = self
|
||||
.config
|
||||
.language_id_for(path)
|
||||
.ok_or_else(|| LspError::UnsupportedDocument(path.to_path_buf()))?;
|
||||
|
||||
self.notify(
|
||||
"textDocument/didOpen",
|
||||
json!({
|
||||
"textDocument": {
|
||||
"uri": uri,
|
||||
"languageId": language_id,
|
||||
"version": 1,
|
||||
"text": text,
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
self.open_documents
|
||||
.lock()
|
||||
.await
|
||||
.insert(path.to_path_buf(), 1);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn change_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
|
||||
if !self.is_document_open(path).await {
|
||||
return self.open_document(path, text).await;
|
||||
}
|
||||
|
||||
let uri = file_url(path)?;
|
||||
let next_version = {
|
||||
let mut open_documents = self.open_documents.lock().await;
|
||||
let version = open_documents
|
||||
.entry(path.to_path_buf())
|
||||
.and_modify(|value| *value += 1)
|
||||
.or_insert(1);
|
||||
*version
|
||||
};
|
||||
|
||||
self.notify(
|
||||
"textDocument/didChange",
|
||||
json!({
|
||||
"textDocument": {
|
||||
"uri": uri,
|
||||
"version": next_version,
|
||||
},
|
||||
"contentChanges": [{
|
||||
"text": text,
|
||||
}],
|
||||
}),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn save_document(&self, path: &Path) -> Result<(), LspError> {
|
||||
if !self.is_document_open(path).await {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
self.notify(
|
||||
"textDocument/didSave",
|
||||
json!({
|
||||
"textDocument": {
|
||||
"uri": file_url(path)?,
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn close_document(&self, path: &Path) -> Result<(), LspError> {
|
||||
if !self.is_document_open(path).await {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
self.notify(
|
||||
"textDocument/didClose",
|
||||
json!({
|
||||
"textDocument": {
|
||||
"uri": file_url(path)?,
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
self.open_documents.lock().await.remove(path);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) async fn is_document_open(&self, path: &Path) -> bool {
|
||||
self.open_documents.lock().await.contains_key(path)
|
||||
}
|
||||
|
||||
pub(crate) async fn go_to_definition(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
) -> Result<Vec<SymbolLocation>, LspError> {
|
||||
self.ensure_document_open(path).await?;
|
||||
let response = self
|
||||
.request::<Option<GotoDefinitionResponse>>(
|
||||
"textDocument/definition",
|
||||
json!({
|
||||
"textDocument": { "uri": file_url(path)? },
|
||||
"position": position,
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(match response {
|
||||
Some(GotoDefinitionResponse::Scalar(location)) => {
|
||||
location_to_symbol_locations(vec![location])
|
||||
}
|
||||
Some(GotoDefinitionResponse::Array(locations)) => location_to_symbol_locations(locations),
|
||||
Some(GotoDefinitionResponse::Link(links)) => location_links_to_symbol_locations(links),
|
||||
None => Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn find_references(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
include_declaration: bool,
|
||||
) -> Result<Vec<SymbolLocation>, LspError> {
|
||||
self.ensure_document_open(path).await?;
|
||||
let response = self
|
||||
.request::<Option<Vec<Location>>>(
|
||||
"textDocument/references",
|
||||
json!({
|
||||
"textDocument": { "uri": file_url(path)? },
|
||||
"position": position,
|
||||
"context": {
|
||||
"includeDeclaration": include_declaration,
|
||||
},
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(location_to_symbol_locations(response.unwrap_or_default()))
|
||||
}
|
||||
|
||||
pub(crate) async fn diagnostics_snapshot(&self) -> BTreeMap<String, Vec<Diagnostic>> {
|
||||
self.diagnostics.lock().await.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn shutdown(&self) -> Result<(), LspError> {
|
||||
let _ = self.request::<Value>("shutdown", json!({})).await;
|
||||
let _ = self.notify("exit", Value::Null).await;
|
||||
|
||||
let mut child = self.child.lock().await;
|
||||
if child.kill().await.is_err() {
|
||||
let _ = child.wait().await;
|
||||
return Ok(());
|
||||
}
|
||||
let _ = child.wait().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn spawn_reader(&self, stdout: ChildStdout) {
|
||||
let diagnostics = &self.diagnostics;
|
||||
let pending_requests = &self.pending_requests;
|
||||
|
||||
let diagnostics = diagnostics.clone();
|
||||
let pending_requests = pending_requests.clone();
|
||||
tokio::spawn(async move {
|
||||
let mut reader = BufReader::new(stdout);
|
||||
let result = async {
|
||||
while let Some(message) = read_message(&mut reader).await? {
|
||||
if let Some(id) = message.get("id").and_then(Value::as_i64) {
|
||||
let response = if let Some(error) = message.get("error") {
|
||||
Err(LspError::Protocol(error.to_string()))
|
||||
} else {
|
||||
Ok(message.get("result").cloned().unwrap_or(Value::Null))
|
||||
};
|
||||
|
||||
if let Some(sender) = pending_requests.lock().await.remove(&id) {
|
||||
let _ = sender.send(response);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
let Some(method) = message.get("method").and_then(Value::as_str) else {
|
||||
continue;
|
||||
};
|
||||
if method != "textDocument/publishDiagnostics" {
|
||||
continue;
|
||||
}
|
||||
|
||||
let params = message.get("params").cloned().unwrap_or(Value::Null);
|
||||
let notification = serde_json::from_value::<PublishDiagnosticsParams>(params)?;
|
||||
let mut diagnostics_map = diagnostics.lock().await;
|
||||
if notification.diagnostics.is_empty() {
|
||||
diagnostics_map.remove(¬ification.uri.to_string());
|
||||
} else {
|
||||
diagnostics_map.insert(notification.uri.to_string(), notification.diagnostics);
|
||||
}
|
||||
}
|
||||
Ok::<(), LspError>(())
|
||||
}
|
||||
.await;
|
||||
|
||||
if let Err(error) = result {
|
||||
let mut pending = pending_requests.lock().await;
|
||||
let drained = pending
|
||||
.keys()
|
||||
.copied()
|
||||
.collect::<Vec<_>>();
|
||||
for id in drained {
|
||||
if let Some(sender) = pending.remove(&id) {
|
||||
let _ = sender.send(Err(LspError::Protocol(error.to_string())));
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn spawn_stderr_drain<R>(stderr: R)
|
||||
where
|
||||
R: AsyncRead + Unpin + Send + 'static,
|
||||
{
|
||||
tokio::spawn(async move {
|
||||
let mut reader = BufReader::new(stderr);
|
||||
let mut sink = Vec::new();
|
||||
let _ = reader.read_to_end(&mut sink).await;
|
||||
});
|
||||
}
|
||||
|
||||
async fn initialize(&self) -> Result<(), LspError> {
|
||||
let workspace_uri = file_url(&self.config.workspace_root)?;
|
||||
let _ = self
|
||||
.request::<Value>(
|
||||
"initialize",
|
||||
json!({
|
||||
"processId": std::process::id(),
|
||||
"rootUri": workspace_uri,
|
||||
"rootPath": self.config.workspace_root,
|
||||
"workspaceFolders": [{
|
||||
"uri": workspace_uri,
|
||||
"name": self.config.name,
|
||||
}],
|
||||
"initializationOptions": self.config.initialization_options.clone().unwrap_or(Value::Null),
|
||||
"capabilities": {
|
||||
"textDocument": {
|
||||
"publishDiagnostics": {
|
||||
"relatedInformation": true,
|
||||
},
|
||||
"definition": {
|
||||
"linkSupport": true,
|
||||
},
|
||||
"references": {}
|
||||
},
|
||||
"workspace": {
|
||||
"configuration": false,
|
||||
"workspaceFolders": true,
|
||||
},
|
||||
"general": {
|
||||
"positionEncodings": ["utf-16"],
|
||||
}
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await?;
|
||||
self.notify("initialized", json!({})).await
|
||||
}
|
||||
|
||||
async fn request<T>(&self, method: &str, params: Value) -> Result<T, LspError>
|
||||
where
|
||||
T: for<'de> serde::Deserialize<'de>,
|
||||
{
|
||||
let id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
|
||||
let (sender, receiver) = oneshot::channel();
|
||||
self.pending_requests.lock().await.insert(id, sender);
|
||||
|
||||
if let Err(error) = self
|
||||
.send_message(&json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": id,
|
||||
"method": method,
|
||||
"params": params,
|
||||
}))
|
||||
.await
|
||||
{
|
||||
self.pending_requests.lock().await.remove(&id);
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
let response = receiver
|
||||
.await
|
||||
.map_err(|_| LspError::Protocol(format!("request channel closed for {method}")))??;
|
||||
Ok(serde_json::from_value(response)?)
|
||||
}
|
||||
|
||||
async fn notify(&self, method: &str, params: Value) -> Result<(), LspError> {
|
||||
self.send_message(&json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": method,
|
||||
"params": params,
|
||||
}))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_message(&self, payload: &Value) -> Result<(), LspError> {
|
||||
let body = serde_json::to_vec(payload)?;
|
||||
let mut writer = self.writer.lock().await;
|
||||
writer
|
||||
.write_all(format!("Content-Length: {}\r\n\r\n", body.len()).as_bytes())
|
||||
.await?;
|
||||
writer.write_all(&body).await?;
|
||||
writer.flush().await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_message<R>(reader: &mut BufReader<R>) -> Result<Option<Value>, LspError>
|
||||
where
|
||||
R: AsyncRead + Unpin,
|
||||
{
|
||||
let mut content_length = None;
|
||||
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
let read = reader.read_line(&mut line).await?;
|
||||
if read == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
if line == "\r\n" {
|
||||
break;
|
||||
}
|
||||
|
||||
let trimmed = line.trim_end_matches(['\r', '\n']);
|
||||
if let Some((name, value)) = trimmed.split_once(':') {
|
||||
if name.eq_ignore_ascii_case("Content-Length") {
|
||||
let value = value.trim().to_string();
|
||||
content_length = Some(
|
||||
value
|
||||
.parse::<usize>()
|
||||
.map_err(|_| LspError::InvalidContentLength(value.clone()))?,
|
||||
);
|
||||
}
|
||||
} else {
|
||||
return Err(LspError::InvalidHeader(trimmed.to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
let content_length = content_length.ok_or(LspError::MissingContentLength)?;
|
||||
let mut body = vec![0_u8; content_length];
|
||||
reader.read_exact(&mut body).await?;
|
||||
Ok(Some(serde_json::from_slice(&body)?))
|
||||
}
|
||||
|
||||
fn file_url(path: &Path) -> Result<String, LspError> {
|
||||
url::Url::from_file_path(path)
|
||||
.map(|url| url.to_string())
|
||||
.map_err(|()| LspError::PathToUrl(path.to_path_buf()))
|
||||
}
|
||||
|
||||
fn location_to_symbol_locations(locations: Vec<Location>) -> Vec<SymbolLocation> {
|
||||
locations
|
||||
.into_iter()
|
||||
.filter_map(|location| {
|
||||
uri_to_path(&location.uri.to_string()).map(|path| SymbolLocation {
|
||||
path,
|
||||
range: location.range,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn location_links_to_symbol_locations(links: Vec<LocationLink>) -> Vec<SymbolLocation> {
|
||||
links.into_iter()
|
||||
.filter_map(|link| {
|
||||
uri_to_path(&link.target_uri.to_string()).map(|path| SymbolLocation {
|
||||
path,
|
||||
range: link.target_selection_range,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn uri_to_path(uri: &str) -> Option<PathBuf> {
|
||||
url::Url::parse(uri).ok()?.to_file_path().ok()
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum LspError {
|
||||
Io(std::io::Error),
|
||||
Json(serde_json::Error),
|
||||
InvalidHeader(String),
|
||||
MissingContentLength,
|
||||
InvalidContentLength(String),
|
||||
UnsupportedDocument(PathBuf),
|
||||
UnknownServer(String),
|
||||
DuplicateExtension {
|
||||
extension: String,
|
||||
existing_server: String,
|
||||
new_server: String,
|
||||
},
|
||||
PathToUrl(PathBuf),
|
||||
Protocol(String),
|
||||
}
|
||||
|
||||
impl Display for LspError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Io(error) => write!(f, "{error}"),
|
||||
Self::Json(error) => write!(f, "{error}"),
|
||||
Self::InvalidHeader(header) => write!(f, "invalid LSP header: {header}"),
|
||||
Self::MissingContentLength => write!(f, "missing LSP Content-Length header"),
|
||||
Self::InvalidContentLength(value) => {
|
||||
write!(f, "invalid LSP Content-Length value: {value}")
|
||||
}
|
||||
Self::UnsupportedDocument(path) => {
|
||||
write!(f, "no LSP server configured for {}", path.display())
|
||||
}
|
||||
Self::UnknownServer(name) => write!(f, "unknown LSP server: {name}"),
|
||||
Self::DuplicateExtension {
|
||||
extension,
|
||||
existing_server,
|
||||
new_server,
|
||||
} => write!(
|
||||
f,
|
||||
"duplicate LSP extension mapping for {extension}: {existing_server} and {new_server}"
|
||||
),
|
||||
Self::PathToUrl(path) => write!(f, "failed to convert path to file URL: {}", path.display()),
|
||||
Self::Protocol(message) => write!(f, "LSP protocol error: {message}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for LspError {}
|
||||
|
||||
impl From<std::io::Error> for LspError {
|
||||
fn from(value: std::io::Error) -> Self {
|
||||
Self::Io(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for LspError {
|
||||
fn from(value: serde_json::Error) -> Self {
|
||||
Self::Json(value)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,284 @@
|
||||
mod client;
|
||||
mod error;
|
||||
mod manager;
|
||||
mod types;
|
||||
|
||||
pub use error::LspError;
|
||||
pub use manager::LspManager;
|
||||
pub use types::{
|
||||
FileDiagnostics, LspContextEnrichment, LspServerConfig, SymbolLocation, WorkspaceDiagnostics,
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use lsp_types::{DiagnosticSeverity, Position};
|
||||
|
||||
use crate::{LspManager, LspServerConfig};
|
||||
|
||||
fn temp_dir(label: &str) -> PathBuf {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!("lsp-{label}-{nanos}"))
|
||||
}
|
||||
|
||||
fn python3_path() -> Option<String> {
|
||||
let candidates = ["python3", "/usr/bin/python3"];
|
||||
candidates.iter().find_map(|candidate| {
|
||||
Command::new(candidate)
|
||||
.arg("--version")
|
||||
.output()
|
||||
.ok()
|
||||
.filter(|output| output.status.success())
|
||||
.map(|_| (*candidate).to_string())
|
||||
})
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_lines)]
|
||||
fn write_mock_server_script(root: &std::path::Path) -> PathBuf {
|
||||
let script_path = root.join("mock_lsp_server.py");
|
||||
fs::write(
|
||||
&script_path,
|
||||
r#"import json
|
||||
import sys
|
||||
|
||||
|
||||
def read_message():
|
||||
headers = {}
|
||||
while True:
|
||||
line = sys.stdin.buffer.readline()
|
||||
if not line:
|
||||
return None
|
||||
if line == b"\r\n":
|
||||
break
|
||||
key, value = line.decode("utf-8").split(":", 1)
|
||||
headers[key.lower()] = value.strip()
|
||||
length = int(headers["content-length"])
|
||||
body = sys.stdin.buffer.read(length)
|
||||
return json.loads(body)
|
||||
|
||||
|
||||
def write_message(payload):
|
||||
raw = json.dumps(payload).encode("utf-8")
|
||||
sys.stdout.buffer.write(f"Content-Length: {len(raw)}\r\n\r\n".encode("utf-8"))
|
||||
sys.stdout.buffer.write(raw)
|
||||
sys.stdout.buffer.flush()
|
||||
|
||||
|
||||
while True:
|
||||
message = read_message()
|
||||
if message is None:
|
||||
break
|
||||
|
||||
method = message.get("method")
|
||||
if method == "initialize":
|
||||
write_message({
|
||||
"jsonrpc": "2.0",
|
||||
"id": message["id"],
|
||||
"result": {
|
||||
"capabilities": {
|
||||
"definitionProvider": True,
|
||||
"referencesProvider": True,
|
||||
"textDocumentSync": 1,
|
||||
}
|
||||
},
|
||||
})
|
||||
elif method == "initialized":
|
||||
continue
|
||||
elif method == "textDocument/didOpen":
|
||||
document = message["params"]["textDocument"]
|
||||
write_message({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "textDocument/publishDiagnostics",
|
||||
"params": {
|
||||
"uri": document["uri"],
|
||||
"diagnostics": [
|
||||
{
|
||||
"range": {
|
||||
"start": {"line": 0, "character": 0},
|
||||
"end": {"line": 0, "character": 3},
|
||||
},
|
||||
"severity": 1,
|
||||
"source": "mock-server",
|
||||
"message": "mock error",
|
||||
}
|
||||
],
|
||||
},
|
||||
})
|
||||
elif method == "textDocument/didChange":
|
||||
continue
|
||||
elif method == "textDocument/didSave":
|
||||
continue
|
||||
elif method == "textDocument/definition":
|
||||
uri = message["params"]["textDocument"]["uri"]
|
||||
write_message({
|
||||
"jsonrpc": "2.0",
|
||||
"id": message["id"],
|
||||
"result": [
|
||||
{
|
||||
"uri": uri,
|
||||
"range": {
|
||||
"start": {"line": 0, "character": 0},
|
||||
"end": {"line": 0, "character": 3},
|
||||
},
|
||||
}
|
||||
],
|
||||
})
|
||||
elif method == "textDocument/references":
|
||||
uri = message["params"]["textDocument"]["uri"]
|
||||
write_message({
|
||||
"jsonrpc": "2.0",
|
||||
"id": message["id"],
|
||||
"result": [
|
||||
{
|
||||
"uri": uri,
|
||||
"range": {
|
||||
"start": {"line": 0, "character": 0},
|
||||
"end": {"line": 0, "character": 3},
|
||||
},
|
||||
},
|
||||
{
|
||||
"uri": uri,
|
||||
"range": {
|
||||
"start": {"line": 1, "character": 4},
|
||||
"end": {"line": 1, "character": 7},
|
||||
},
|
||||
},
|
||||
],
|
||||
})
|
||||
elif method == "shutdown":
|
||||
write_message({"jsonrpc": "2.0", "id": message["id"], "result": None})
|
||||
elif method == "exit":
|
||||
break
|
||||
"#,
|
||||
)
|
||||
.expect("mock server should be written");
|
||||
script_path
|
||||
}
|
||||
|
||||
async fn wait_for_diagnostics(manager: &LspManager) {
|
||||
tokio::time::timeout(Duration::from_secs(2), async {
|
||||
loop {
|
||||
if manager
|
||||
.collect_workspace_diagnostics()
|
||||
.await
|
||||
.expect("diagnostics snapshot should load")
|
||||
.total_diagnostics()
|
||||
> 0
|
||||
{
|
||||
break;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("diagnostics should arrive from mock server");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn collects_diagnostics_and_symbol_navigation_from_mock_server() {
|
||||
let Some(python) = python3_path() else {
|
||||
return;
|
||||
};
|
||||
|
||||
// given
|
||||
let root = temp_dir("manager");
|
||||
fs::create_dir_all(root.join("src")).expect("workspace root should exist");
|
||||
let script_path = write_mock_server_script(&root);
|
||||
let source_path = root.join("src").join("main.rs");
|
||||
fs::write(&source_path, "fn main() {}\nlet value = 1;\n").expect("source file should exist");
|
||||
let manager = LspManager::new(vec![LspServerConfig {
|
||||
name: "rust-analyzer".to_string(),
|
||||
command: python,
|
||||
args: vec![script_path.display().to_string()],
|
||||
env: BTreeMap::new(),
|
||||
workspace_root: root.clone(),
|
||||
initialization_options: None,
|
||||
extension_to_language: BTreeMap::from([(".rs".to_string(), "rust".to_string())]),
|
||||
}])
|
||||
.expect("manager should build");
|
||||
manager
|
||||
.open_document(&source_path, &fs::read_to_string(&source_path).expect("source read should succeed"))
|
||||
.await
|
||||
.expect("document should open");
|
||||
wait_for_diagnostics(&manager).await;
|
||||
|
||||
// when
|
||||
let diagnostics = manager
|
||||
.collect_workspace_diagnostics()
|
||||
.await
|
||||
.expect("diagnostics should be available");
|
||||
let definitions = manager
|
||||
.go_to_definition(&source_path, Position::new(0, 0))
|
||||
.await
|
||||
.expect("definition request should succeed");
|
||||
let references = manager
|
||||
.find_references(&source_path, Position::new(0, 0), true)
|
||||
.await
|
||||
.expect("references request should succeed");
|
||||
|
||||
// then
|
||||
assert_eq!(diagnostics.files.len(), 1);
|
||||
assert_eq!(diagnostics.total_diagnostics(), 1);
|
||||
assert_eq!(diagnostics.files[0].diagnostics[0].severity, Some(DiagnosticSeverity::ERROR));
|
||||
assert_eq!(definitions.len(), 1);
|
||||
assert_eq!(definitions[0].start_line(), 1);
|
||||
assert_eq!(references.len(), 2);
|
||||
|
||||
manager.shutdown().await.expect("shutdown should succeed");
|
||||
fs::remove_dir_all(root).expect("temp workspace should be removed");
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn renders_runtime_context_enrichment_for_prompt_usage() {
|
||||
let Some(python) = python3_path() else {
|
||||
return;
|
||||
};
|
||||
|
||||
// given
|
||||
let root = temp_dir("prompt");
|
||||
fs::create_dir_all(root.join("src")).expect("workspace root should exist");
|
||||
let script_path = write_mock_server_script(&root);
|
||||
let source_path = root.join("src").join("lib.rs");
|
||||
fs::write(&source_path, "pub fn answer() -> i32 { 42 }\n").expect("source file should exist");
|
||||
let manager = LspManager::new(vec![LspServerConfig {
|
||||
name: "rust-analyzer".to_string(),
|
||||
command: python,
|
||||
args: vec![script_path.display().to_string()],
|
||||
env: BTreeMap::new(),
|
||||
workspace_root: root.clone(),
|
||||
initialization_options: None,
|
||||
extension_to_language: BTreeMap::from([(".rs".to_string(), "rust".to_string())]),
|
||||
}])
|
||||
.expect("manager should build");
|
||||
manager
|
||||
.open_document(&source_path, &fs::read_to_string(&source_path).expect("source read should succeed"))
|
||||
.await
|
||||
.expect("document should open");
|
||||
wait_for_diagnostics(&manager).await;
|
||||
|
||||
// when
|
||||
let enrichment = manager
|
||||
.context_enrichment(&source_path, Position::new(0, 0))
|
||||
.await
|
||||
.expect("context enrichment should succeed");
|
||||
let rendered = enrichment.render_prompt_section();
|
||||
|
||||
// then
|
||||
assert!(rendered.contains("# LSP context"));
|
||||
assert!(rendered.contains("Workspace diagnostics: 1 across 1 file(s)"));
|
||||
assert!(rendered.contains("Definitions:"));
|
||||
assert!(rendered.contains("References:"));
|
||||
assert!(rendered.contains("mock error"));
|
||||
|
||||
manager.shutdown().await.expect("shutdown should succeed");
|
||||
fs::remove_dir_all(root).expect("temp workspace should be removed");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,191 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use lsp_types::Position;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::client::LspClient;
|
||||
use crate::error::LspError;
|
||||
use crate::types::{
|
||||
normalize_extension, FileDiagnostics, LspContextEnrichment, LspServerConfig, SymbolLocation,
|
||||
WorkspaceDiagnostics,
|
||||
};
|
||||
|
||||
pub struct LspManager {
|
||||
server_configs: BTreeMap<String, LspServerConfig>,
|
||||
extension_map: BTreeMap<String, String>,
|
||||
clients: Mutex<BTreeMap<String, Arc<LspClient>>>,
|
||||
}
|
||||
|
||||
impl LspManager {
|
||||
pub fn new(server_configs: Vec<LspServerConfig>) -> Result<Self, LspError> {
|
||||
let mut configs_by_name = BTreeMap::new();
|
||||
let mut extension_map = BTreeMap::new();
|
||||
|
||||
for config in server_configs {
|
||||
for extension in config.extension_to_language.keys() {
|
||||
let normalized = normalize_extension(extension);
|
||||
if let Some(existing_server) = extension_map.insert(normalized.clone(), config.name.clone()) {
|
||||
return Err(LspError::DuplicateExtension {
|
||||
extension: normalized,
|
||||
existing_server,
|
||||
new_server: config.name.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
configs_by_name.insert(config.name.clone(), config);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
server_configs: configs_by_name,
|
||||
extension_map,
|
||||
clients: Mutex::new(BTreeMap::new()),
|
||||
})
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn supports_path(&self, path: &Path) -> bool {
|
||||
path.extension().is_some_and(|extension| {
|
||||
let normalized = normalize_extension(extension.to_string_lossy().as_ref());
|
||||
self.extension_map.contains_key(&normalized)
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn open_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
|
||||
self.client_for_path(path).await?.open_document(path, text).await
|
||||
}
|
||||
|
||||
pub async fn sync_document_from_disk(&self, path: &Path) -> Result<(), LspError> {
|
||||
let contents = std::fs::read_to_string(path)?;
|
||||
self.change_document(path, &contents).await?;
|
||||
self.save_document(path).await
|
||||
}
|
||||
|
||||
pub async fn change_document(&self, path: &Path, text: &str) -> Result<(), LspError> {
|
||||
self.client_for_path(path).await?.change_document(path, text).await
|
||||
}
|
||||
|
||||
pub async fn save_document(&self, path: &Path) -> Result<(), LspError> {
|
||||
self.client_for_path(path).await?.save_document(path).await
|
||||
}
|
||||
|
||||
pub async fn close_document(&self, path: &Path) -> Result<(), LspError> {
|
||||
self.client_for_path(path).await?.close_document(path).await
|
||||
}
|
||||
|
||||
pub async fn go_to_definition(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
) -> Result<Vec<SymbolLocation>, LspError> {
|
||||
let mut locations = self.client_for_path(path).await?.go_to_definition(path, position).await?;
|
||||
dedupe_locations(&mut locations);
|
||||
Ok(locations)
|
||||
}
|
||||
|
||||
pub async fn find_references(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
include_declaration: bool,
|
||||
) -> Result<Vec<SymbolLocation>, LspError> {
|
||||
let mut locations = self
|
||||
.client_for_path(path)
|
||||
.await?
|
||||
.find_references(path, position, include_declaration)
|
||||
.await?;
|
||||
dedupe_locations(&mut locations);
|
||||
Ok(locations)
|
||||
}
|
||||
|
||||
pub async fn collect_workspace_diagnostics(&self) -> Result<WorkspaceDiagnostics, LspError> {
|
||||
let clients = self.clients.lock().await.values().cloned().collect::<Vec<_>>();
|
||||
let mut files = Vec::new();
|
||||
|
||||
for client in clients {
|
||||
for (uri, diagnostics) in client.diagnostics_snapshot().await {
|
||||
let Ok(path) = url::Url::parse(&uri)
|
||||
.and_then(|url| url.to_file_path().map_err(|()| url::ParseError::RelativeUrlWithoutBase))
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
if diagnostics.is_empty() {
|
||||
continue;
|
||||
}
|
||||
files.push(FileDiagnostics {
|
||||
path,
|
||||
uri,
|
||||
diagnostics,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
files.sort_by(|left, right| left.path.cmp(&right.path));
|
||||
Ok(WorkspaceDiagnostics { files })
|
||||
}
|
||||
|
||||
pub async fn context_enrichment(
|
||||
&self,
|
||||
path: &Path,
|
||||
position: Position,
|
||||
) -> Result<LspContextEnrichment, LspError> {
|
||||
Ok(LspContextEnrichment {
|
||||
file_path: path.to_path_buf(),
|
||||
diagnostics: self.collect_workspace_diagnostics().await?,
|
||||
definitions: self.go_to_definition(path, position).await?,
|
||||
references: self.find_references(path, position, true).await?,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn shutdown(&self) -> Result<(), LspError> {
|
||||
let mut clients = self.clients.lock().await;
|
||||
let drained = clients.values().cloned().collect::<Vec<_>>();
|
||||
clients.clear();
|
||||
drop(clients);
|
||||
|
||||
for client in drained {
|
||||
client.shutdown().await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn client_for_path(&self, path: &Path) -> Result<Arc<LspClient>, LspError> {
|
||||
let extension = path
|
||||
.extension()
|
||||
.map(|extension| normalize_extension(extension.to_string_lossy().as_ref()))
|
||||
.ok_or_else(|| LspError::UnsupportedDocument(path.to_path_buf()))?;
|
||||
let server_name = self
|
||||
.extension_map
|
||||
.get(&extension)
|
||||
.cloned()
|
||||
.ok_or_else(|| LspError::UnsupportedDocument(path.to_path_buf()))?;
|
||||
|
||||
let mut clients = self.clients.lock().await;
|
||||
if let Some(client) = clients.get(&server_name) {
|
||||
return Ok(client.clone());
|
||||
}
|
||||
|
||||
let config = self
|
||||
.server_configs
|
||||
.get(&server_name)
|
||||
.cloned()
|
||||
.ok_or_else(|| LspError::UnknownServer(server_name.clone()))?;
|
||||
let client = Arc::new(LspClient::connect(config).await?);
|
||||
clients.insert(server_name, client.clone());
|
||||
Ok(client)
|
||||
}
|
||||
}
|
||||
|
||||
fn dedupe_locations(locations: &mut Vec<SymbolLocation>) {
|
||||
let mut seen = BTreeSet::new();
|
||||
locations.retain(|location| {
|
||||
seen.insert((
|
||||
location.path.clone(),
|
||||
location.range.start.line,
|
||||
location.range.start.character,
|
||||
location.range.end.line,
|
||||
location.range.end.character,
|
||||
))
|
||||
});
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use lsp_types::{Diagnostic, Range};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LspServerConfig {
|
||||
pub name: String,
|
||||
pub command: String,
|
||||
pub args: Vec<String>,
|
||||
pub env: BTreeMap<String, String>,
|
||||
pub workspace_root: PathBuf,
|
||||
pub initialization_options: Option<Value>,
|
||||
pub extension_to_language: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
impl LspServerConfig {
|
||||
#[must_use]
|
||||
pub fn language_id_for(&self, path: &Path) -> Option<&str> {
|
||||
let extension = normalize_extension(path.extension()?.to_string_lossy().as_ref());
|
||||
self.extension_to_language
|
||||
.get(&extension)
|
||||
.map(String::as_str)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct FileDiagnostics {
|
||||
pub path: PathBuf,
|
||||
pub uri: String,
|
||||
pub diagnostics: Vec<Diagnostic>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct WorkspaceDiagnostics {
|
||||
pub files: Vec<FileDiagnostics>,
|
||||
}
|
||||
|
||||
impl WorkspaceDiagnostics {
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.files.is_empty()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn total_diagnostics(&self) -> usize {
|
||||
self.files.iter().map(|file| file.diagnostics.len()).sum()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SymbolLocation {
|
||||
pub path: PathBuf,
|
||||
pub range: Range,
|
||||
}
|
||||
|
||||
impl SymbolLocation {
|
||||
#[must_use]
|
||||
pub fn start_line(&self) -> u32 {
|
||||
self.range.start.line + 1
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn start_character(&self) -> u32 {
|
||||
self.range.start.character + 1
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for SymbolLocation {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"{}:{}:{}",
|
||||
self.path.display(),
|
||||
self.start_line(),
|
||||
self.start_character()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct LspContextEnrichment {
|
||||
pub file_path: PathBuf,
|
||||
pub diagnostics: WorkspaceDiagnostics,
|
||||
pub definitions: Vec<SymbolLocation>,
|
||||
pub references: Vec<SymbolLocation>,
|
||||
}
|
||||
|
||||
impl LspContextEnrichment {
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.diagnostics.is_empty() && self.definitions.is_empty() && self.references.is_empty()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn render_prompt_section(&self) -> String {
|
||||
const MAX_RENDERED_DIAGNOSTICS: usize = 12;
|
||||
const MAX_RENDERED_LOCATIONS: usize = 12;
|
||||
|
||||
let mut lines = vec!["# LSP context".to_string()];
|
||||
lines.push(format!(" - Focus file: {}", self.file_path.display()));
|
||||
lines.push(format!(
|
||||
" - Workspace diagnostics: {} across {} file(s)",
|
||||
self.diagnostics.total_diagnostics(),
|
||||
self.diagnostics.files.len()
|
||||
));
|
||||
|
||||
if !self.diagnostics.files.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("Diagnostics:".to_string());
|
||||
let mut rendered = 0usize;
|
||||
for file in &self.diagnostics.files {
|
||||
for diagnostic in &file.diagnostics {
|
||||
if rendered == MAX_RENDERED_DIAGNOSTICS {
|
||||
lines.push(" - Additional diagnostics omitted for brevity.".to_string());
|
||||
break;
|
||||
}
|
||||
let severity = diagnostic_severity_label(diagnostic.severity);
|
||||
lines.push(format!(
|
||||
" - {}:{}:{} [{}] {}",
|
||||
file.path.display(),
|
||||
diagnostic.range.start.line + 1,
|
||||
diagnostic.range.start.character + 1,
|
||||
severity,
|
||||
diagnostic.message.replace('\n', " ")
|
||||
));
|
||||
rendered += 1;
|
||||
}
|
||||
if rendered == MAX_RENDERED_DIAGNOSTICS {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if !self.definitions.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("Definitions:".to_string());
|
||||
lines.extend(
|
||||
self.definitions
|
||||
.iter()
|
||||
.take(MAX_RENDERED_LOCATIONS)
|
||||
.map(|location| format!(" - {location}")),
|
||||
);
|
||||
if self.definitions.len() > MAX_RENDERED_LOCATIONS {
|
||||
lines.push(" - Additional definitions omitted for brevity.".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
if !self.references.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("References:".to_string());
|
||||
lines.extend(
|
||||
self.references
|
||||
.iter()
|
||||
.take(MAX_RENDERED_LOCATIONS)
|
||||
.map(|location| format!(" - {location}")),
|
||||
);
|
||||
if self.references.len() > MAX_RENDERED_LOCATIONS {
|
||||
lines.push(" - Additional references omitted for brevity.".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
lines.join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub(crate) fn normalize_extension(extension: &str) -> String {
|
||||
if extension.starts_with('.') {
|
||||
extension.to_ascii_lowercase()
|
||||
} else {
|
||||
format!(".{}", extension.to_ascii_lowercase())
|
||||
}
|
||||
}
|
||||
|
||||
fn diagnostic_severity_label(severity: Option<lsp_types::DiagnosticSeverity>) -> &'static str {
|
||||
match severity {
|
||||
Some(lsp_types::DiagnosticSeverity::ERROR) => "error",
|
||||
Some(lsp_types::DiagnosticSeverity::WARNING) => "warning",
|
||||
Some(lsp_types::DiagnosticSeverity::INFORMATION) => "info",
|
||||
Some(lsp_types::DiagnosticSeverity::HINT) => "hint",
|
||||
_ => "unknown",
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
[package]
|
||||
name = "mock-anthropic-service"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
publish.workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "mock-anthropic-service"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
api = { path = "../api" }
|
||||
serde_json.workspace = true
|
||||
tokio = { version = "1", features = ["io-util", "macros", "net", "rt-multi-thread", "signal", "sync"] }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,34 @@
|
||||
use std::env;
|
||||
|
||||
use mock_anthropic_service::MockAnthropicService;
|
||||
|
||||
#[tokio::main(flavor = "multi_thread")]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let mut bind_addr = String::from("127.0.0.1:0");
|
||||
let mut args = env::args().skip(1);
|
||||
while let Some(arg) = args.next() {
|
||||
match arg.as_str() {
|
||||
"--bind" => {
|
||||
bind_addr = args
|
||||
.next()
|
||||
.ok_or_else(|| "missing value for --bind".to_string())?;
|
||||
}
|
||||
flag if flag.starts_with("--bind=") => {
|
||||
bind_addr = flag[7..].to_string();
|
||||
}
|
||||
"--help" | "-h" => {
|
||||
println!("Usage: mock-anthropic-service [--bind HOST:PORT]");
|
||||
return Ok(());
|
||||
}
|
||||
other => {
|
||||
return Err(format!("unsupported argument: {other}").into());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let server = MockAnthropicService::spawn_on(&bind_addr).await?;
|
||||
println!("MOCK_ANTHROPIC_BASE_URL={}", server.base_url());
|
||||
tokio::signal::ctrl_c().await?;
|
||||
drop(server);
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
[package]
|
||||
name = "plugins"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
publish.workspace = true
|
||||
|
||||
[dependencies]
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json.workspace = true
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,59 @@
|
||||
# 插件模块 (plugins)
|
||||
|
||||
本模块实现了 Claw 的插件系统,允许通过外部扩展来增强 AI 的功能,包括自定义工具、命令以及在工具执行前后运行的钩子 (Hooks)。
|
||||
|
||||
## 概览
|
||||
|
||||
`plugins` 模块的主要职责是:
|
||||
- 定义插件的清单格式 (`plugin.json`) 和元数据结构。
|
||||
- 管理插件的完整生命周期:安装、加载、初始化、启用/禁用、更新和卸载。
|
||||
- 提供插件类型的抽象:
|
||||
- **Builtin (内置)**:编译在程序内部的插件。
|
||||
- **Bundled (绑定)**:随应用程序分发但作为独立文件存在的插件。
|
||||
- **External (外部)**:用户自行安装或从远程仓库下载的插件。
|
||||
- 实现插件隔离与执行机制,支持插件定义的自定义工具。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **插件清单 (PluginManifest)**:每个插件必须包含一个 `plugin.json`,详细说明其名称、版本、所需权限、钩子、生命周期脚本以及它所暴露的工具。
|
||||
- **自定义工具 (PluginTool)**:插件可以定义全新的工具供 AI 调用。这些工具在执行时被作为独立的外部进程启动。
|
||||
- **钩子系统 (Hooks)**:支持 `PreToolUse` 和 `PostToolUse` 钩子,允许插件在 AI 调用任何工具之前或之后执行特定的逻辑。
|
||||
- **生命周期管理**:提供 `Init` 和 `Shutdown` 阶段,允许插件在加载时进行环境准备,在卸载或关闭时进行清理。
|
||||
- **权限模型**:强制要求插件声明权限(`read`, `write`, `execute`),并为定义的工具指定安全级别(`read-only`, `workspace-write`, `danger-full-access`)。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心模块
|
||||
|
||||
- **`lib.rs`**: 包含了插件定义的各种结构体(Manifest, Metadata, Tool, Permission 等)以及插件特性的定义。
|
||||
- **`manager.rs`**: 实现了 `PluginManager`,负责插件在磁盘上的组织、注册表的维护以及安装/更新逻辑。
|
||||
- **`hooks.rs`**: 实现了钩子执行器 (`HookRunner`),负责在正确的时机触发插件定义的脚本。
|
||||
|
||||
### 插件加载与执行流程
|
||||
|
||||
1. `PluginManager` 扫描指定的目录(内置、绑定及外部安装目录)。
|
||||
2. 读取并验证每个插件的 `plugin.json`。
|
||||
3. 如果插件被启用,则初始化该插件并将其定义的工具注册到系统的全局工具注册表中。
|
||||
4. 当 AI 调用插件工具时,系统根据清单中定义的命令行信息启动一个子进程,并通过标准输入/环境变量传递参数。
|
||||
5. 钩子逻辑会在工具执行的生命周期内被自动触发。
|
||||
|
||||
## 使用示例 (插件定义样例)
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "my-custom-plugin",
|
||||
"version": "1.0.0",
|
||||
"description": "一个演示插件",
|
||||
"permissions": ["read", "execute"],
|
||||
"tools": [
|
||||
{
|
||||
"name": "custom_search",
|
||||
"description": "执行自定义搜索",
|
||||
"inputSchema": { "type": "object", "properties": { "query": { "type": "string" } } },
|
||||
"command": "python3",
|
||||
"args": ["search_script.py"],
|
||||
"requiredPermission": "read-only"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
@@ -0,0 +1,10 @@
|
||||
{
|
||||
"name": "example-bundled",
|
||||
"version": "0.1.0",
|
||||
"description": "Example bundled plugin scaffold for the Rust plugin system",
|
||||
"defaultEnabled": false,
|
||||
"hooks": {
|
||||
"PreToolUse": ["./hooks/pre.sh"],
|
||||
"PostToolUse": ["./hooks/post.sh"]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
#!/bin/sh
|
||||
printf '%s\n' 'example bundled post hook'
|
||||
@@ -0,0 +1,2 @@
|
||||
#!/bin/sh
|
||||
printf '%s\n' 'example bundled pre hook'
|
||||
@@ -0,0 +1,10 @@
|
||||
{
|
||||
"name": "sample-hooks",
|
||||
"version": "0.1.0",
|
||||
"description": "Bundled sample plugin scaffold for hook integration tests.",
|
||||
"defaultEnabled": false,
|
||||
"hooks": {
|
||||
"PreToolUse": ["./hooks/pre.sh"],
|
||||
"PostToolUse": ["./hooks/post.sh"]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
#!/bin/sh
|
||||
printf 'sample bundled post hook'
|
||||
@@ -0,0 +1,2 @@
|
||||
#!/bin/sh
|
||||
printf 'sample bundled pre hook'
|
||||
@@ -0,0 +1,563 @@
|
||||
use std::ffi::OsStr;
|
||||
use std::process::Command;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use crate::{PluginError, PluginHooks, PluginRegistry};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum HookEvent {
|
||||
PreToolUse,
|
||||
PostToolUse,
|
||||
PostToolUseFailure,
|
||||
}
|
||||
|
||||
impl HookEvent {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::PreToolUse => "PreToolUse",
|
||||
Self::PostToolUse => "PostToolUse",
|
||||
Self::PostToolUseFailure => "PostToolUseFailure",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct HookRunResult {
|
||||
denied: bool,
|
||||
failed: bool,
|
||||
messages: Vec<String>,
|
||||
}
|
||||
|
||||
impl HookRunResult {
|
||||
#[must_use]
|
||||
pub fn allow(messages: Vec<String>) -> Self {
|
||||
Self {
|
||||
denied: false,
|
||||
failed: false,
|
||||
messages,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_denied(&self) -> bool {
|
||||
self.denied
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_failed(&self) -> bool {
|
||||
self.failed
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn messages(&self) -> &[String] {
|
||||
&self.messages
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct HookRunner {
|
||||
hooks: PluginHooks,
|
||||
}
|
||||
|
||||
impl HookRunner {
|
||||
#[must_use]
|
||||
pub fn new(hooks: PluginHooks) -> Self {
|
||||
Self { hooks }
|
||||
}
|
||||
|
||||
pub fn from_registry(plugin_registry: &PluginRegistry) -> Result<Self, PluginError> {
|
||||
Ok(Self::new(plugin_registry.aggregated_hooks()?))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_pre_tool_use(&self, tool_name: &str, tool_input: &str) -> HookRunResult {
|
||||
Self::run_commands(
|
||||
HookEvent::PreToolUse,
|
||||
&self.hooks.pre_tool_use,
|
||||
tool_name,
|
||||
tool_input,
|
||||
None,
|
||||
false,
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_post_tool_use(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: &str,
|
||||
is_error: bool,
|
||||
) -> HookRunResult {
|
||||
Self::run_commands(
|
||||
HookEvent::PostToolUse,
|
||||
&self.hooks.post_tool_use,
|
||||
tool_name,
|
||||
tool_input,
|
||||
Some(tool_output),
|
||||
is_error,
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_post_tool_use_failure(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_error: &str,
|
||||
) -> HookRunResult {
|
||||
Self::run_commands(
|
||||
HookEvent::PostToolUseFailure,
|
||||
&self.hooks.post_tool_use_failure,
|
||||
tool_name,
|
||||
tool_input,
|
||||
Some(tool_error),
|
||||
true,
|
||||
)
|
||||
}
|
||||
|
||||
fn run_commands(
|
||||
event: HookEvent,
|
||||
commands: &[String],
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: Option<&str>,
|
||||
is_error: bool,
|
||||
) -> HookRunResult {
|
||||
if commands.is_empty() {
|
||||
return HookRunResult::allow(Vec::new());
|
||||
}
|
||||
|
||||
let payload = hook_payload(event, tool_name, tool_input, tool_output, is_error).to_string();
|
||||
|
||||
let mut messages = Vec::new();
|
||||
|
||||
for command in commands {
|
||||
match Self::run_command(
|
||||
command,
|
||||
event,
|
||||
tool_name,
|
||||
tool_input,
|
||||
tool_output,
|
||||
is_error,
|
||||
&payload,
|
||||
) {
|
||||
HookCommandOutcome::Allow { message } => {
|
||||
if let Some(message) = message {
|
||||
messages.push(message);
|
||||
}
|
||||
}
|
||||
HookCommandOutcome::Deny { message } => {
|
||||
messages.push(message.unwrap_or_else(|| {
|
||||
format!("{} hook denied tool `{tool_name}`", event.as_str())
|
||||
}));
|
||||
return HookRunResult {
|
||||
denied: true,
|
||||
failed: false,
|
||||
messages,
|
||||
};
|
||||
}
|
||||
HookCommandOutcome::Failed { message } => {
|
||||
messages.push(message);
|
||||
return HookRunResult {
|
||||
denied: false,
|
||||
failed: true,
|
||||
messages,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
HookRunResult::allow(messages)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn run_command(
|
||||
command: &str,
|
||||
event: HookEvent,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: Option<&str>,
|
||||
is_error: bool,
|
||||
payload: &str,
|
||||
) -> HookCommandOutcome {
|
||||
let mut child = shell_command(command);
|
||||
child.stdin(std::process::Stdio::piped());
|
||||
child.stdout(std::process::Stdio::piped());
|
||||
child.stderr(std::process::Stdio::piped());
|
||||
child.env("HOOK_EVENT", event.as_str());
|
||||
child.env("HOOK_TOOL_NAME", tool_name);
|
||||
child.env("HOOK_TOOL_INPUT", tool_input);
|
||||
child.env("HOOK_TOOL_IS_ERROR", if is_error { "1" } else { "0" });
|
||||
if let Some(tool_output) = tool_output {
|
||||
child.env("HOOK_TOOL_OUTPUT", tool_output);
|
||||
}
|
||||
|
||||
match child.output_with_stdin(payload.as_bytes()) {
|
||||
Ok(output) => {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
|
||||
let message = (!stdout.is_empty()).then_some(stdout);
|
||||
match output.status.code() {
|
||||
Some(0) => HookCommandOutcome::Allow { message },
|
||||
Some(2) => HookCommandOutcome::Deny { message },
|
||||
Some(code) => HookCommandOutcome::Failed {
|
||||
message: format_hook_warning(
|
||||
command,
|
||||
code,
|
||||
message.as_deref(),
|
||||
stderr.as_str(),
|
||||
),
|
||||
},
|
||||
None => HookCommandOutcome::Failed {
|
||||
message: format!(
|
||||
"{} hook `{command}` terminated by signal while handling `{tool_name}`",
|
||||
event.as_str()
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
Err(error) => HookCommandOutcome::Failed {
|
||||
message: format!(
|
||||
"{} hook `{command}` failed to start for `{tool_name}`: {error}",
|
||||
event.as_str()
|
||||
),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum HookCommandOutcome {
|
||||
Allow { message: Option<String> },
|
||||
Deny { message: Option<String> },
|
||||
Failed { message: String },
|
||||
}
|
||||
|
||||
fn hook_payload(
|
||||
event: HookEvent,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: Option<&str>,
|
||||
is_error: bool,
|
||||
) -> serde_json::Value {
|
||||
match event {
|
||||
HookEvent::PostToolUseFailure => json!({
|
||||
"hook_event_name": event.as_str(),
|
||||
"tool_name": tool_name,
|
||||
"tool_input": parse_tool_input(tool_input),
|
||||
"tool_input_json": tool_input,
|
||||
"tool_error": tool_output,
|
||||
"tool_result_is_error": true,
|
||||
}),
|
||||
_ => json!({
|
||||
"hook_event_name": event.as_str(),
|
||||
"tool_name": tool_name,
|
||||
"tool_input": parse_tool_input(tool_input),
|
||||
"tool_input_json": tool_input,
|
||||
"tool_output": tool_output,
|
||||
"tool_result_is_error": is_error,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_tool_input(tool_input: &str) -> serde_json::Value {
|
||||
serde_json::from_str(tool_input).unwrap_or_else(|_| json!({ "raw": tool_input }))
|
||||
}
|
||||
|
||||
fn format_hook_warning(command: &str, code: i32, stdout: Option<&str>, stderr: &str) -> String {
|
||||
let mut message = format!("Hook `{command}` exited with status {code}");
|
||||
if let Some(stdout) = stdout.filter(|stdout| !stdout.is_empty()) {
|
||||
message.push_str(": ");
|
||||
message.push_str(stdout);
|
||||
} else if !stderr.is_empty() {
|
||||
message.push_str(": ");
|
||||
message.push_str(stderr);
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
fn shell_command(command: &str) -> CommandWithStdin {
|
||||
#[cfg(windows)]
|
||||
let command_builder = {
|
||||
let mut command_builder = Command::new("cmd");
|
||||
command_builder.arg("/C").arg(command);
|
||||
CommandWithStdin::new(command_builder)
|
||||
};
|
||||
|
||||
#[cfg(not(windows))]
|
||||
let command_builder = if Path::new(command).exists() {
|
||||
let mut command_builder = Command::new("sh");
|
||||
command_builder.arg(command);
|
||||
CommandWithStdin::new(command_builder)
|
||||
} else {
|
||||
let mut command_builder = Command::new("sh");
|
||||
command_builder.arg("-lc").arg(command);
|
||||
CommandWithStdin::new(command_builder)
|
||||
};
|
||||
|
||||
command_builder
|
||||
}
|
||||
|
||||
struct CommandWithStdin {
|
||||
command: Command,
|
||||
}
|
||||
|
||||
impl CommandWithStdin {
|
||||
fn new(command: Command) -> Self {
|
||||
Self { command }
|
||||
}
|
||||
|
||||
fn stdin(&mut self, cfg: std::process::Stdio) -> &mut Self {
|
||||
self.command.stdin(cfg);
|
||||
self
|
||||
}
|
||||
|
||||
fn stdout(&mut self, cfg: std::process::Stdio) -> &mut Self {
|
||||
self.command.stdout(cfg);
|
||||
self
|
||||
}
|
||||
|
||||
fn stderr(&mut self, cfg: std::process::Stdio) -> &mut Self {
|
||||
self.command.stderr(cfg);
|
||||
self
|
||||
}
|
||||
|
||||
fn env<K, V>(&mut self, key: K, value: V) -> &mut Self
|
||||
where
|
||||
K: AsRef<OsStr>,
|
||||
V: AsRef<OsStr>,
|
||||
{
|
||||
self.command.env(key, value);
|
||||
self
|
||||
}
|
||||
|
||||
fn output_with_stdin(&mut self, stdin: &[u8]) -> std::io::Result<std::process::Output> {
|
||||
let mut child = self.command.spawn()?;
|
||||
if let Some(mut child_stdin) = child.stdin.take() {
|
||||
use std::io::Write as _;
|
||||
// Tolerate BrokenPipe: a hook script that runs to completion
|
||||
// (or exits early without reading stdin) closes its stdin
|
||||
// before the parent finishes writing the JSON payload, and
|
||||
// the kernel raises EPIPE on the parent's write_all. That is
|
||||
// not a hook failure — the child still exited cleanly and we
|
||||
// still need to wait_with_output() to capture stdout/stderr
|
||||
// and the real exit code. Other write errors (e.g. EIO,
|
||||
// permission, OOM) still propagate.
|
||||
//
|
||||
// This was the root cause of the Linux CI flake on
|
||||
// hooks::tests::collects_and_runs_hooks_from_enabled_plugins
|
||||
// (ROADMAP #25, runs 24120271422 / 24120538408 / 24121392171
|
||||
// / 24121776826): the test hook scripts run in microseconds
|
||||
// and the parent's stdin write races against child exit.
|
||||
// macOS pipes happen to buffer the small payload before the
|
||||
// child exits; Linux pipes do not, so the race shows up
|
||||
// deterministically on ubuntu runners.
|
||||
match child_stdin.write_all(stdin) {
|
||||
Ok(()) => {}
|
||||
Err(error) if error.kind() == std::io::ErrorKind::BrokenPipe => {}
|
||||
Err(error) => return Err(error),
|
||||
}
|
||||
}
|
||||
child.wait_with_output()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{HookRunResult, HookRunner};
|
||||
use crate::{PluginManager, PluginManagerConfig};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_dir(label: &str) -> PathBuf {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!("plugins-hook-runner-{label}-{nanos}"))
|
||||
}
|
||||
|
||||
fn make_executable(path: &Path) {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let perms = fs::Permissions::from_mode(0o755);
|
||||
fs::set_permissions(path, perms)
|
||||
.unwrap_or_else(|e| panic!("chmod +x {}: {e}", path.display()));
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
let _ = path;
|
||||
}
|
||||
|
||||
fn write_hook_plugin(
|
||||
root: &Path,
|
||||
name: &str,
|
||||
pre_message: &str,
|
||||
post_message: &str,
|
||||
failure_message: &str,
|
||||
) {
|
||||
fs::create_dir_all(root.join(".claude-plugin")).expect("manifest dir");
|
||||
fs::create_dir_all(root.join("hooks")).expect("hooks dir");
|
||||
|
||||
let pre_path = root.join("hooks").join("pre.sh");
|
||||
fs::write(
|
||||
&pre_path,
|
||||
format!("#!/bin/sh\nprintf '%s\\n' '{pre_message}'\n"),
|
||||
)
|
||||
.expect("write pre hook");
|
||||
make_executable(&pre_path);
|
||||
|
||||
let post_path = root.join("hooks").join("post.sh");
|
||||
fs::write(
|
||||
&post_path,
|
||||
format!("#!/bin/sh\nprintf '%s\\n' '{post_message}'\n"),
|
||||
)
|
||||
.expect("write post hook");
|
||||
make_executable(&post_path);
|
||||
|
||||
let failure_path = root.join("hooks").join("failure.sh");
|
||||
fs::write(
|
||||
&failure_path,
|
||||
format!("#!/bin/sh\nprintf '%s\\n' '{failure_message}'\n"),
|
||||
)
|
||||
.expect("write failure hook");
|
||||
make_executable(&failure_path);
|
||||
fs::write(
|
||||
root.join(".claude-plugin").join("plugin.json"),
|
||||
format!(
|
||||
"{{\n \"name\": \"{name}\",\n \"version\": \"1.0.0\",\n \"description\": \"hook plugin\",\n \"hooks\": {{\n \"PreToolUse\": [\"./hooks/pre.sh\"],\n \"PostToolUse\": [\"./hooks/post.sh\"],\n \"PostToolUseFailure\": [\"./hooks/failure.sh\"]\n }}\n}}"
|
||||
),
|
||||
)
|
||||
.expect("write plugin manifest");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collects_and_runs_hooks_from_enabled_plugins() {
|
||||
// given
|
||||
let config_home = temp_dir("config");
|
||||
let first_source_root = temp_dir("source-a");
|
||||
let second_source_root = temp_dir("source-b");
|
||||
write_hook_plugin(
|
||||
&first_source_root,
|
||||
"first",
|
||||
"plugin pre one",
|
||||
"plugin post one",
|
||||
"plugin failure one",
|
||||
);
|
||||
write_hook_plugin(
|
||||
&second_source_root,
|
||||
"second",
|
||||
"plugin pre two",
|
||||
"plugin post two",
|
||||
"plugin failure two",
|
||||
);
|
||||
|
||||
let mut manager = PluginManager::new(PluginManagerConfig::new(&config_home));
|
||||
manager
|
||||
.install(first_source_root.to_str().expect("utf8 path"))
|
||||
.expect("first plugin install should succeed");
|
||||
manager
|
||||
.install(second_source_root.to_str().expect("utf8 path"))
|
||||
.expect("second plugin install should succeed");
|
||||
let registry = manager.plugin_registry().expect("registry should build");
|
||||
|
||||
// when
|
||||
let runner = HookRunner::from_registry(®istry).expect("plugin hooks should load");
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
runner.run_pre_tool_use("Read", r#"{"path":"README.md"}"#),
|
||||
HookRunResult::allow(vec![
|
||||
"plugin pre one".to_string(),
|
||||
"plugin pre two".to_string(),
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
runner.run_post_tool_use("Read", r#"{"path":"README.md"}"#, "ok", false),
|
||||
HookRunResult::allow(vec![
|
||||
"plugin post one".to_string(),
|
||||
"plugin post two".to_string(),
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
runner.run_post_tool_use_failure("Read", r#"{"path":"README.md"}"#, "tool failed",),
|
||||
HookRunResult::allow(vec![
|
||||
"plugin failure one".to_string(),
|
||||
"plugin failure two".to_string(),
|
||||
])
|
||||
);
|
||||
|
||||
let _ = fs::remove_dir_all(config_home);
|
||||
let _ = fs::remove_dir_all(first_source_root);
|
||||
let _ = fs::remove_dir_all(second_source_root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pre_tool_use_denies_when_plugin_hook_exits_two() {
|
||||
// given
|
||||
let runner = HookRunner::new(crate::PluginHooks {
|
||||
pre_tool_use: vec!["printf 'blocked by plugin'; exit 2".to_string()],
|
||||
post_tool_use: Vec::new(),
|
||||
post_tool_use_failure: Vec::new(),
|
||||
});
|
||||
|
||||
// when
|
||||
let result = runner.run_pre_tool_use("Bash", r#"{"command":"pwd"}"#);
|
||||
|
||||
// then
|
||||
assert!(result.is_denied());
|
||||
assert_eq!(result.messages(), &["blocked by plugin".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn propagates_plugin_hook_failures() {
|
||||
// given
|
||||
let runner = HookRunner::new(crate::PluginHooks {
|
||||
pre_tool_use: vec![
|
||||
"printf 'broken plugin hook'; exit 1".to_string(),
|
||||
"printf 'later plugin hook'".to_string(),
|
||||
],
|
||||
post_tool_use: Vec::new(),
|
||||
post_tool_use_failure: Vec::new(),
|
||||
});
|
||||
|
||||
// when
|
||||
let result = runner.run_pre_tool_use("Bash", r#"{"command":"pwd"}"#);
|
||||
|
||||
// then
|
||||
assert!(result.is_failed());
|
||||
assert!(result
|
||||
.messages()
|
||||
.iter()
|
||||
.any(|message| message.contains("broken plugin hook")));
|
||||
assert!(!result
|
||||
.messages()
|
||||
.iter()
|
||||
.any(|message| message == "later plugin hook"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[cfg(unix)]
|
||||
fn generated_hook_scripts_are_executable() {
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
// given
|
||||
let root = temp_dir("exec-guard");
|
||||
write_hook_plugin(&root, "exec-check", "pre", "post", "fail");
|
||||
|
||||
// then
|
||||
for script in ["pre.sh", "post.sh", "failure.sh"] {
|
||||
let path = root.join("hooks").join(script);
|
||||
let mode = fs::metadata(&path)
|
||||
.unwrap_or_else(|e| panic!("{script} metadata: {e}"))
|
||||
.permissions()
|
||||
.mode();
|
||||
assert!(
|
||||
mode & 0o111 != 0,
|
||||
"{script} must have at least one execute bit set, got mode {mode:#o}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
[package]
|
||||
name = "runtime"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
publish.workspace = true
|
||||
|
||||
[dependencies]
|
||||
sha2 = "0.10"
|
||||
which = "7"
|
||||
glob = "0.3"
|
||||
plugins = { path = "../plugins" }
|
||||
regex = "1"
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json.workspace = true
|
||||
telemetry = { path = "../telemetry" }
|
||||
tokio = { version = "1", features = ["io-std", "io-util", "macros", "process", "rt", "rt-multi-thread", "time"] }
|
||||
walkdir = "2"
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -0,0 +1,60 @@
|
||||
# 运行时模块 (runtime)
|
||||
|
||||
本模块是 Claw 的核心引擎,负责协调 AI 模型、工具执行、会话管理和权限控制之间的所有交互。
|
||||
|
||||
## 概览
|
||||
|
||||
`runtime` 模块是整个系统的“中枢神经”,其主要职责包括:
|
||||
- **对话驱动**:管理“用户-助手-工具”的循环迭代。
|
||||
- **会话持久化**:负责会话的加载、保存及历史记录的压缩 (Compaction)。
|
||||
- **MCP 客户端**:实现模型上下文协议 (Model Context Protocol),支持与外部 MCP 服务器通信。
|
||||
- **安全沙箱与权限**:实施基于策略的工具执行权限检查。
|
||||
- **上下文构建**:动态生成系统提示词 (System Prompt),集成工作区上下文。
|
||||
- **消耗统计**:精确跟踪 Token 使用情况和 Token 缓存状态。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **ConversationRuntime**:核心驱动类,支持流式响应处理和多轮工具调用迭代。
|
||||
- **权限引擎 (Permissions)**:提供多种模式(`ReadOnly`, `WorkspaceWrite`, `DangerFullAccess`),并支持交互式权限确认。
|
||||
- **会话压缩 (Compaction)**:当对话历史过长影响性能或成本时,自动将旧消息总结为摘要,保持上下文精简。
|
||||
- **钩子集成 (Hooks)**:在工具执行的前后触发插件定义的钩子,支持干预工具输入或处理执行结果。
|
||||
- **沙箱执行 (Sandbox)**:为 Bash 等敏感工具提供受限的执行环境。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心子模块
|
||||
|
||||
- **`conversation.rs`**: 定义了核心的 `ConversationRuntime` 和 `ApiClient`/`ToolExecutor` 特性。
|
||||
- **`mcp_stdio.rs` / `mcp_client.rs`**: 实现了完整的 MCP 规范,支持通过标准输入/输出与外部工具服务器交互。
|
||||
- **`session.rs`**: 定义了消息模型 (`ConversationMessage`)、内容块 (`ContentBlock`) 和会话序列化逻辑。
|
||||
- **`permissions.rs`**: 实现了权限审核逻辑和提示器接口。
|
||||
- **`compact.rs`**: 包含了基于 LLM 的会话摘要生成和历史裁剪算法。
|
||||
- **`config.rs`**: 负责加载和合并多层级的配置文件。
|
||||
|
||||
### 对话循环流程 (run_turn)
|
||||
|
||||
1. 将用户输入推入 `Session`。
|
||||
2. 调用 `ApiClient` 发起流式请求。
|
||||
3. 监听 `AssistantEvent`,解析文本内容和工具调用请求。
|
||||
4. **工具权限审核**:针对每个 `ToolUse`,根据 `PermissionPolicy` 决定是允许、拒绝还是询问用户。
|
||||
5. **执行工具**:若允许,则通过 `ToolExecutor`(或 MCP 客户端)执行工具,并运行相关的 `Pre/Post Hooks`。
|
||||
6. 将工具结果反馈给 AI,进入下一轮迭代,直到 AI 给出最终回复。
|
||||
|
||||
## 使用示例 (内部)
|
||||
|
||||
```rust
|
||||
use runtime::{ConversationRuntime, Session, PermissionPolicy, PermissionMode};
|
||||
|
||||
// 初始化运行时
|
||||
let mut runtime = ConversationRuntime::new(
|
||||
Session::new(),
|
||||
api_client,
|
||||
tool_executor,
|
||||
PermissionPolicy::new(PermissionMode::WorkspaceWrite),
|
||||
system_prompt,
|
||||
);
|
||||
|
||||
// 运行一轮对话
|
||||
let summary = runtime.run_turn("帮我重构 src/lib.rs", Some(&mut cli_prompter))?;
|
||||
println!("共迭代 {} 次,消耗 {} tokens", summary.iterations, summary.usage.total_tokens());
|
||||
```
|
||||
@@ -0,0 +1,361 @@
|
||||
use std::env;
|
||||
use std::io;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::time::Duration;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::process::Command as TokioCommand;
|
||||
use tokio::runtime::Builder;
|
||||
use tokio::time::timeout;
|
||||
|
||||
use crate::sandbox::{
|
||||
build_linux_sandbox_command, resolve_sandbox_status_for_request, FilesystemIsolationMode,
|
||||
SandboxConfig, SandboxStatus,
|
||||
};
|
||||
use crate::ConfigLoader;
|
||||
|
||||
/// Input schema for the built-in bash execution tool.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct BashCommandInput {
|
||||
pub command: String,
|
||||
pub timeout: Option<u64>,
|
||||
pub description: Option<String>,
|
||||
#[serde(rename = "run_in_background")]
|
||||
pub run_in_background: Option<bool>,
|
||||
#[serde(rename = "dangerouslyDisableSandbox")]
|
||||
pub dangerously_disable_sandbox: Option<bool>,
|
||||
#[serde(rename = "namespaceRestrictions")]
|
||||
pub namespace_restrictions: Option<bool>,
|
||||
#[serde(rename = "isolateNetwork")]
|
||||
pub isolate_network: Option<bool>,
|
||||
#[serde(rename = "filesystemMode")]
|
||||
pub filesystem_mode: Option<FilesystemIsolationMode>,
|
||||
#[serde(rename = "allowedMounts")]
|
||||
pub allowed_mounts: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
/// Output returned from a bash tool invocation.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct BashCommandOutput {
|
||||
pub stdout: String,
|
||||
pub stderr: String,
|
||||
#[serde(rename = "rawOutputPath")]
|
||||
pub raw_output_path: Option<String>,
|
||||
pub interrupted: bool,
|
||||
#[serde(rename = "isImage")]
|
||||
pub is_image: Option<bool>,
|
||||
#[serde(rename = "backgroundTaskId")]
|
||||
pub background_task_id: Option<String>,
|
||||
#[serde(rename = "backgroundedByUser")]
|
||||
pub backgrounded_by_user: Option<bool>,
|
||||
#[serde(rename = "assistantAutoBackgrounded")]
|
||||
pub assistant_auto_backgrounded: Option<bool>,
|
||||
#[serde(rename = "dangerouslyDisableSandbox")]
|
||||
pub dangerously_disable_sandbox: Option<bool>,
|
||||
#[serde(rename = "returnCodeInterpretation")]
|
||||
pub return_code_interpretation: Option<String>,
|
||||
#[serde(rename = "noOutputExpected")]
|
||||
pub no_output_expected: Option<bool>,
|
||||
#[serde(rename = "structuredContent")]
|
||||
pub structured_content: Option<Vec<serde_json::Value>>,
|
||||
#[serde(rename = "persistedOutputPath")]
|
||||
pub persisted_output_path: Option<String>,
|
||||
#[serde(rename = "persistedOutputSize")]
|
||||
pub persisted_output_size: Option<u64>,
|
||||
#[serde(rename = "sandboxStatus")]
|
||||
pub sandbox_status: Option<SandboxStatus>,
|
||||
}
|
||||
|
||||
/// Executes a shell command with the requested sandbox settings.
|
||||
pub fn execute_bash(input: BashCommandInput) -> io::Result<BashCommandOutput> {
|
||||
let cwd = env::current_dir()?;
|
||||
let sandbox_status = sandbox_status_for_input(&input, &cwd);
|
||||
|
||||
if input.run_in_background.unwrap_or(false) {
|
||||
let mut child = prepare_command(&input.command, &cwd, &sandbox_status, false);
|
||||
let child = child
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.spawn()?;
|
||||
|
||||
return Ok(BashCommandOutput {
|
||||
stdout: String::new(),
|
||||
stderr: String::new(),
|
||||
raw_output_path: None,
|
||||
interrupted: false,
|
||||
is_image: None,
|
||||
background_task_id: Some(child.id().to_string()),
|
||||
backgrounded_by_user: Some(false),
|
||||
assistant_auto_backgrounded: Some(false),
|
||||
dangerously_disable_sandbox: input.dangerously_disable_sandbox,
|
||||
return_code_interpretation: None,
|
||||
no_output_expected: Some(true),
|
||||
structured_content: None,
|
||||
persisted_output_path: None,
|
||||
persisted_output_size: None,
|
||||
sandbox_status: Some(sandbox_status),
|
||||
});
|
||||
}
|
||||
|
||||
let runtime = Builder::new_current_thread().enable_all().build()?;
|
||||
runtime.block_on(execute_bash_async(input, sandbox_status, cwd))
|
||||
}
|
||||
|
||||
async fn execute_bash_async(
|
||||
input: BashCommandInput,
|
||||
sandbox_status: SandboxStatus,
|
||||
cwd: std::path::PathBuf,
|
||||
) -> io::Result<BashCommandOutput> {
|
||||
let mut command = prepare_tokio_command(&input.command, &cwd, &sandbox_status, true);
|
||||
|
||||
let output_result = if let Some(timeout_ms) = input.timeout {
|
||||
match timeout(Duration::from_millis(timeout_ms), command.output()).await {
|
||||
Ok(result) => (result?, false),
|
||||
Err(_) => {
|
||||
return Ok(BashCommandOutput {
|
||||
stdout: String::new(),
|
||||
stderr: format!("Command exceeded timeout of {timeout_ms} ms"),
|
||||
raw_output_path: None,
|
||||
interrupted: true,
|
||||
is_image: None,
|
||||
background_task_id: None,
|
||||
backgrounded_by_user: None,
|
||||
assistant_auto_backgrounded: None,
|
||||
dangerously_disable_sandbox: input.dangerously_disable_sandbox,
|
||||
return_code_interpretation: Some(String::from("timeout")),
|
||||
no_output_expected: Some(true),
|
||||
structured_content: None,
|
||||
persisted_output_path: None,
|
||||
persisted_output_size: None,
|
||||
sandbox_status: Some(sandbox_status),
|
||||
});
|
||||
}
|
||||
}
|
||||
} else {
|
||||
(command.output().await?, false)
|
||||
};
|
||||
|
||||
let (output, interrupted) = output_result;
|
||||
let stdout = truncate_output(&String::from_utf8_lossy(&output.stdout));
|
||||
let stderr = truncate_output(&String::from_utf8_lossy(&output.stderr));
|
||||
let no_output_expected = Some(stdout.trim().is_empty() && stderr.trim().is_empty());
|
||||
let return_code_interpretation = output.status.code().and_then(|code| {
|
||||
if code == 0 {
|
||||
None
|
||||
} else {
|
||||
Some(format!("exit_code:{code}"))
|
||||
}
|
||||
});
|
||||
|
||||
Ok(BashCommandOutput {
|
||||
stdout,
|
||||
stderr,
|
||||
raw_output_path: None,
|
||||
interrupted,
|
||||
is_image: None,
|
||||
background_task_id: None,
|
||||
backgrounded_by_user: None,
|
||||
assistant_auto_backgrounded: None,
|
||||
dangerously_disable_sandbox: input.dangerously_disable_sandbox,
|
||||
return_code_interpretation,
|
||||
no_output_expected,
|
||||
structured_content: None,
|
||||
persisted_output_path: None,
|
||||
persisted_output_size: None,
|
||||
sandbox_status: Some(sandbox_status),
|
||||
})
|
||||
}
|
||||
|
||||
fn sandbox_status_for_input(input: &BashCommandInput, cwd: &std::path::Path) -> SandboxStatus {
|
||||
let config = ConfigLoader::default_for(cwd).load().map_or_else(
|
||||
|_| SandboxConfig::default(),
|
||||
|runtime_config| runtime_config.sandbox().clone(),
|
||||
);
|
||||
let request = config.resolve_request(
|
||||
input.dangerously_disable_sandbox.map(|disabled| !disabled),
|
||||
input.namespace_restrictions,
|
||||
input.isolate_network,
|
||||
input.filesystem_mode,
|
||||
input.allowed_mounts.clone(),
|
||||
);
|
||||
resolve_sandbox_status_for_request(&request, cwd)
|
||||
}
|
||||
|
||||
fn prepare_command(
|
||||
command: &str,
|
||||
cwd: &std::path::Path,
|
||||
sandbox_status: &SandboxStatus,
|
||||
create_dirs: bool,
|
||||
) -> Command {
|
||||
if create_dirs {
|
||||
prepare_sandbox_dirs(cwd);
|
||||
}
|
||||
|
||||
if let Some(launcher) = build_linux_sandbox_command(command, cwd, sandbox_status) {
|
||||
let mut prepared = Command::new(launcher.program);
|
||||
prepared.args(launcher.args);
|
||||
prepared.current_dir(cwd);
|
||||
prepared.envs(launcher.env);
|
||||
return prepared;
|
||||
}
|
||||
|
||||
let prepared = if cfg!(target_os = "windows") && !sh_exists() {
|
||||
let mut p = Command::new("cmd");
|
||||
p.arg("/C").arg(command).current_dir(cwd);
|
||||
p
|
||||
} else {
|
||||
let mut p = Command::new("sh");
|
||||
p.arg("-lc").arg(command).current_dir(cwd);
|
||||
if sandbox_status.filesystem_active {
|
||||
p.env("HOME", cwd.join(".sandbox-home"));
|
||||
p.env("TMPDIR", cwd.join(".sandbox-tmp"));
|
||||
}
|
||||
p
|
||||
};
|
||||
prepared
|
||||
}
|
||||
|
||||
fn sh_exists() -> bool {
|
||||
#[cfg(windows)]
|
||||
{
|
||||
which::which("sh").is_ok()
|
||||
}
|
||||
#[cfg(not(windows))]
|
||||
{
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
fn prepare_tokio_command(
|
||||
command: &str,
|
||||
cwd: &std::path::Path,
|
||||
sandbox_status: &SandboxStatus,
|
||||
create_dirs: bool,
|
||||
) -> TokioCommand {
|
||||
if create_dirs {
|
||||
prepare_sandbox_dirs(cwd);
|
||||
}
|
||||
|
||||
if let Some(launcher) = build_linux_sandbox_command(command, cwd, sandbox_status) {
|
||||
let mut prepared = TokioCommand::new(launcher.program);
|
||||
prepared.args(launcher.args);
|
||||
prepared.current_dir(cwd);
|
||||
prepared.envs(launcher.env);
|
||||
return prepared;
|
||||
}
|
||||
|
||||
let prepared = if cfg!(target_os = "windows") && !sh_exists() {
|
||||
let mut p = TokioCommand::new("cmd");
|
||||
p.arg("/C").arg(command).current_dir(cwd);
|
||||
p
|
||||
} else {
|
||||
let mut p = TokioCommand::new("sh");
|
||||
p.arg("-lc").arg(command).current_dir(cwd);
|
||||
if sandbox_status.filesystem_active {
|
||||
p.env("HOME", cwd.join(".sandbox-home"));
|
||||
p.env("TMPDIR", cwd.join(".sandbox-tmp"));
|
||||
}
|
||||
p
|
||||
};
|
||||
prepared
|
||||
}
|
||||
|
||||
fn prepare_sandbox_dirs(cwd: &std::path::Path) {
|
||||
let _ = std::fs::create_dir_all(cwd.join(".sandbox-home"));
|
||||
let _ = std::fs::create_dir_all(cwd.join(".sandbox-tmp"));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{execute_bash, BashCommandInput};
|
||||
use crate::sandbox::FilesystemIsolationMode;
|
||||
|
||||
#[test]
|
||||
fn executes_simple_command() {
|
||||
let output = execute_bash(BashCommandInput {
|
||||
command: String::from("printf 'hello'"),
|
||||
timeout: Some(1_000),
|
||||
description: None,
|
||||
run_in_background: Some(false),
|
||||
dangerously_disable_sandbox: Some(false),
|
||||
namespace_restrictions: Some(false),
|
||||
isolate_network: Some(false),
|
||||
filesystem_mode: Some(FilesystemIsolationMode::WorkspaceOnly),
|
||||
allowed_mounts: None,
|
||||
})
|
||||
.expect("bash command should execute");
|
||||
|
||||
assert_eq!(output.stdout, "hello");
|
||||
assert!(!output.interrupted);
|
||||
assert!(output.sandbox_status.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disables_sandbox_when_requested() {
|
||||
let output = execute_bash(BashCommandInput {
|
||||
command: String::from("printf 'hello'"),
|
||||
timeout: Some(1_000),
|
||||
description: None,
|
||||
run_in_background: Some(false),
|
||||
dangerously_disable_sandbox: Some(true),
|
||||
namespace_restrictions: None,
|
||||
isolate_network: None,
|
||||
filesystem_mode: None,
|
||||
allowed_mounts: None,
|
||||
})
|
||||
.expect("bash command should execute");
|
||||
|
||||
assert!(!output.sandbox_status.expect("sandbox status").enabled);
|
||||
}
|
||||
}
|
||||
|
||||
/// Maximum output bytes before truncation (16 KiB, matching upstream).
|
||||
const MAX_OUTPUT_BYTES: usize = 16_384;
|
||||
|
||||
/// Truncate output to `MAX_OUTPUT_BYTES`, appending a marker when trimmed.
|
||||
fn truncate_output(s: &str) -> String {
|
||||
if s.len() <= MAX_OUTPUT_BYTES {
|
||||
return s.to_string();
|
||||
}
|
||||
// Find the last valid UTF-8 boundary at or before MAX_OUTPUT_BYTES
|
||||
let mut end = MAX_OUTPUT_BYTES;
|
||||
while end > 0 && !s.is_char_boundary(end) {
|
||||
end -= 1;
|
||||
}
|
||||
let mut truncated = s[..end].to_string();
|
||||
truncated.push_str("\n\n[output truncated — exceeded 16384 bytes]");
|
||||
truncated
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod truncation_tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn short_output_unchanged() {
|
||||
let s = "hello world";
|
||||
assert_eq!(truncate_output(s), s);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn long_output_truncated() {
|
||||
let s = "x".repeat(20_000);
|
||||
let result = truncate_output(&s);
|
||||
assert!(result.len() < 20_000);
|
||||
assert!(result.ends_with("[output truncated — exceeded 16384 bytes]"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exact_boundary_unchanged() {
|
||||
let s = "a".repeat(MAX_OUTPUT_BYTES);
|
||||
assert_eq!(truncate_output(&s), s);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn one_over_boundary_truncated() {
|
||||
let s = "a".repeat(MAX_OUTPUT_BYTES + 1);
|
||||
let result = truncate_output(&s);
|
||||
assert!(result.contains("[output truncated"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum BootstrapPhase {
|
||||
CliEntry,
|
||||
FastPathVersion,
|
||||
StartupProfiler,
|
||||
SystemPromptFastPath,
|
||||
ChromeMcpFastPath,
|
||||
DaemonWorkerFastPath,
|
||||
BridgeFastPath,
|
||||
DaemonFastPath,
|
||||
BackgroundSessionFastPath,
|
||||
TemplateFastPath,
|
||||
EnvironmentRunnerFastPath,
|
||||
MainRuntime,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct BootstrapPlan {
|
||||
phases: Vec<BootstrapPhase>,
|
||||
}
|
||||
|
||||
impl BootstrapPlan {
|
||||
#[must_use]
|
||||
pub fn claude_code_default() -> Self {
|
||||
Self::from_phases(vec![
|
||||
BootstrapPhase::CliEntry,
|
||||
BootstrapPhase::FastPathVersion,
|
||||
BootstrapPhase::StartupProfiler,
|
||||
BootstrapPhase::SystemPromptFastPath,
|
||||
BootstrapPhase::ChromeMcpFastPath,
|
||||
BootstrapPhase::DaemonWorkerFastPath,
|
||||
BootstrapPhase::BridgeFastPath,
|
||||
BootstrapPhase::DaemonFastPath,
|
||||
BootstrapPhase::BackgroundSessionFastPath,
|
||||
BootstrapPhase::TemplateFastPath,
|
||||
BootstrapPhase::EnvironmentRunnerFastPath,
|
||||
BootstrapPhase::MainRuntime,
|
||||
])
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn from_phases(phases: Vec<BootstrapPhase>) -> Self {
|
||||
let mut deduped = Vec::new();
|
||||
for phase in phases {
|
||||
if !deduped.contains(&phase) {
|
||||
deduped.push(phase);
|
||||
}
|
||||
}
|
||||
Self { phases: deduped }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn phases(&self) -> &[BootstrapPhase] {
|
||||
&self.phases
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{BootstrapPhase, BootstrapPlan};
|
||||
|
||||
#[test]
|
||||
fn from_phases_deduplicates_while_preserving_order() {
|
||||
// given
|
||||
let phases = vec![
|
||||
BootstrapPhase::CliEntry,
|
||||
BootstrapPhase::FastPathVersion,
|
||||
BootstrapPhase::CliEntry,
|
||||
BootstrapPhase::MainRuntime,
|
||||
BootstrapPhase::FastPathVersion,
|
||||
];
|
||||
|
||||
// when
|
||||
let plan = BootstrapPlan::from_phases(phases);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
plan.phases(),
|
||||
&[
|
||||
BootstrapPhase::CliEntry,
|
||||
BootstrapPhase::FastPathVersion,
|
||||
BootstrapPhase::MainRuntime,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_code_default_covers_each_phase_once() {
|
||||
// given
|
||||
let expected = [
|
||||
BootstrapPhase::CliEntry,
|
||||
BootstrapPhase::FastPathVersion,
|
||||
BootstrapPhase::StartupProfiler,
|
||||
BootstrapPhase::SystemPromptFastPath,
|
||||
BootstrapPhase::ChromeMcpFastPath,
|
||||
BootstrapPhase::DaemonWorkerFastPath,
|
||||
BootstrapPhase::BridgeFastPath,
|
||||
BootstrapPhase::DaemonFastPath,
|
||||
BootstrapPhase::BackgroundSessionFastPath,
|
||||
BootstrapPhase::TemplateFastPath,
|
||||
BootstrapPhase::EnvironmentRunnerFastPath,
|
||||
BootstrapPhase::MainRuntime,
|
||||
];
|
||||
|
||||
// when
|
||||
let plan = BootstrapPlan::claude_code_default();
|
||||
|
||||
// then
|
||||
assert_eq!(plan.phases(), &expected);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BranchLockIntent {
|
||||
#[serde(rename = "laneId")]
|
||||
pub lane_id: String,
|
||||
pub branch: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub worktree: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub modules: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct BranchLockCollision {
|
||||
pub branch: String,
|
||||
pub module: String,
|
||||
#[serde(rename = "laneIds")]
|
||||
pub lane_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn detect_branch_lock_collisions(intents: &[BranchLockIntent]) -> Vec<BranchLockCollision> {
|
||||
let mut collisions = Vec::new();
|
||||
|
||||
for (index, left) in intents.iter().enumerate() {
|
||||
for right in &intents[index + 1..] {
|
||||
if left.branch != right.branch {
|
||||
continue;
|
||||
}
|
||||
for module in overlapping_modules(&left.modules, &right.modules) {
|
||||
collisions.push(BranchLockCollision {
|
||||
branch: left.branch.clone(),
|
||||
module,
|
||||
lane_ids: vec![left.lane_id.clone(), right.lane_id.clone()],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
collisions.sort_by(|a, b| {
|
||||
a.branch
|
||||
.cmp(&b.branch)
|
||||
.then(a.module.cmp(&b.module))
|
||||
.then(a.lane_ids.cmp(&b.lane_ids))
|
||||
});
|
||||
collisions.dedup();
|
||||
collisions
|
||||
}
|
||||
|
||||
fn overlapping_modules(left: &[String], right: &[String]) -> Vec<String> {
|
||||
let mut overlaps = Vec::new();
|
||||
for left_module in left {
|
||||
for right_module in right {
|
||||
if modules_overlap(left_module, right_module) {
|
||||
overlaps.push(shared_scope(left_module, right_module));
|
||||
}
|
||||
}
|
||||
}
|
||||
overlaps.sort();
|
||||
overlaps.dedup();
|
||||
overlaps
|
||||
}
|
||||
|
||||
fn modules_overlap(left: &str, right: &str) -> bool {
|
||||
left == right
|
||||
|| left.starts_with(&format!("{right}/"))
|
||||
|| right.starts_with(&format!("{left}/"))
|
||||
}
|
||||
|
||||
fn shared_scope(left: &str, right: &str) -> String {
|
||||
if left.starts_with(&format!("{right}/")) || left == right {
|
||||
right.to_string()
|
||||
} else {
|
||||
left.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{detect_branch_lock_collisions, BranchLockIntent};
|
||||
|
||||
#[test]
|
||||
fn detects_same_branch_same_module_collisions() {
|
||||
let collisions = detect_branch_lock_collisions(&[
|
||||
BranchLockIntent {
|
||||
lane_id: "lane-a".to_string(),
|
||||
branch: "feature/lock".to_string(),
|
||||
worktree: Some("wt-a".to_string()),
|
||||
modules: vec!["runtime/mcp".to_string()],
|
||||
},
|
||||
BranchLockIntent {
|
||||
lane_id: "lane-b".to_string(),
|
||||
branch: "feature/lock".to_string(),
|
||||
worktree: Some("wt-b".to_string()),
|
||||
modules: vec!["runtime/mcp".to_string()],
|
||||
},
|
||||
]);
|
||||
|
||||
assert_eq!(collisions.len(), 1);
|
||||
assert_eq!(collisions[0].branch, "feature/lock");
|
||||
assert_eq!(collisions[0].module, "runtime/mcp");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_nested_module_scope_collisions() {
|
||||
let collisions = detect_branch_lock_collisions(&[
|
||||
BranchLockIntent {
|
||||
lane_id: "lane-a".to_string(),
|
||||
branch: "feature/lock".to_string(),
|
||||
worktree: None,
|
||||
modules: vec!["runtime".to_string()],
|
||||
},
|
||||
BranchLockIntent {
|
||||
lane_id: "lane-b".to_string(),
|
||||
branch: "feature/lock".to_string(),
|
||||
worktree: None,
|
||||
modules: vec!["runtime/mcp".to_string()],
|
||||
},
|
||||
]);
|
||||
|
||||
assert_eq!(collisions[0].module, "runtime");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_different_branches() {
|
||||
let collisions = detect_branch_lock_collisions(&[
|
||||
BranchLockIntent {
|
||||
lane_id: "lane-a".to_string(),
|
||||
branch: "feature/a".to_string(),
|
||||
worktree: None,
|
||||
modules: vec!["runtime/mcp".to_string()],
|
||||
},
|
||||
BranchLockIntent {
|
||||
lane_id: "lane-b".to_string(),
|
||||
branch: "feature/b".to_string(),
|
||||
worktree: None,
|
||||
modules: vec!["runtime/mcp".to_string()],
|
||||
},
|
||||
]);
|
||||
|
||||
assert!(collisions.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,827 @@
|
||||
use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
|
||||
|
||||
const COMPACT_CONTINUATION_PREAMBLE: &str =
|
||||
"This session is being continued from a previous conversation that ran out of context. The summary below covers the earlier portion of the conversation.\n\n";
|
||||
const COMPACT_RECENT_MESSAGES_NOTE: &str = "Recent messages are preserved verbatim.";
|
||||
const COMPACT_DIRECT_RESUME_INSTRUCTION: &str = "Continue the conversation from where it left off without asking the user any further questions. Resume directly — do not acknowledge the summary, do not recap what was happening, and do not preface with continuation text.";
|
||||
|
||||
/// Thresholds controlling when and how a session is compacted.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct CompactionConfig {
|
||||
pub preserve_recent_messages: usize,
|
||||
pub max_estimated_tokens: usize,
|
||||
}
|
||||
|
||||
impl Default for CompactionConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
preserve_recent_messages: 4,
|
||||
max_estimated_tokens: 10_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of compacting a session into a summary plus preserved tail messages.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CompactionResult {
|
||||
pub summary: String,
|
||||
pub formatted_summary: String,
|
||||
pub compacted_session: Session,
|
||||
pub removed_message_count: usize,
|
||||
}
|
||||
|
||||
/// Roughly estimates the token footprint of the current session transcript.
|
||||
#[must_use]
|
||||
pub fn estimate_session_tokens(session: &Session) -> usize {
|
||||
session.messages.iter().map(estimate_message_tokens).sum()
|
||||
}
|
||||
|
||||
/// Returns `true` when the session exceeds the configured compaction budget.
|
||||
#[must_use]
|
||||
pub fn should_compact(session: &Session, config: CompactionConfig) -> bool {
|
||||
let start = compacted_summary_prefix_len(session);
|
||||
let compactable = &session.messages[start..];
|
||||
|
||||
compactable.len() > config.preserve_recent_messages
|
||||
&& compactable
|
||||
.iter()
|
||||
.map(estimate_message_tokens)
|
||||
.sum::<usize>()
|
||||
>= config.max_estimated_tokens
|
||||
}
|
||||
|
||||
/// Normalizes a compaction summary into user-facing continuation text.
|
||||
#[must_use]
|
||||
pub fn format_compact_summary(summary: &str) -> String {
|
||||
let without_analysis = strip_tag_block(summary, "analysis");
|
||||
let formatted = if let Some(content) = extract_tag_block(&without_analysis, "summary") {
|
||||
without_analysis.replace(
|
||||
&format!("<summary>{content}</summary>"),
|
||||
&format!("Summary:\n{}", content.trim()),
|
||||
)
|
||||
} else {
|
||||
without_analysis
|
||||
};
|
||||
|
||||
collapse_blank_lines(&formatted).trim().to_string()
|
||||
}
|
||||
|
||||
/// Builds the synthetic system message used after session compaction.
|
||||
#[must_use]
|
||||
pub fn get_compact_continuation_message(
|
||||
summary: &str,
|
||||
suppress_follow_up_questions: bool,
|
||||
recent_messages_preserved: bool,
|
||||
) -> String {
|
||||
let mut base = format!(
|
||||
"{COMPACT_CONTINUATION_PREAMBLE}{}",
|
||||
format_compact_summary(summary)
|
||||
);
|
||||
|
||||
if recent_messages_preserved {
|
||||
base.push_str("\n\n");
|
||||
base.push_str(COMPACT_RECENT_MESSAGES_NOTE);
|
||||
}
|
||||
|
||||
if suppress_follow_up_questions {
|
||||
base.push('\n');
|
||||
base.push_str(COMPACT_DIRECT_RESUME_INSTRUCTION);
|
||||
}
|
||||
|
||||
base
|
||||
}
|
||||
|
||||
/// Compacts a session by summarizing older messages and preserving the recent tail.
|
||||
#[must_use]
|
||||
pub fn compact_session(session: &Session, config: CompactionConfig) -> CompactionResult {
|
||||
if !should_compact(session, config) {
|
||||
return CompactionResult {
|
||||
summary: String::new(),
|
||||
formatted_summary: String::new(),
|
||||
compacted_session: session.clone(),
|
||||
removed_message_count: 0,
|
||||
};
|
||||
}
|
||||
|
||||
let existing_summary = session
|
||||
.messages
|
||||
.first()
|
||||
.and_then(extract_existing_compacted_summary);
|
||||
let compacted_prefix_len = usize::from(existing_summary.is_some());
|
||||
let raw_keep_from = session
|
||||
.messages
|
||||
.len()
|
||||
.saturating_sub(config.preserve_recent_messages);
|
||||
// Ensure we do not split a tool-use / tool-result pair at the compaction
|
||||
// boundary. If the first preserved message is a user message whose first
|
||||
// block is a ToolResult, the assistant message with the matching ToolUse
|
||||
// was slated for removal — that produces an orphaned tool role message on
|
||||
// the OpenAI-compat path (400: tool message must follow assistant with
|
||||
// tool_calls). Walk the boundary back until we start at a safe point.
|
||||
let keep_from = {
|
||||
let mut k = raw_keep_from;
|
||||
// If the first preserved message is a tool-result turn, ensure its
|
||||
// paired assistant tool-use turn is preserved too. Without this fix,
|
||||
// the OpenAI-compat adapter sends an orphaned 'tool' role message
|
||||
// with no preceding assistant 'tool_calls', which providers reject
|
||||
// with a 400. We walk back only if the immediately preceding message
|
||||
// is NOT an assistant message that contains a ToolUse block (i.e. the
|
||||
// pair is actually broken at the boundary).
|
||||
loop {
|
||||
if k == 0 || k <= compacted_prefix_len {
|
||||
break;
|
||||
}
|
||||
let first_preserved = &session.messages[k];
|
||||
let starts_with_tool_result = first_preserved
|
||||
.blocks
|
||||
.first()
|
||||
.map(|b| matches!(b, ContentBlock::ToolResult { .. }))
|
||||
.unwrap_or(false);
|
||||
if !starts_with_tool_result {
|
||||
break;
|
||||
}
|
||||
// Check the message just before the current boundary.
|
||||
let preceding = &session.messages[k - 1];
|
||||
let preceding_has_tool_use = preceding
|
||||
.blocks
|
||||
.iter()
|
||||
.any(|b| matches!(b, ContentBlock::ToolUse { .. }));
|
||||
if preceding_has_tool_use {
|
||||
// Pair is intact — walk back one more to include the assistant turn.
|
||||
k = k.saturating_sub(1);
|
||||
break;
|
||||
}
|
||||
// Preceding message has no ToolUse but we have a ToolResult —
|
||||
// this is already an orphaned pair; walk back to try to fix it.
|
||||
k = k.saturating_sub(1);
|
||||
}
|
||||
k
|
||||
};
|
||||
let removed = &session.messages[compacted_prefix_len..keep_from];
|
||||
let preserved = session.messages[keep_from..].to_vec();
|
||||
let summary =
|
||||
merge_compact_summaries(existing_summary.as_deref(), &summarize_messages(removed));
|
||||
let formatted_summary = format_compact_summary(&summary);
|
||||
let continuation = get_compact_continuation_message(&summary, true, !preserved.is_empty());
|
||||
|
||||
let mut compacted_messages = vec![ConversationMessage {
|
||||
role: MessageRole::System,
|
||||
blocks: vec![ContentBlock::Text { text: continuation }],
|
||||
usage: None,
|
||||
}];
|
||||
compacted_messages.extend(preserved);
|
||||
|
||||
let mut compacted_session = session.clone();
|
||||
compacted_session.messages = compacted_messages;
|
||||
compacted_session.record_compaction(summary.clone(), removed.len());
|
||||
|
||||
CompactionResult {
|
||||
summary,
|
||||
formatted_summary,
|
||||
compacted_session,
|
||||
removed_message_count: removed.len(),
|
||||
}
|
||||
}
|
||||
|
||||
fn compacted_summary_prefix_len(session: &Session) -> usize {
|
||||
usize::from(
|
||||
session
|
||||
.messages
|
||||
.first()
|
||||
.and_then(extract_existing_compacted_summary)
|
||||
.is_some(),
|
||||
)
|
||||
}
|
||||
|
||||
fn summarize_messages(messages: &[ConversationMessage]) -> String {
|
||||
let user_messages = messages
|
||||
.iter()
|
||||
.filter(|message| message.role == MessageRole::User)
|
||||
.count();
|
||||
let assistant_messages = messages
|
||||
.iter()
|
||||
.filter(|message| message.role == MessageRole::Assistant)
|
||||
.count();
|
||||
let tool_messages = messages
|
||||
.iter()
|
||||
.filter(|message| message.role == MessageRole::Tool)
|
||||
.count();
|
||||
|
||||
let mut tool_names = messages
|
||||
.iter()
|
||||
.flat_map(|message| message.blocks.iter())
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::ToolUse { name, .. } => Some(name.as_str()),
|
||||
ContentBlock::ToolResult { tool_name, .. } => Some(tool_name.as_str()),
|
||||
ContentBlock::Text { .. } => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
tool_names.sort_unstable();
|
||||
tool_names.dedup();
|
||||
|
||||
let mut lines = vec![
|
||||
"<summary>".to_string(),
|
||||
"Conversation summary:".to_string(),
|
||||
format!(
|
||||
"- Scope: {} earlier messages compacted (user={}, assistant={}, tool={}).",
|
||||
messages.len(),
|
||||
user_messages,
|
||||
assistant_messages,
|
||||
tool_messages
|
||||
),
|
||||
];
|
||||
|
||||
if !tool_names.is_empty() {
|
||||
lines.push(format!("- Tools mentioned: {}.", tool_names.join(", ")));
|
||||
}
|
||||
|
||||
let recent_user_requests = collect_recent_role_summaries(messages, MessageRole::User, 3);
|
||||
if !recent_user_requests.is_empty() {
|
||||
lines.push("- Recent user requests:".to_string());
|
||||
lines.extend(
|
||||
recent_user_requests
|
||||
.into_iter()
|
||||
.map(|request| format!(" - {request}")),
|
||||
);
|
||||
}
|
||||
|
||||
let pending_work = infer_pending_work(messages);
|
||||
if !pending_work.is_empty() {
|
||||
lines.push("- Pending work:".to_string());
|
||||
lines.extend(pending_work.into_iter().map(|item| format!(" - {item}")));
|
||||
}
|
||||
|
||||
let key_files = collect_key_files(messages);
|
||||
if !key_files.is_empty() {
|
||||
lines.push(format!("- Key files referenced: {}.", key_files.join(", ")));
|
||||
}
|
||||
|
||||
if let Some(current_work) = infer_current_work(messages) {
|
||||
lines.push(format!("- Current work: {current_work}"));
|
||||
}
|
||||
|
||||
lines.push("- Key timeline:".to_string());
|
||||
for message in messages {
|
||||
let role = match message.role {
|
||||
MessageRole::System => "system",
|
||||
MessageRole::User => "user",
|
||||
MessageRole::Assistant => "assistant",
|
||||
MessageRole::Tool => "tool",
|
||||
};
|
||||
let content = message
|
||||
.blocks
|
||||
.iter()
|
||||
.map(summarize_block)
|
||||
.collect::<Vec<_>>()
|
||||
.join(" | ");
|
||||
lines.push(format!(" - {role}: {content}"));
|
||||
}
|
||||
lines.push("</summary>".to_string());
|
||||
lines.join("\n")
|
||||
}
|
||||
|
||||
fn merge_compact_summaries(existing_summary: Option<&str>, new_summary: &str) -> String {
|
||||
let Some(existing_summary) = existing_summary else {
|
||||
return new_summary.to_string();
|
||||
};
|
||||
|
||||
let previous_highlights = extract_summary_highlights(existing_summary);
|
||||
let new_formatted_summary = format_compact_summary(new_summary);
|
||||
let new_highlights = extract_summary_highlights(&new_formatted_summary);
|
||||
let new_timeline = extract_summary_timeline(&new_formatted_summary);
|
||||
|
||||
let mut lines = vec!["<summary>".to_string(), "Conversation summary:".to_string()];
|
||||
|
||||
if !previous_highlights.is_empty() {
|
||||
lines.push("- Previously compacted context:".to_string());
|
||||
lines.extend(
|
||||
previous_highlights
|
||||
.into_iter()
|
||||
.map(|line| format!(" {line}")),
|
||||
);
|
||||
}
|
||||
|
||||
if !new_highlights.is_empty() {
|
||||
lines.push("- Newly compacted context:".to_string());
|
||||
lines.extend(new_highlights.into_iter().map(|line| format!(" {line}")));
|
||||
}
|
||||
|
||||
if !new_timeline.is_empty() {
|
||||
lines.push("- Key timeline:".to_string());
|
||||
lines.extend(new_timeline.into_iter().map(|line| format!(" {line}")));
|
||||
}
|
||||
|
||||
lines.push("</summary>".to_string());
|
||||
lines.join("\n")
|
||||
}
|
||||
|
||||
fn summarize_block(block: &ContentBlock) -> String {
|
||||
let raw = match block {
|
||||
ContentBlock::Text { text } => text.clone(),
|
||||
ContentBlock::ToolUse { name, input, .. } => format!("tool_use {name}({input})"),
|
||||
ContentBlock::ToolResult {
|
||||
tool_name,
|
||||
output,
|
||||
is_error,
|
||||
..
|
||||
} => format!(
|
||||
"tool_result {tool_name}: {}{output}",
|
||||
if *is_error { "error " } else { "" }
|
||||
),
|
||||
};
|
||||
truncate_summary(&raw, 160)
|
||||
}
|
||||
|
||||
fn collect_recent_role_summaries(
|
||||
messages: &[ConversationMessage],
|
||||
role: MessageRole,
|
||||
limit: usize,
|
||||
) -> Vec<String> {
|
||||
messages
|
||||
.iter()
|
||||
.filter(|message| message.role == role)
|
||||
.rev()
|
||||
.filter_map(|message| first_text_block(message))
|
||||
.take(limit)
|
||||
.map(|text| truncate_summary(text, 160))
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn infer_pending_work(messages: &[ConversationMessage]) -> Vec<String> {
|
||||
messages
|
||||
.iter()
|
||||
.rev()
|
||||
.filter_map(first_text_block)
|
||||
.filter(|text| {
|
||||
let lowered = text.to_ascii_lowercase();
|
||||
lowered.contains("todo")
|
||||
|| lowered.contains("next")
|
||||
|| lowered.contains("pending")
|
||||
|| lowered.contains("follow up")
|
||||
|| lowered.contains("remaining")
|
||||
})
|
||||
.take(3)
|
||||
.map(|text| truncate_summary(text, 160))
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn collect_key_files(messages: &[ConversationMessage]) -> Vec<String> {
|
||||
let mut files = messages
|
||||
.iter()
|
||||
.flat_map(|message| message.blocks.iter())
|
||||
.map(|block| match block {
|
||||
ContentBlock::Text { text } => text.as_str(),
|
||||
ContentBlock::ToolUse { input, .. } => input.as_str(),
|
||||
ContentBlock::ToolResult { output, .. } => output.as_str(),
|
||||
})
|
||||
.flat_map(extract_file_candidates)
|
||||
.collect::<Vec<_>>();
|
||||
files.sort();
|
||||
files.dedup();
|
||||
files.into_iter().take(8).collect()
|
||||
}
|
||||
|
||||
fn infer_current_work(messages: &[ConversationMessage]) -> Option<String> {
|
||||
messages
|
||||
.iter()
|
||||
.rev()
|
||||
.filter_map(first_text_block)
|
||||
.find(|text| !text.trim().is_empty())
|
||||
.map(|text| truncate_summary(text, 200))
|
||||
}
|
||||
|
||||
fn first_text_block(message: &ConversationMessage) -> Option<&str> {
|
||||
message.blocks.iter().find_map(|block| match block {
|
||||
ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.as_str()),
|
||||
ContentBlock::ToolUse { .. }
|
||||
| ContentBlock::ToolResult { .. }
|
||||
| ContentBlock::Text { .. } => None,
|
||||
})
|
||||
}
|
||||
|
||||
fn has_interesting_extension(candidate: &str) -> bool {
|
||||
std::path::Path::new(candidate)
|
||||
.extension()
|
||||
.and_then(|extension| extension.to_str())
|
||||
.is_some_and(|extension| {
|
||||
["rs", "ts", "tsx", "js", "json", "md"]
|
||||
.iter()
|
||||
.any(|expected| extension.eq_ignore_ascii_case(expected))
|
||||
})
|
||||
}
|
||||
|
||||
fn extract_file_candidates(content: &str) -> Vec<String> {
|
||||
content
|
||||
.split_whitespace()
|
||||
.filter_map(|token| {
|
||||
let candidate = token.trim_matches(|char: char| {
|
||||
matches!(char, ',' | '.' | ':' | ';' | ')' | '(' | '"' | '\'' | '`')
|
||||
});
|
||||
if candidate.contains('/') && has_interesting_extension(candidate) {
|
||||
Some(candidate.to_string())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn truncate_summary(content: &str, max_chars: usize) -> String {
|
||||
if content.chars().count() <= max_chars {
|
||||
return content.to_string();
|
||||
}
|
||||
let mut truncated = content.chars().take(max_chars).collect::<String>();
|
||||
truncated.push('…');
|
||||
truncated
|
||||
}
|
||||
|
||||
fn estimate_message_tokens(message: &ConversationMessage) -> usize {
|
||||
message
|
||||
.blocks
|
||||
.iter()
|
||||
.map(|block| match block {
|
||||
ContentBlock::Text { text } => text.len() / 4 + 1,
|
||||
ContentBlock::ToolUse { name, input, .. } => (name.len() + input.len()) / 4 + 1,
|
||||
ContentBlock::ToolResult {
|
||||
tool_name, output, ..
|
||||
} => (tool_name.len() + output.len()) / 4 + 1,
|
||||
})
|
||||
.sum()
|
||||
}
|
||||
|
||||
fn extract_tag_block(content: &str, tag: &str) -> Option<String> {
|
||||
let start = format!("<{tag}>");
|
||||
let end = format!("</{tag}>");
|
||||
let start_index = content.find(&start)? + start.len();
|
||||
let end_index = content[start_index..].find(&end)? + start_index;
|
||||
Some(content[start_index..end_index].to_string())
|
||||
}
|
||||
|
||||
fn strip_tag_block(content: &str, tag: &str) -> String {
|
||||
let start = format!("<{tag}>");
|
||||
let end = format!("</{tag}>");
|
||||
if let (Some(start_index), Some(end_index_rel)) = (content.find(&start), content.find(&end)) {
|
||||
let end_index = end_index_rel + end.len();
|
||||
let mut stripped = String::new();
|
||||
stripped.push_str(&content[..start_index]);
|
||||
stripped.push_str(&content[end_index..]);
|
||||
stripped
|
||||
} else {
|
||||
content.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn collapse_blank_lines(content: &str) -> String {
|
||||
let mut result = String::new();
|
||||
let mut last_blank = false;
|
||||
for line in content.lines() {
|
||||
let is_blank = line.trim().is_empty();
|
||||
if is_blank && last_blank {
|
||||
continue;
|
||||
}
|
||||
result.push_str(line);
|
||||
result.push('\n');
|
||||
last_blank = is_blank;
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
fn extract_existing_compacted_summary(message: &ConversationMessage) -> Option<String> {
|
||||
if message.role != MessageRole::System {
|
||||
return None;
|
||||
}
|
||||
|
||||
let text = first_text_block(message)?;
|
||||
let summary = text.strip_prefix(COMPACT_CONTINUATION_PREAMBLE)?;
|
||||
let summary = summary
|
||||
.split_once(&format!("\n\n{COMPACT_RECENT_MESSAGES_NOTE}"))
|
||||
.map_or(summary, |(value, _)| value);
|
||||
let summary = summary
|
||||
.split_once(&format!("\n{COMPACT_DIRECT_RESUME_INSTRUCTION}"))
|
||||
.map_or(summary, |(value, _)| value);
|
||||
Some(summary.trim().to_string())
|
||||
}
|
||||
|
||||
fn extract_summary_highlights(summary: &str) -> Vec<String> {
|
||||
let mut lines = Vec::new();
|
||||
let mut in_timeline = false;
|
||||
|
||||
for line in format_compact_summary(summary).lines() {
|
||||
let trimmed = line.trim_end();
|
||||
if trimmed.is_empty() || trimmed == "Summary:" || trimmed == "Conversation summary:" {
|
||||
continue;
|
||||
}
|
||||
if trimmed == "- Key timeline:" {
|
||||
in_timeline = true;
|
||||
continue;
|
||||
}
|
||||
if in_timeline {
|
||||
continue;
|
||||
}
|
||||
lines.push(trimmed.to_string());
|
||||
}
|
||||
|
||||
lines
|
||||
}
|
||||
|
||||
fn extract_summary_timeline(summary: &str) -> Vec<String> {
|
||||
let mut lines = Vec::new();
|
||||
let mut in_timeline = false;
|
||||
|
||||
for line in format_compact_summary(summary).lines() {
|
||||
let trimmed = line.trim_end();
|
||||
if trimmed == "- Key timeline:" {
|
||||
in_timeline = true;
|
||||
continue;
|
||||
}
|
||||
if !in_timeline {
|
||||
continue;
|
||||
}
|
||||
if trimmed.is_empty() {
|
||||
break;
|
||||
}
|
||||
lines.push(trimmed.to_string());
|
||||
}
|
||||
|
||||
lines
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
collect_key_files, compact_session, format_compact_summary,
|
||||
get_compact_continuation_message, infer_pending_work, should_compact, CompactionConfig,
|
||||
};
|
||||
use crate::session::{ContentBlock, ConversationMessage, MessageRole, Session};
|
||||
|
||||
#[test]
|
||||
fn formats_compact_summary_like_upstream() {
|
||||
let summary = "<analysis>scratch</analysis>\n<summary>Kept work</summary>";
|
||||
assert_eq!(format_compact_summary(summary), "Summary:\nKept work");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn leaves_small_sessions_unchanged() {
|
||||
let mut session = Session::new();
|
||||
session.messages = vec![ConversationMessage::user_text("hello")];
|
||||
|
||||
let result = compact_session(&session, CompactionConfig::default());
|
||||
assert_eq!(result.removed_message_count, 0);
|
||||
assert_eq!(result.compacted_session, session);
|
||||
assert!(result.summary.is_empty());
|
||||
assert!(result.formatted_summary.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compacts_older_messages_into_a_system_summary() {
|
||||
let mut session = Session::new();
|
||||
session.messages = vec![
|
||||
ConversationMessage::user_text("one ".repeat(200)),
|
||||
ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "two ".repeat(200),
|
||||
}]),
|
||||
ConversationMessage::tool_result("1", "bash", "ok ".repeat(200), false),
|
||||
ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
blocks: vec![ContentBlock::Text {
|
||||
text: "recent".to_string(),
|
||||
}],
|
||||
usage: None,
|
||||
},
|
||||
];
|
||||
|
||||
let result = compact_session(
|
||||
&session,
|
||||
CompactionConfig {
|
||||
preserve_recent_messages: 2,
|
||||
max_estimated_tokens: 1,
|
||||
},
|
||||
);
|
||||
|
||||
// With the tool-use/tool-result boundary fix, the compaction preserves
|
||||
// one extra message to avoid an orphaned tool result at the boundary.
|
||||
// messages[1] (assistant) must be kept along with messages[2] (tool result).
|
||||
assert!(
|
||||
result.removed_message_count <= 2,
|
||||
"expected at most 2 removed, got {}",
|
||||
result.removed_message_count
|
||||
);
|
||||
assert_eq!(
|
||||
result.compacted_session.messages[0].role,
|
||||
MessageRole::System
|
||||
);
|
||||
assert!(matches!(
|
||||
&result.compacted_session.messages[0].blocks[0],
|
||||
ContentBlock::Text { text } if text.contains("Summary:")
|
||||
));
|
||||
assert!(result.formatted_summary.contains("Scope:"));
|
||||
assert!(result.formatted_summary.contains("Key timeline:"));
|
||||
assert!(should_compact(
|
||||
&session,
|
||||
CompactionConfig {
|
||||
preserve_recent_messages: 2,
|
||||
max_estimated_tokens: 1,
|
||||
}
|
||||
));
|
||||
// Note: with the tool-use/tool-result boundary guard the compacted session
|
||||
// may preserve one extra message at the boundary, so token reduction is
|
||||
// not guaranteed for small sessions. The invariant that matters is that
|
||||
// the removed_message_count is non-zero (something was compacted).
|
||||
assert!(
|
||||
result.removed_message_count > 0,
|
||||
"compaction must remove at least one message"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_previous_compacted_context_when_compacting_again() {
|
||||
let mut initial_session = Session::new();
|
||||
initial_session.messages = vec![
|
||||
ConversationMessage::user_text("Investigate rust/crates/runtime/src/compact.rs"),
|
||||
ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "I will inspect the compact flow.".to_string(),
|
||||
}]),
|
||||
ConversationMessage::user_text("Also update rust/crates/runtime/src/conversation.rs"),
|
||||
ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "Next: preserve prior summary context during auto compact.".to_string(),
|
||||
}]),
|
||||
];
|
||||
let config = CompactionConfig {
|
||||
preserve_recent_messages: 2,
|
||||
max_estimated_tokens: 1,
|
||||
};
|
||||
|
||||
let first = compact_session(&initial_session, config);
|
||||
let mut follow_up_messages = first.compacted_session.messages.clone();
|
||||
follow_up_messages.extend([
|
||||
ConversationMessage::user_text("Please add regression tests for compaction."),
|
||||
ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "Working on regression coverage now.".to_string(),
|
||||
}]),
|
||||
]);
|
||||
|
||||
let mut second_session = Session::new();
|
||||
second_session.messages = follow_up_messages;
|
||||
let second = compact_session(&second_session, config);
|
||||
|
||||
assert!(second
|
||||
.formatted_summary
|
||||
.contains("Previously compacted context:"));
|
||||
assert!(second
|
||||
.formatted_summary
|
||||
.contains("Scope: 2 earlier messages compacted"));
|
||||
assert!(second
|
||||
.formatted_summary
|
||||
.contains("Newly compacted context:"));
|
||||
assert!(second
|
||||
.formatted_summary
|
||||
.contains("Also update rust/crates/runtime/src/conversation.rs"));
|
||||
assert!(matches!(
|
||||
&second.compacted_session.messages[0].blocks[0],
|
||||
ContentBlock::Text { text }
|
||||
if text.contains("Previously compacted context:")
|
||||
&& text.contains("Newly compacted context:")
|
||||
));
|
||||
assert!(matches!(
|
||||
&second.compacted_session.messages[1].blocks[0],
|
||||
ContentBlock::Text { text } if text.contains("Please add regression tests for compaction.")
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_existing_compacted_summary_when_deciding_to_recompact() {
|
||||
let summary = "<summary>Conversation summary:\n- Scope: earlier work preserved.\n- Key timeline:\n - user: large preserved context\n</summary>";
|
||||
let mut session = Session::new();
|
||||
session.messages = vec![
|
||||
ConversationMessage {
|
||||
role: MessageRole::System,
|
||||
blocks: vec![ContentBlock::Text {
|
||||
text: get_compact_continuation_message(summary, true, true),
|
||||
}],
|
||||
usage: None,
|
||||
},
|
||||
ConversationMessage::user_text("tiny"),
|
||||
ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "recent".to_string(),
|
||||
}]),
|
||||
];
|
||||
|
||||
assert!(!should_compact(
|
||||
&session,
|
||||
CompactionConfig {
|
||||
preserve_recent_messages: 2,
|
||||
max_estimated_tokens: 1,
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncates_long_blocks_in_summary() {
|
||||
let summary = super::summarize_block(&ContentBlock::Text {
|
||||
text: "x".repeat(400),
|
||||
});
|
||||
assert!(summary.ends_with('…'));
|
||||
assert!(summary.chars().count() <= 161);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extracts_key_files_from_message_content() {
|
||||
let files = collect_key_files(&[ConversationMessage::user_text(
|
||||
"Update rust/crates/runtime/src/compact.rs and rust/crates/rusty-claude-cli/src/main.rs next.",
|
||||
)]);
|
||||
assert!(files.contains(&"rust/crates/runtime/src/compact.rs".to_string()));
|
||||
assert!(files.contains(&"rust/crates/rusty-claude-cli/src/main.rs".to_string()));
|
||||
}
|
||||
|
||||
/// Regression: compaction must not split an assistant(ToolUse) /
|
||||
/// user(ToolResult) pair at the boundary. An orphaned tool-result message
|
||||
/// without the preceding assistant tool_calls causes a 400 on the
|
||||
/// OpenAI-compat path (gaebal-gajae repro 2026-04-09).
|
||||
#[test]
|
||||
fn compaction_does_not_split_tool_use_tool_result_pair() {
|
||||
use crate::session::{ContentBlock, Session};
|
||||
|
||||
let tool_id = "call_abc";
|
||||
let mut session = Session::default();
|
||||
// Turn 1: user prompt
|
||||
session
|
||||
.push_message(ConversationMessage::user_text("Search for files"))
|
||||
.unwrap();
|
||||
// Turn 2: assistant calls a tool
|
||||
session
|
||||
.push_message(ConversationMessage::assistant(vec![
|
||||
ContentBlock::ToolUse {
|
||||
id: tool_id.to_string(),
|
||||
name: "search".to_string(),
|
||||
input: "{\"q\":\"*.rs\"}".to_string(),
|
||||
},
|
||||
]))
|
||||
.unwrap();
|
||||
// Turn 3: tool result
|
||||
session
|
||||
.push_message(ConversationMessage::tool_result(
|
||||
tool_id,
|
||||
"search",
|
||||
"found 5 files",
|
||||
false,
|
||||
))
|
||||
.unwrap();
|
||||
// Turn 4: assistant final response
|
||||
session
|
||||
.push_message(ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "Done.".to_string(),
|
||||
}]))
|
||||
.unwrap();
|
||||
|
||||
// Compact preserving only 1 recent message — without the fix this
|
||||
// would cut the boundary so that the tool result (turn 3) is first,
|
||||
// without its preceding assistant tool_calls (turn 2).
|
||||
let config = CompactionConfig {
|
||||
preserve_recent_messages: 1,
|
||||
..CompactionConfig::default()
|
||||
};
|
||||
let result = compact_session(&session, config);
|
||||
// After compaction, no two consecutive messages should have the pattern
|
||||
// tool_result immediately following a non-assistant message (i.e. an
|
||||
// orphaned tool result without a preceding assistant ToolUse).
|
||||
let messages = &result.compacted_session.messages;
|
||||
for i in 1..messages.len() {
|
||||
let curr_is_tool_result = messages[i]
|
||||
.blocks
|
||||
.first()
|
||||
.map(|b| matches!(b, ContentBlock::ToolResult { .. }))
|
||||
.unwrap_or(false);
|
||||
if curr_is_tool_result {
|
||||
let prev_has_tool_use = messages[i - 1]
|
||||
.blocks
|
||||
.iter()
|
||||
.any(|b| matches!(b, ContentBlock::ToolUse { .. }));
|
||||
assert!(
|
||||
prev_has_tool_use,
|
||||
"message[{}] is a ToolResult but message[{}] has no ToolUse: {:?}",
|
||||
i,
|
||||
i - 1,
|
||||
&messages[i - 1].blocks
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn infers_pending_work_from_recent_messages() {
|
||||
let pending = infer_pending_work(&[
|
||||
ConversationMessage::user_text("done"),
|
||||
ConversationMessage::assistant(vec![ContentBlock::Text {
|
||||
text: "Next: update tests and follow up on remaining CLI polish.".to_string(),
|
||||
}]),
|
||||
]);
|
||||
assert_eq!(pending.len(), 1);
|
||||
assert!(pending[0].contains("Next: update tests"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,901 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::config::ConfigError;
|
||||
use crate::json::JsonValue;
|
||||
|
||||
/// Diagnostic emitted when a config file contains a suspect field.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ConfigDiagnostic {
|
||||
pub path: String,
|
||||
pub field: String,
|
||||
pub line: Option<usize>,
|
||||
pub kind: DiagnosticKind,
|
||||
}
|
||||
|
||||
/// Classification of the diagnostic.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum DiagnosticKind {
|
||||
UnknownKey {
|
||||
suggestion: Option<String>,
|
||||
},
|
||||
WrongType {
|
||||
expected: &'static str,
|
||||
got: &'static str,
|
||||
},
|
||||
Deprecated {
|
||||
replacement: &'static str,
|
||||
},
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ConfigDiagnostic {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
let location = self
|
||||
.line
|
||||
.map_or_else(String::new, |line| format!(" (line {line})"));
|
||||
match &self.kind {
|
||||
DiagnosticKind::UnknownKey { suggestion: None } => {
|
||||
write!(f, "{}: unknown key \"{}\"{location}", self.path, self.field)
|
||||
}
|
||||
DiagnosticKind::UnknownKey {
|
||||
suggestion: Some(hint),
|
||||
} => {
|
||||
write!(
|
||||
f,
|
||||
"{}: unknown key \"{}\"{location}. Did you mean \"{}\"?",
|
||||
self.path, self.field, hint
|
||||
)
|
||||
}
|
||||
DiagnosticKind::WrongType { expected, got } => {
|
||||
write!(
|
||||
f,
|
||||
"{}: field \"{}\" must be {expected}, got {got}{location}",
|
||||
self.path, self.field
|
||||
)
|
||||
}
|
||||
DiagnosticKind::Deprecated { replacement } => {
|
||||
write!(
|
||||
f,
|
||||
"{}: field \"{}\" is deprecated{location}. Use \"{replacement}\" instead",
|
||||
self.path, self.field
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of validating a single config file.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ValidationResult {
|
||||
pub errors: Vec<ConfigDiagnostic>,
|
||||
pub warnings: Vec<ConfigDiagnostic>,
|
||||
}
|
||||
|
||||
impl ValidationResult {
|
||||
#[must_use]
|
||||
pub fn is_ok(&self) -> bool {
|
||||
self.errors.is_empty()
|
||||
}
|
||||
|
||||
fn merge(&mut self, other: Self) {
|
||||
self.errors.extend(other.errors);
|
||||
self.warnings.extend(other.warnings);
|
||||
}
|
||||
}
|
||||
|
||||
// ---- known-key schema ----
|
||||
|
||||
/// Expected type for a config field.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum FieldType {
|
||||
String,
|
||||
Bool,
|
||||
Object,
|
||||
StringArray,
|
||||
Number,
|
||||
}
|
||||
|
||||
impl FieldType {
|
||||
fn label(self) -> &'static str {
|
||||
match self {
|
||||
Self::String => "a string",
|
||||
Self::Bool => "a boolean",
|
||||
Self::Object => "an object",
|
||||
Self::StringArray => "an array of strings",
|
||||
Self::Number => "a number",
|
||||
}
|
||||
}
|
||||
|
||||
fn matches(self, value: &JsonValue) -> bool {
|
||||
match self {
|
||||
Self::String => value.as_str().is_some(),
|
||||
Self::Bool => value.as_bool().is_some(),
|
||||
Self::Object => value.as_object().is_some(),
|
||||
Self::StringArray => value
|
||||
.as_array()
|
||||
.is_some_and(|arr| arr.iter().all(|v| v.as_str().is_some())),
|
||||
Self::Number => value.as_i64().is_some(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn json_type_label(value: &JsonValue) -> &'static str {
|
||||
match value {
|
||||
JsonValue::Null => "null",
|
||||
JsonValue::Bool(_) => "a boolean",
|
||||
JsonValue::Number(_) => "a number",
|
||||
JsonValue::String(_) => "a string",
|
||||
JsonValue::Array(_) => "an array",
|
||||
JsonValue::Object(_) => "an object",
|
||||
}
|
||||
}
|
||||
|
||||
struct FieldSpec {
|
||||
name: &'static str,
|
||||
expected: FieldType,
|
||||
}
|
||||
|
||||
struct DeprecatedField {
|
||||
name: &'static str,
|
||||
replacement: &'static str,
|
||||
}
|
||||
|
||||
const TOP_LEVEL_FIELDS: &[FieldSpec] = &[
|
||||
FieldSpec {
|
||||
name: "$schema",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "model",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "hooks",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "permissions",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "permissionMode",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "mcpServers",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "oauth",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "enabledPlugins",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "plugins",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "sandbox",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "env",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "aliases",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "providerFallbacks",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "trustedRoots",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
];
|
||||
|
||||
const HOOKS_FIELDS: &[FieldSpec] = &[
|
||||
FieldSpec {
|
||||
name: "PreToolUse",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "PostToolUse",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "PostToolUseFailure",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
];
|
||||
|
||||
const PERMISSIONS_FIELDS: &[FieldSpec] = &[
|
||||
FieldSpec {
|
||||
name: "defaultMode",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "allow",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "deny",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "ask",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
];
|
||||
|
||||
const PLUGINS_FIELDS: &[FieldSpec] = &[
|
||||
FieldSpec {
|
||||
name: "enabled",
|
||||
expected: FieldType::Object,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "externalDirectories",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "installRoot",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "registryPath",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "bundledRoot",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "maxOutputTokens",
|
||||
expected: FieldType::Number,
|
||||
},
|
||||
];
|
||||
|
||||
const SANDBOX_FIELDS: &[FieldSpec] = &[
|
||||
FieldSpec {
|
||||
name: "enabled",
|
||||
expected: FieldType::Bool,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "namespaceRestrictions",
|
||||
expected: FieldType::Bool,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "networkIsolation",
|
||||
expected: FieldType::Bool,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "filesystemMode",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "allowedMounts",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
];
|
||||
|
||||
const OAUTH_FIELDS: &[FieldSpec] = &[
|
||||
FieldSpec {
|
||||
name: "clientId",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "authorizeUrl",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "tokenUrl",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "callbackPort",
|
||||
expected: FieldType::Number,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "manualRedirectUrl",
|
||||
expected: FieldType::String,
|
||||
},
|
||||
FieldSpec {
|
||||
name: "scopes",
|
||||
expected: FieldType::StringArray,
|
||||
},
|
||||
];
|
||||
|
||||
const DEPRECATED_FIELDS: &[DeprecatedField] = &[
|
||||
DeprecatedField {
|
||||
name: "permissionMode",
|
||||
replacement: "permissions.defaultMode",
|
||||
},
|
||||
DeprecatedField {
|
||||
name: "enabledPlugins",
|
||||
replacement: "plugins.enabled",
|
||||
},
|
||||
];
|
||||
|
||||
// ---- line-number resolution ----
|
||||
|
||||
/// Find the 1-based line number where a JSON key first appears in the raw source.
|
||||
fn find_key_line(source: &str, key: &str) -> Option<usize> {
|
||||
// Search for `"key"` followed by optional whitespace and a colon.
|
||||
let needle = format!("\"{key}\"");
|
||||
let mut search_start = 0;
|
||||
while let Some(offset) = source[search_start..].find(&needle) {
|
||||
let absolute = search_start + offset;
|
||||
let after = absolute + needle.len();
|
||||
// Verify the next non-whitespace char is `:` to confirm this is a key, not a value.
|
||||
if source[after..].chars().find(|ch| !ch.is_ascii_whitespace()) == Some(':') {
|
||||
return Some(source[..absolute].chars().filter(|&ch| ch == '\n').count() + 1);
|
||||
}
|
||||
search_start = after;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
// ---- core validation ----
|
||||
|
||||
fn validate_object_keys(
|
||||
object: &BTreeMap<String, JsonValue>,
|
||||
known_fields: &[FieldSpec],
|
||||
prefix: &str,
|
||||
source: &str,
|
||||
path_display: &str,
|
||||
) -> ValidationResult {
|
||||
let mut result = ValidationResult {
|
||||
errors: Vec::new(),
|
||||
warnings: Vec::new(),
|
||||
};
|
||||
|
||||
let known_names: Vec<&str> = known_fields.iter().map(|f| f.name).collect();
|
||||
|
||||
for (key, value) in object {
|
||||
let field_path = if prefix.is_empty() {
|
||||
key.clone()
|
||||
} else {
|
||||
format!("{prefix}.{key}")
|
||||
};
|
||||
|
||||
if let Some(spec) = known_fields.iter().find(|f| f.name == key) {
|
||||
// Type check.
|
||||
if !spec.expected.matches(value) {
|
||||
result.errors.push(ConfigDiagnostic {
|
||||
path: path_display.to_string(),
|
||||
field: field_path,
|
||||
line: find_key_line(source, key),
|
||||
kind: DiagnosticKind::WrongType {
|
||||
expected: spec.expected.label(),
|
||||
got: json_type_label(value),
|
||||
},
|
||||
});
|
||||
}
|
||||
} else if DEPRECATED_FIELDS.iter().any(|d| d.name == key) {
|
||||
// Deprecated key — handled separately, not an unknown-key error.
|
||||
} else {
|
||||
// Unknown key.
|
||||
let suggestion = suggest_field(key, &known_names);
|
||||
result.errors.push(ConfigDiagnostic {
|
||||
path: path_display.to_string(),
|
||||
field: field_path,
|
||||
line: find_key_line(source, key),
|
||||
kind: DiagnosticKind::UnknownKey { suggestion },
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn suggest_field(input: &str, candidates: &[&str]) -> Option<String> {
|
||||
let input_lower = input.to_ascii_lowercase();
|
||||
candidates
|
||||
.iter()
|
||||
.filter_map(|candidate| {
|
||||
let distance = simple_edit_distance(&input_lower, &candidate.to_ascii_lowercase());
|
||||
(distance <= 3).then_some((distance, *candidate))
|
||||
})
|
||||
.min_by_key(|(distance, _)| *distance)
|
||||
.map(|(_, name)| name.to_string())
|
||||
}
|
||||
|
||||
fn simple_edit_distance(left: &str, right: &str) -> usize {
|
||||
if left.is_empty() {
|
||||
return right.len();
|
||||
}
|
||||
if right.is_empty() {
|
||||
return left.len();
|
||||
}
|
||||
let right_chars: Vec<char> = right.chars().collect();
|
||||
let mut previous: Vec<usize> = (0..=right_chars.len()).collect();
|
||||
let mut current = vec![0; right_chars.len() + 1];
|
||||
|
||||
for (left_index, left_char) in left.chars().enumerate() {
|
||||
current[0] = left_index + 1;
|
||||
for (right_index, right_char) in right_chars.iter().enumerate() {
|
||||
let cost = usize::from(left_char != *right_char);
|
||||
current[right_index + 1] = (previous[right_index + 1] + 1)
|
||||
.min(current[right_index] + 1)
|
||||
.min(previous[right_index] + cost);
|
||||
}
|
||||
previous.clone_from(¤t);
|
||||
}
|
||||
|
||||
previous[right_chars.len()]
|
||||
}
|
||||
|
||||
/// Validate a parsed config file's keys and types against the known schema.
|
||||
///
|
||||
/// Returns diagnostics (errors and deprecation warnings) without blocking the load.
|
||||
pub fn validate_config_file(
|
||||
object: &BTreeMap<String, JsonValue>,
|
||||
source: &str,
|
||||
file_path: &Path,
|
||||
) -> ValidationResult {
|
||||
let path_display = file_path.display().to_string();
|
||||
let mut result = validate_object_keys(object, TOP_LEVEL_FIELDS, "", source, &path_display);
|
||||
|
||||
// Check deprecated fields.
|
||||
for deprecated in DEPRECATED_FIELDS {
|
||||
if object.contains_key(deprecated.name) {
|
||||
result.warnings.push(ConfigDiagnostic {
|
||||
path: path_display.clone(),
|
||||
field: deprecated.name.to_string(),
|
||||
line: find_key_line(source, deprecated.name),
|
||||
kind: DiagnosticKind::Deprecated {
|
||||
replacement: deprecated.replacement,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Validate known nested objects.
|
||||
if let Some(hooks) = object.get("hooks").and_then(JsonValue::as_object) {
|
||||
result.merge(validate_object_keys(
|
||||
hooks,
|
||||
HOOKS_FIELDS,
|
||||
"hooks",
|
||||
source,
|
||||
&path_display,
|
||||
));
|
||||
}
|
||||
if let Some(permissions) = object.get("permissions").and_then(JsonValue::as_object) {
|
||||
result.merge(validate_object_keys(
|
||||
permissions,
|
||||
PERMISSIONS_FIELDS,
|
||||
"permissions",
|
||||
source,
|
||||
&path_display,
|
||||
));
|
||||
}
|
||||
if let Some(plugins) = object.get("plugins").and_then(JsonValue::as_object) {
|
||||
result.merge(validate_object_keys(
|
||||
plugins,
|
||||
PLUGINS_FIELDS,
|
||||
"plugins",
|
||||
source,
|
||||
&path_display,
|
||||
));
|
||||
}
|
||||
if let Some(sandbox) = object.get("sandbox").and_then(JsonValue::as_object) {
|
||||
result.merge(validate_object_keys(
|
||||
sandbox,
|
||||
SANDBOX_FIELDS,
|
||||
"sandbox",
|
||||
source,
|
||||
&path_display,
|
||||
));
|
||||
}
|
||||
if let Some(oauth) = object.get("oauth").and_then(JsonValue::as_object) {
|
||||
result.merge(validate_object_keys(
|
||||
oauth,
|
||||
OAUTH_FIELDS,
|
||||
"oauth",
|
||||
source,
|
||||
&path_display,
|
||||
));
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Check whether a file path uses an unsupported config format (e.g. TOML).
|
||||
pub fn check_unsupported_format(file_path: &Path) -> Result<(), ConfigError> {
|
||||
if let Some(ext) = file_path.extension().and_then(|e| e.to_str()) {
|
||||
if ext.eq_ignore_ascii_case("toml") {
|
||||
return Err(ConfigError::Parse(format!(
|
||||
"{}: TOML config files are not supported. Use JSON (settings.json) instead",
|
||||
file_path.display()
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Format all diagnostics into a human-readable report.
|
||||
#[must_use]
|
||||
pub fn format_diagnostics(result: &ValidationResult) -> String {
|
||||
let mut lines = Vec::new();
|
||||
for warning in &result.warnings {
|
||||
lines.push(format!("warning: {warning}"));
|
||||
}
|
||||
for error in &result.errors {
|
||||
lines.push(format!("error: {error}"));
|
||||
}
|
||||
lines.join("\n")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::path::PathBuf;
|
||||
|
||||
fn test_path() -> PathBuf {
|
||||
PathBuf::from("/test/settings.json")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_unknown_top_level_key() {
|
||||
// given
|
||||
let source = r#"{"model": "opus", "unknownField": true}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].field, "unknownField");
|
||||
assert!(matches!(
|
||||
result.errors[0].kind,
|
||||
DiagnosticKind::UnknownKey { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_wrong_type_for_model() {
|
||||
// given
|
||||
let source = r#"{"model": 123}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].field, "model");
|
||||
assert!(matches!(
|
||||
result.errors[0].kind,
|
||||
DiagnosticKind::WrongType {
|
||||
expected: "a string",
|
||||
got: "a number"
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_deprecated_permission_mode() {
|
||||
// given
|
||||
let source = r#"{"permissionMode": "plan"}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.warnings.len(), 1);
|
||||
assert_eq!(result.warnings[0].field, "permissionMode");
|
||||
assert!(matches!(
|
||||
result.warnings[0].kind,
|
||||
DiagnosticKind::Deprecated {
|
||||
replacement: "permissions.defaultMode"
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_deprecated_enabled_plugins() {
|
||||
// given
|
||||
let source = r#"{"enabledPlugins": {"tool-guard@builtin": true}}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.warnings.len(), 1);
|
||||
assert_eq!(result.warnings[0].field, "enabledPlugins");
|
||||
assert!(matches!(
|
||||
result.warnings[0].kind,
|
||||
DiagnosticKind::Deprecated {
|
||||
replacement: "plugins.enabled"
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_line_number_for_unknown_key() {
|
||||
// given
|
||||
let source = "{\n \"model\": \"opus\",\n \"badKey\": true\n}";
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].line, Some(3));
|
||||
assert_eq!(result.errors[0].field, "badKey");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reports_line_number_for_wrong_type() {
|
||||
// given
|
||||
let source = "{\n \"model\": 42\n}";
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].line, Some(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_nested_hooks_keys() {
|
||||
// given
|
||||
let source = r#"{"hooks": {"PreToolUse": ["cmd"], "BadHook": ["x"]}}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].field, "hooks.BadHook");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_nested_permissions_keys() {
|
||||
// given
|
||||
let source = r#"{"permissions": {"allow": ["Read"], "denyAll": true}}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].field, "permissions.denyAll");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_nested_sandbox_keys() {
|
||||
// given
|
||||
let source = r#"{"sandbox": {"enabled": true, "containerMode": "strict"}}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].field, "sandbox.containerMode");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_nested_plugins_keys() {
|
||||
// given
|
||||
let source = r#"{"plugins": {"installRoot": "/tmp", "autoUpdate": true}}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].field, "plugins.autoUpdate");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_nested_oauth_keys() {
|
||||
// given
|
||||
let source = r#"{"oauth": {"clientId": "abc", "secret": "hidden"}}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].field, "oauth.secret");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_config_produces_no_diagnostics() {
|
||||
// given
|
||||
let source = r#"{
|
||||
"model": "opus",
|
||||
"hooks": {"PreToolUse": ["guard"]},
|
||||
"permissions": {"defaultMode": "plan", "allow": ["Read"]},
|
||||
"mcpServers": {},
|
||||
"sandbox": {"enabled": false}
|
||||
}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert!(result.is_ok());
|
||||
assert!(result.warnings.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn suggests_close_field_name() {
|
||||
// given
|
||||
let source = r#"{"modle": "opus"}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
match &result.errors[0].kind {
|
||||
DiagnosticKind::UnknownKey {
|
||||
suggestion: Some(s),
|
||||
} => assert_eq!(s, "model"),
|
||||
other => panic!("expected suggestion, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_diagnostics_includes_all_entries() {
|
||||
// given
|
||||
let source = r#"{"permissionMode": "plan", "badKey": 1}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// when
|
||||
let output = format_diagnostics(&result);
|
||||
|
||||
// then
|
||||
assert!(output.contains("warning:"));
|
||||
assert!(output.contains("error:"));
|
||||
assert!(output.contains("badKey"));
|
||||
assert!(output.contains("permissionMode"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_unsupported_format_rejects_toml() {
|
||||
// given
|
||||
let path = PathBuf::from("/home/.claw/settings.toml");
|
||||
|
||||
// when
|
||||
let result = check_unsupported_format(&path);
|
||||
|
||||
// then
|
||||
assert!(result.is_err());
|
||||
let message = result.unwrap_err().to_string();
|
||||
assert!(message.contains("TOML"));
|
||||
assert!(message.contains("settings.toml"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_unsupported_format_allows_json() {
|
||||
// given
|
||||
let path = PathBuf::from("/home/.claw/settings.json");
|
||||
|
||||
// when / then
|
||||
assert!(check_unsupported_format(&path).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrong_type_in_nested_sandbox_field() {
|
||||
// given
|
||||
let source = r#"{"sandbox": {"enabled": "yes"}}"#;
|
||||
let parsed = JsonValue::parse(source).expect("valid json");
|
||||
let object = parsed.as_object().expect("object");
|
||||
|
||||
// when
|
||||
let result = validate_config_file(object, source, &test_path());
|
||||
|
||||
// then
|
||||
assert_eq!(result.errors.len(), 1);
|
||||
assert_eq!(result.errors[0].field, "sandbox.enabled");
|
||||
assert!(matches!(
|
||||
result.errors[0].kind,
|
||||
DiagnosticKind::WrongType {
|
||||
expected: "a boolean",
|
||||
got: "a string"
|
||||
}
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn display_format_unknown_key_with_line() {
|
||||
// given
|
||||
let diag = ConfigDiagnostic {
|
||||
path: "/test/settings.json".to_string(),
|
||||
field: "badKey".to_string(),
|
||||
line: Some(5),
|
||||
kind: DiagnosticKind::UnknownKey { suggestion: None },
|
||||
};
|
||||
|
||||
// when
|
||||
let output = diag.to_string();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
output,
|
||||
r#"/test/settings.json: unknown key "badKey" (line 5)"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn display_format_wrong_type_with_line() {
|
||||
// given
|
||||
let diag = ConfigDiagnostic {
|
||||
path: "/test/settings.json".to_string(),
|
||||
field: "model".to_string(),
|
||||
line: Some(2),
|
||||
kind: DiagnosticKind::WrongType {
|
||||
expected: "a string",
|
||||
got: "a number",
|
||||
},
|
||||
};
|
||||
|
||||
// when
|
||||
let output = diag.to_string();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
output,
|
||||
r#"/test/settings.json: field "model" must be a string, got a number (line 2)"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn display_format_deprecated_with_line() {
|
||||
// given
|
||||
let diag = ConfigDiagnostic {
|
||||
path: "/test/settings.json".to_string(),
|
||||
field: "permissionMode".to_string(),
|
||||
line: Some(3),
|
||||
kind: DiagnosticKind::Deprecated {
|
||||
replacement: "permissions.defaultMode",
|
||||
},
|
||||
};
|
||||
|
||||
// when
|
||||
let output = diag.to_string();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
output,
|
||||
r#"/test/settings.json: field "permissionMode" is deprecated (line 3). Use "permissions.defaultMode" instead"#
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,41 @@ use regex::RegexBuilder;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use walkdir::WalkDir;
|
||||
|
||||
/// Maximum file size that can be read (10 MB).
|
||||
const MAX_READ_SIZE: u64 = 10 * 1024 * 1024;
|
||||
|
||||
/// Maximum file size that can be written (10 MB).
|
||||
const MAX_WRITE_SIZE: usize = 10 * 1024 * 1024;
|
||||
|
||||
/// Check whether a file appears to contain binary content by examining
|
||||
/// the first chunk for NUL bytes.
|
||||
fn is_binary_file(path: &Path) -> io::Result<bool> {
|
||||
use std::io::Read;
|
||||
let mut file = fs::File::open(path)?;
|
||||
let mut buffer = [0u8; 8192];
|
||||
let bytes_read = file.read(&mut buffer)?;
|
||||
Ok(buffer[..bytes_read].contains(&0))
|
||||
}
|
||||
|
||||
/// Validate that a resolved path stays within the given workspace root.
|
||||
/// Returns the canonical path on success, or an error if the path escapes
|
||||
/// the workspace boundary (e.g. via `../` traversal or symlink).
|
||||
#[allow(dead_code)]
|
||||
fn validate_workspace_boundary(resolved: &Path, workspace_root: &Path) -> io::Result<()> {
|
||||
if !resolved.starts_with(workspace_root) {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::PermissionDenied,
|
||||
format!(
|
||||
"path {} escapes workspace boundary {}",
|
||||
resolved.display(),
|
||||
workspace_root.display()
|
||||
),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Text payload returned by file-reading operations.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct TextFilePayload {
|
||||
#[serde(rename = "filePath")]
|
||||
@@ -22,6 +57,7 @@ pub struct TextFilePayload {
|
||||
pub total_lines: usize,
|
||||
}
|
||||
|
||||
/// Output envelope for the `read_file` tool.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct ReadFileOutput {
|
||||
#[serde(rename = "type")]
|
||||
@@ -29,6 +65,7 @@ pub struct ReadFileOutput {
|
||||
pub file: TextFilePayload,
|
||||
}
|
||||
|
||||
/// Structured patch hunk emitted by write and edit operations.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct StructuredPatchHunk {
|
||||
#[serde(rename = "oldStart")]
|
||||
@@ -42,6 +79,7 @@ pub struct StructuredPatchHunk {
|
||||
pub lines: Vec<String>,
|
||||
}
|
||||
|
||||
/// Output envelope for full-file write operations.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct WriteFileOutput {
|
||||
#[serde(rename = "type")]
|
||||
@@ -57,6 +95,7 @@ pub struct WriteFileOutput {
|
||||
pub git_diff: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Output envelope for targeted string-replacement edits.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct EditFileOutput {
|
||||
#[serde(rename = "filePath")]
|
||||
@@ -77,6 +116,7 @@ pub struct EditFileOutput {
|
||||
pub git_diff: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Result of a glob-based filename search.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||
pub struct GlobSearchOutput {
|
||||
#[serde(rename = "durationMs")]
|
||||
@@ -87,6 +127,7 @@ pub struct GlobSearchOutput {
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
/// Parameters accepted by the grep-style search tool.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct GrepSearchInput {
|
||||
pub pattern: String,
|
||||
@@ -112,6 +153,7 @@ pub struct GrepSearchInput {
|
||||
pub multiline: Option<bool>,
|
||||
}
|
||||
|
||||
/// Result payload returned by the grep-style search tool.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
|
||||
pub struct GrepSearchOutput {
|
||||
pub mode: Option<String>,
|
||||
@@ -129,12 +171,35 @@ pub struct GrepSearchOutput {
|
||||
pub applied_offset: Option<usize>,
|
||||
}
|
||||
|
||||
/// Reads a text file and returns a line-windowed payload.
|
||||
pub fn read_file(
|
||||
path: &str,
|
||||
offset: Option<usize>,
|
||||
limit: Option<usize>,
|
||||
) -> io::Result<ReadFileOutput> {
|
||||
let absolute_path = normalize_path(path)?;
|
||||
|
||||
// Check file size before reading
|
||||
let metadata = fs::metadata(&absolute_path)?;
|
||||
if metadata.len() > MAX_READ_SIZE {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
format!(
|
||||
"file is too large ({} bytes, max {} bytes)",
|
||||
metadata.len(),
|
||||
MAX_READ_SIZE
|
||||
),
|
||||
));
|
||||
}
|
||||
|
||||
// Detect binary files
|
||||
if is_binary_file(&absolute_path)? {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
"file appears to be binary",
|
||||
));
|
||||
}
|
||||
|
||||
let content = fs::read_to_string(&absolute_path)?;
|
||||
let lines: Vec<&str> = content.lines().collect();
|
||||
let start_index = offset.unwrap_or(0).min(lines.len());
|
||||
@@ -155,7 +220,19 @@ pub fn read_file(
|
||||
})
|
||||
}
|
||||
|
||||
/// Replaces a file's contents and returns patch metadata.
|
||||
pub fn write_file(path: &str, content: &str) -> io::Result<WriteFileOutput> {
|
||||
if content.len() > MAX_WRITE_SIZE {
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::InvalidData,
|
||||
format!(
|
||||
"content is too large ({} bytes, max {} bytes)",
|
||||
content.len(),
|
||||
MAX_WRITE_SIZE
|
||||
),
|
||||
));
|
||||
}
|
||||
|
||||
let absolute_path = normalize_path_allow_missing(path)?;
|
||||
let original_file = fs::read_to_string(&absolute_path).ok();
|
||||
if let Some(parent) = absolute_path.parent() {
|
||||
@@ -177,6 +254,7 @@ pub fn write_file(path: &str, content: &str) -> io::Result<WriteFileOutput> {
|
||||
})
|
||||
}
|
||||
|
||||
/// Performs an in-file string replacement and returns patch metadata.
|
||||
pub fn edit_file(
|
||||
path: &str,
|
||||
old_string: &str,
|
||||
@@ -217,6 +295,7 @@ pub fn edit_file(
|
||||
})
|
||||
}
|
||||
|
||||
/// Expands a glob pattern and returns matching filenames.
|
||||
pub fn glob_search(pattern: &str, path: Option<&str>) -> io::Result<GlobSearchOutput> {
|
||||
let started = Instant::now();
|
||||
let base_dir = path
|
||||
@@ -229,14 +308,22 @@ pub fn glob_search(pattern: &str, path: Option<&str>) -> io::Result<GlobSearchOu
|
||||
base_dir.join(pattern).to_string_lossy().into_owned()
|
||||
};
|
||||
|
||||
// The `glob` crate does not support brace expansion ({a,b,c}).
|
||||
// Expand braces into multiple patterns so patterns like
|
||||
// `Assets/**/*.{cs,uxml,uss}` work correctly.
|
||||
let expanded = expand_braces(&search_pattern);
|
||||
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
let mut matches = Vec::new();
|
||||
let entries = glob::glob(&search_pattern)
|
||||
for pat in &expanded {
|
||||
let entries = glob::glob(pat)
|
||||
.map_err(|error| io::Error::new(io::ErrorKind::InvalidInput, error.to_string()))?;
|
||||
for entry in entries.flatten() {
|
||||
if entry.is_file() {
|
||||
if entry.is_file() && seen.insert(entry.clone()) {
|
||||
matches.push(entry);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
matches.sort_by_key(|path| {
|
||||
fs::metadata(path)
|
||||
@@ -260,6 +347,7 @@ pub fn glob_search(pattern: &str, path: Option<&str>) -> io::Result<GlobSearchOu
|
||||
})
|
||||
}
|
||||
|
||||
/// Runs a regex search over workspace files with optional context lines.
|
||||
pub fn grep_search(input: &GrepSearchInput) -> io::Result<GrepSearchOutput> {
|
||||
let base_path = input
|
||||
.path
|
||||
@@ -477,11 +565,98 @@ fn normalize_path_allow_missing(path: &str) -> io::Result<PathBuf> {
|
||||
Ok(candidate)
|
||||
}
|
||||
|
||||
/// Read a file with workspace boundary enforcement.
|
||||
#[allow(dead_code)]
|
||||
pub fn read_file_in_workspace(
|
||||
path: &str,
|
||||
offset: Option<usize>,
|
||||
limit: Option<usize>,
|
||||
workspace_root: &Path,
|
||||
) -> io::Result<ReadFileOutput> {
|
||||
let absolute_path = normalize_path(path)?;
|
||||
let canonical_root = workspace_root
|
||||
.canonicalize()
|
||||
.unwrap_or_else(|_| workspace_root.to_path_buf());
|
||||
validate_workspace_boundary(&absolute_path, &canonical_root)?;
|
||||
read_file(path, offset, limit)
|
||||
}
|
||||
|
||||
/// Write a file with workspace boundary enforcement.
|
||||
#[allow(dead_code)]
|
||||
pub fn write_file_in_workspace(
|
||||
path: &str,
|
||||
content: &str,
|
||||
workspace_root: &Path,
|
||||
) -> io::Result<WriteFileOutput> {
|
||||
let absolute_path = normalize_path_allow_missing(path)?;
|
||||
let canonical_root = workspace_root
|
||||
.canonicalize()
|
||||
.unwrap_or_else(|_| workspace_root.to_path_buf());
|
||||
validate_workspace_boundary(&absolute_path, &canonical_root)?;
|
||||
write_file(path, content)
|
||||
}
|
||||
|
||||
/// Edit a file with workspace boundary enforcement.
|
||||
#[allow(dead_code)]
|
||||
pub fn edit_file_in_workspace(
|
||||
path: &str,
|
||||
old_string: &str,
|
||||
new_string: &str,
|
||||
replace_all: bool,
|
||||
workspace_root: &Path,
|
||||
) -> io::Result<EditFileOutput> {
|
||||
let absolute_path = normalize_path(path)?;
|
||||
let canonical_root = workspace_root
|
||||
.canonicalize()
|
||||
.unwrap_or_else(|_| workspace_root.to_path_buf());
|
||||
validate_workspace_boundary(&absolute_path, &canonical_root)?;
|
||||
edit_file(path, old_string, new_string, replace_all)
|
||||
}
|
||||
|
||||
/// Check whether a path is a symlink that resolves outside the workspace.
|
||||
#[allow(dead_code)]
|
||||
pub fn is_symlink_escape(path: &Path, workspace_root: &Path) -> io::Result<bool> {
|
||||
let metadata = fs::symlink_metadata(path)?;
|
||||
if !metadata.is_symlink() {
|
||||
return Ok(false);
|
||||
}
|
||||
let resolved = path.canonicalize()?;
|
||||
let canonical_root = workspace_root
|
||||
.canonicalize()
|
||||
.unwrap_or_else(|_| workspace_root.to_path_buf());
|
||||
Ok(!resolved.starts_with(&canonical_root))
|
||||
}
|
||||
|
||||
/// Expand shell-style brace groups in a glob pattern.
|
||||
///
|
||||
/// Handles one level of braces: `foo.{a,b,c}` → `["foo.a", "foo.b", "foo.c"]`.
|
||||
/// Nested braces are not expanded (uncommon in practice).
|
||||
/// Patterns without braces pass through unchanged.
|
||||
fn expand_braces(pattern: &str) -> Vec<String> {
|
||||
let Some(open) = pattern.find('{') else {
|
||||
return vec![pattern.to_owned()];
|
||||
};
|
||||
let Some(close) = pattern[open..].find('}').map(|i| open + i) else {
|
||||
// Unmatched brace — treat as literal.
|
||||
return vec![pattern.to_owned()];
|
||||
};
|
||||
let prefix = &pattern[..open];
|
||||
let suffix = &pattern[close + 1..];
|
||||
let alternatives = &pattern[open + 1..close];
|
||||
alternatives
|
||||
.split(',')
|
||||
.flat_map(|alt| expand_braces(&format!("{prefix}{alt}{suffix}")))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use super::{edit_file, glob_search, grep_search, read_file, write_file, GrepSearchInput};
|
||||
use super::{
|
||||
edit_file, expand_braces, glob_search, grep_search, is_symlink_escape, read_file,
|
||||
read_file_in_workspace, write_file, GrepSearchInput, MAX_WRITE_SIZE,
|
||||
};
|
||||
|
||||
fn temp_path(name: &str) -> std::path::PathBuf {
|
||||
let unique = SystemTime::now()
|
||||
@@ -513,6 +688,73 @@ mod tests {
|
||||
assert!(output.replace_all);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_binary_files() {
|
||||
let path = temp_path("binary-test.bin");
|
||||
std::fs::write(&path, b"\x00\x01\x02\x03binary content").expect("write should succeed");
|
||||
let result = read_file(path.to_string_lossy().as_ref(), None, None);
|
||||
assert!(result.is_err());
|
||||
let error = result.unwrap_err();
|
||||
assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
|
||||
assert!(error.to_string().contains("binary"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_oversized_writes() {
|
||||
let path = temp_path("oversize-write.txt");
|
||||
let huge = "x".repeat(MAX_WRITE_SIZE + 1);
|
||||
let result = write_file(path.to_string_lossy().as_ref(), &huge);
|
||||
assert!(result.is_err());
|
||||
let error = result.unwrap_err();
|
||||
assert_eq!(error.kind(), std::io::ErrorKind::InvalidData);
|
||||
assert!(error.to_string().contains("too large"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enforces_workspace_boundary() {
|
||||
let workspace = temp_path("workspace-boundary");
|
||||
std::fs::create_dir_all(&workspace).expect("workspace dir should be created");
|
||||
let inside = workspace.join("inside.txt");
|
||||
write_file(inside.to_string_lossy().as_ref(), "safe content")
|
||||
.expect("write inside workspace should succeed");
|
||||
|
||||
// Reading inside workspace should succeed
|
||||
let result =
|
||||
read_file_in_workspace(inside.to_string_lossy().as_ref(), None, None, &workspace);
|
||||
assert!(result.is_ok());
|
||||
|
||||
// Reading outside workspace should fail
|
||||
let outside = temp_path("outside-boundary.txt");
|
||||
write_file(outside.to_string_lossy().as_ref(), "unsafe content")
|
||||
.expect("write outside should succeed");
|
||||
let result =
|
||||
read_file_in_workspace(outside.to_string_lossy().as_ref(), None, None, &workspace);
|
||||
assert!(result.is_err());
|
||||
let error = result.unwrap_err();
|
||||
assert_eq!(error.kind(), std::io::ErrorKind::PermissionDenied);
|
||||
assert!(error.to_string().contains("escapes workspace"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_symlink_escape() {
|
||||
let workspace = temp_path("symlink-workspace");
|
||||
std::fs::create_dir_all(&workspace).expect("workspace dir should be created");
|
||||
let outside = temp_path("symlink-target.txt");
|
||||
std::fs::write(&outside, "target content").expect("target should write");
|
||||
|
||||
let _link_path = workspace.join("escape-link.txt");
|
||||
#[cfg(unix)]
|
||||
{
|
||||
std::os::unix::fs::symlink(&outside, &link_path).expect("symlink should create");
|
||||
assert!(is_symlink_escape(&link_path, &workspace).expect("check should succeed"));
|
||||
}
|
||||
|
||||
// Non-symlink file should not be an escape
|
||||
let normal = workspace.join("normal.txt");
|
||||
std::fs::write(&normal, "normal content").expect("normal file should write");
|
||||
assert!(!is_symlink_escape(&normal, &workspace).expect("check should succeed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn globs_and_greps_directory() {
|
||||
let dir = temp_path("search-dir");
|
||||
@@ -547,4 +789,51 @@ mod tests {
|
||||
.expect("grep should succeed");
|
||||
assert!(grep_output.content.unwrap_or_default().contains("hello"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_braces_no_braces() {
|
||||
assert_eq!(expand_braces("*.rs"), vec!["*.rs"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_braces_single_group() {
|
||||
let mut result = expand_braces("Assets/**/*.{cs,uxml,uss}");
|
||||
result.sort();
|
||||
assert_eq!(
|
||||
result,
|
||||
vec!["Assets/**/*.cs", "Assets/**/*.uss", "Assets/**/*.uxml",]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_braces_nested() {
|
||||
let mut result = expand_braces("src/{a,b}.{rs,toml}");
|
||||
result.sort();
|
||||
assert_eq!(
|
||||
result,
|
||||
vec!["src/a.rs", "src/a.toml", "src/b.rs", "src/b.toml"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expand_braces_unmatched() {
|
||||
assert_eq!(expand_braces("foo.{bar"), vec!["foo.{bar"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn glob_search_with_braces_finds_files() {
|
||||
let dir = temp_path("glob-braces");
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
std::fs::write(dir.join("a.rs"), "fn main() {}").unwrap();
|
||||
std::fs::write(dir.join("b.toml"), "[package]").unwrap();
|
||||
std::fs::write(dir.join("c.txt"), "hello").unwrap();
|
||||
|
||||
let result =
|
||||
glob_search("*.{rs,toml}", Some(dir.to_str().unwrap())).expect("glob should succeed");
|
||||
assert_eq!(
|
||||
result.num_files, 2,
|
||||
"should match .rs and .toml but not .txt"
|
||||
);
|
||||
let _ = std::fs::remove_dir_all(&dir);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,324 @@
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
/// A single git commit entry from the log.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct GitCommitEntry {
|
||||
pub hash: String,
|
||||
pub subject: String,
|
||||
}
|
||||
|
||||
/// Git-aware context gathered at startup for injection into the system prompt.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct GitContext {
|
||||
pub branch: Option<String>,
|
||||
pub recent_commits: Vec<GitCommitEntry>,
|
||||
pub staged_files: Vec<String>,
|
||||
}
|
||||
|
||||
const MAX_RECENT_COMMITS: usize = 5;
|
||||
|
||||
impl GitContext {
|
||||
/// Detect the git context from the given working directory.
|
||||
///
|
||||
/// Returns `None` when the directory is not inside a git repository.
|
||||
#[must_use]
|
||||
pub fn detect(cwd: &Path) -> Option<Self> {
|
||||
// Quick gate: is this a git repo at all?
|
||||
let rev_parse = Command::new("git")
|
||||
.args(["rev-parse", "--is-inside-work-tree"])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
.ok()?;
|
||||
if !rev_parse.status.success() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(Self {
|
||||
branch: read_branch(cwd),
|
||||
recent_commits: read_recent_commits(cwd),
|
||||
staged_files: read_staged_files(cwd),
|
||||
})
|
||||
}
|
||||
|
||||
/// Render a human-readable summary suitable for system-prompt injection.
|
||||
#[must_use]
|
||||
pub fn render(&self) -> String {
|
||||
let mut lines = Vec::new();
|
||||
|
||||
if let Some(branch) = &self.branch {
|
||||
lines.push(format!("Git branch: {branch}"));
|
||||
}
|
||||
|
||||
if !self.recent_commits.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("Recent commits:".to_string());
|
||||
for entry in &self.recent_commits {
|
||||
lines.push(format!(" {} {}", entry.hash, entry.subject));
|
||||
}
|
||||
}
|
||||
|
||||
if !self.staged_files.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("Staged files:".to_string());
|
||||
for file in &self.staged_files {
|
||||
lines.push(format!(" {file}"));
|
||||
}
|
||||
}
|
||||
|
||||
lines.join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
fn read_branch(cwd: &Path) -> Option<String> {
|
||||
let output = Command::new("git")
|
||||
.args(["rev-parse", "--abbrev-ref", "HEAD"])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
.ok()?;
|
||||
if !output.status.success() {
|
||||
return None;
|
||||
}
|
||||
let branch = String::from_utf8(output.stdout).ok()?;
|
||||
let trimmed = branch.trim();
|
||||
if trimmed.is_empty() || trimmed == "HEAD" {
|
||||
None
|
||||
} else {
|
||||
Some(trimmed.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn read_recent_commits(cwd: &Path) -> Vec<GitCommitEntry> {
|
||||
let output = Command::new("git")
|
||||
.args([
|
||||
"--no-optional-locks",
|
||||
"log",
|
||||
"--oneline",
|
||||
"-n",
|
||||
&MAX_RECENT_COMMITS.to_string(),
|
||||
"--no-decorate",
|
||||
])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
.ok();
|
||||
let Some(output) = output else {
|
||||
return Vec::new();
|
||||
};
|
||||
if !output.status.success() {
|
||||
return Vec::new();
|
||||
}
|
||||
let stdout = String::from_utf8(output.stdout).unwrap_or_default();
|
||||
stdout
|
||||
.lines()
|
||||
.filter_map(|line| {
|
||||
let line = line.trim();
|
||||
if line.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let (hash, subject) = line.split_once(' ')?;
|
||||
Some(GitCommitEntry {
|
||||
hash: hash.to_string(),
|
||||
subject: subject.to_string(),
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn read_staged_files(cwd: &Path) -> Vec<String> {
|
||||
let output = Command::new("git")
|
||||
.args(["--no-optional-locks", "diff", "--cached", "--name-only"])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
.ok();
|
||||
let Some(output) = output else {
|
||||
return Vec::new();
|
||||
};
|
||||
if !output.status.success() {
|
||||
return Vec::new();
|
||||
}
|
||||
let stdout = String::from_utf8(output.stdout).unwrap_or_default();
|
||||
stdout
|
||||
.lines()
|
||||
.filter(|line| !line.trim().is_empty())
|
||||
.map(|line| line.trim().to_string())
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{GitCommitEntry, GitContext};
|
||||
use std::fs;
|
||||
use std::process::Command;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_dir(label: &str) -> std::path::PathBuf {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!("runtime-git-context-{label}-{nanos}"))
|
||||
}
|
||||
|
||||
fn env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
crate::test_env_lock()
|
||||
}
|
||||
|
||||
fn ensure_valid_cwd() {
|
||||
if std::env::current_dir().is_err() {
|
||||
std::env::set_current_dir(env!("CARGO_MANIFEST_DIR"))
|
||||
.expect("test cwd should be recoverable");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_for_non_git_directory() {
|
||||
// given
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir("non-git");
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
|
||||
// when
|
||||
let context = GitContext::detect(&root);
|
||||
|
||||
// then
|
||||
assert!(context.is_none());
|
||||
fs::remove_dir_all(root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_branch_name_and_commits() {
|
||||
// given
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir("branch-commits");
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
git(&root, &["init", "--quiet", "--initial-branch=main"]);
|
||||
git(&root, &["config", "user.email", "tests@example.com"]);
|
||||
git(&root, &["config", "user.name", "Git Context Tests"]);
|
||||
fs::write(root.join("a.txt"), "a\n").expect("write a");
|
||||
git(&root, &["add", "a.txt"]);
|
||||
git(&root, &["commit", "-m", "first commit", "--quiet"]);
|
||||
fs::write(root.join("b.txt"), "b\n").expect("write b");
|
||||
git(&root, &["add", "b.txt"]);
|
||||
git(&root, &["commit", "-m", "second commit", "--quiet"]);
|
||||
|
||||
// when
|
||||
let context = GitContext::detect(&root).expect("should detect git repo");
|
||||
|
||||
// then
|
||||
assert_eq!(context.branch.as_deref(), Some("main"));
|
||||
assert_eq!(context.recent_commits.len(), 2);
|
||||
assert_eq!(context.recent_commits[0].subject, "second commit");
|
||||
assert_eq!(context.recent_commits[1].subject, "first commit");
|
||||
assert!(context.staged_files.is_empty());
|
||||
fs::remove_dir_all(root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_staged_files() {
|
||||
// given
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir("staged");
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
git(&root, &["init", "--quiet", "--initial-branch=main"]);
|
||||
git(&root, &["config", "user.email", "tests@example.com"]);
|
||||
git(&root, &["config", "user.name", "Git Context Tests"]);
|
||||
fs::write(root.join("init.txt"), "init\n").expect("write init");
|
||||
git(&root, &["add", "init.txt"]);
|
||||
git(&root, &["commit", "-m", "initial", "--quiet"]);
|
||||
fs::write(root.join("staged.txt"), "staged\n").expect("write staged");
|
||||
git(&root, &["add", "staged.txt"]);
|
||||
|
||||
// when
|
||||
let context = GitContext::detect(&root).expect("should detect git repo");
|
||||
|
||||
// then
|
||||
assert_eq!(context.staged_files, vec!["staged.txt"]);
|
||||
fs::remove_dir_all(root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_formats_all_sections() {
|
||||
// given
|
||||
let context = GitContext {
|
||||
branch: Some("feat/test".to_string()),
|
||||
recent_commits: vec![
|
||||
GitCommitEntry {
|
||||
hash: "abc1234".to_string(),
|
||||
subject: "add feature".to_string(),
|
||||
},
|
||||
GitCommitEntry {
|
||||
hash: "def5678".to_string(),
|
||||
subject: "fix bug".to_string(),
|
||||
},
|
||||
],
|
||||
staged_files: vec!["src/main.rs".to_string()],
|
||||
};
|
||||
|
||||
// when
|
||||
let rendered = context.render();
|
||||
|
||||
// then
|
||||
assert!(rendered.contains("Git branch: feat/test"));
|
||||
assert!(rendered.contains("abc1234 add feature"));
|
||||
assert!(rendered.contains("def5678 fix bug"));
|
||||
assert!(rendered.contains("src/main.rs"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn render_omits_empty_sections() {
|
||||
// given
|
||||
let context = GitContext {
|
||||
branch: Some("main".to_string()),
|
||||
recent_commits: Vec::new(),
|
||||
staged_files: Vec::new(),
|
||||
};
|
||||
|
||||
// when
|
||||
let rendered = context.render();
|
||||
|
||||
// then
|
||||
assert!(rendered.contains("Git branch: main"));
|
||||
assert!(!rendered.contains("Recent commits:"));
|
||||
assert!(!rendered.contains("Staged files:"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn limits_to_five_recent_commits() {
|
||||
// given
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir("five-commits");
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
git(&root, &["init", "--quiet", "--initial-branch=main"]);
|
||||
git(&root, &["config", "user.email", "tests@example.com"]);
|
||||
git(&root, &["config", "user.name", "Git Context Tests"]);
|
||||
for i in 1..=8 {
|
||||
let name = format!("file{i}.txt");
|
||||
fs::write(root.join(&name), format!("{i}\n")).expect("write file");
|
||||
git(&root, &["add", &name]);
|
||||
git(&root, &["commit", "-m", &format!("commit {i}"), "--quiet"]);
|
||||
}
|
||||
|
||||
// when
|
||||
let context = GitContext::detect(&root).expect("should detect git repo");
|
||||
|
||||
// then
|
||||
assert_eq!(context.recent_commits.len(), 5);
|
||||
assert_eq!(context.recent_commits[0].subject, "commit 8");
|
||||
assert_eq!(context.recent_commits[4].subject, "commit 4");
|
||||
fs::remove_dir_all(root).expect("cleanup");
|
||||
}
|
||||
|
||||
fn git(cwd: &std::path::Path, args: &[&str]) {
|
||||
let status = Command::new("git")
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
.unwrap_or_else(|_| panic!("git {args:?} should run"))
|
||||
.status;
|
||||
assert!(status.success(), "git {args:?} failed");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum GreenLevel {
|
||||
TargetedTests,
|
||||
Package,
|
||||
Workspace,
|
||||
MergeReady,
|
||||
}
|
||||
|
||||
impl GreenLevel {
|
||||
#[must_use]
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::TargetedTests => "targeted_tests",
|
||||
Self::Package => "package",
|
||||
Self::Workspace => "workspace",
|
||||
Self::MergeReady => "merge_ready",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for GreenLevel {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct GreenContract {
|
||||
pub required_level: GreenLevel,
|
||||
}
|
||||
|
||||
impl GreenContract {
|
||||
#[must_use]
|
||||
pub fn new(required_level: GreenLevel) -> Self {
|
||||
Self { required_level }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn evaluate(self, observed_level: Option<GreenLevel>) -> GreenContractOutcome {
|
||||
match observed_level {
|
||||
Some(level) if level >= self.required_level => GreenContractOutcome::Satisfied {
|
||||
required_level: self.required_level,
|
||||
observed_level: level,
|
||||
},
|
||||
_ => GreenContractOutcome::Unsatisfied {
|
||||
required_level: self.required_level,
|
||||
observed_level,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_satisfied_by(self, observed_level: GreenLevel) -> bool {
|
||||
observed_level >= self.required_level
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "outcome", rename_all = "snake_case")]
|
||||
pub enum GreenContractOutcome {
|
||||
Satisfied {
|
||||
required_level: GreenLevel,
|
||||
observed_level: GreenLevel,
|
||||
},
|
||||
Unsatisfied {
|
||||
required_level: GreenLevel,
|
||||
observed_level: Option<GreenLevel>,
|
||||
},
|
||||
}
|
||||
|
||||
impl GreenContractOutcome {
|
||||
#[must_use]
|
||||
pub fn is_satisfied(&self) -> bool {
|
||||
matches!(self, Self::Satisfied { .. })
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn given_matching_level_when_evaluating_contract_then_it_is_satisfied() {
|
||||
// given
|
||||
let contract = GreenContract::new(GreenLevel::Package);
|
||||
|
||||
// when
|
||||
let outcome = contract.evaluate(Some(GreenLevel::Package));
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
outcome,
|
||||
GreenContractOutcome::Satisfied {
|
||||
required_level: GreenLevel::Package,
|
||||
observed_level: GreenLevel::Package,
|
||||
}
|
||||
);
|
||||
assert!(outcome.is_satisfied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_higher_level_when_checking_requirement_then_it_still_satisfies_contract() {
|
||||
// given
|
||||
let contract = GreenContract::new(GreenLevel::TargetedTests);
|
||||
|
||||
// when
|
||||
let is_satisfied = contract.is_satisfied_by(GreenLevel::Workspace);
|
||||
|
||||
// then
|
||||
assert!(is_satisfied);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_lower_level_when_evaluating_contract_then_it_is_unsatisfied() {
|
||||
// given
|
||||
let contract = GreenContract::new(GreenLevel::Workspace);
|
||||
|
||||
// when
|
||||
let outcome = contract.evaluate(Some(GreenLevel::Package));
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
outcome,
|
||||
GreenContractOutcome::Unsatisfied {
|
||||
required_level: GreenLevel::Workspace,
|
||||
observed_level: Some(GreenLevel::Package),
|
||||
}
|
||||
);
|
||||
assert!(!outcome.is_satisfied());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_no_green_level_when_evaluating_contract_then_contract_is_unsatisfied() {
|
||||
// given
|
||||
let contract = GreenContract::new(GreenLevel::MergeReady);
|
||||
|
||||
// when
|
||||
let outcome = contract.evaluate(None);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
outcome,
|
||||
GreenContractOutcome::Unsatisfied {
|
||||
required_level: GreenLevel::MergeReady,
|
||||
observed_level: None,
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,988 @@
|
||||
use std::ffi::OsStr;
|
||||
use std::io::Write;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::sync::{
|
||||
atomic::{AtomicBool, Ordering},
|
||||
Arc,
|
||||
};
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::config::{RuntimeFeatureConfig, RuntimeHookConfig};
|
||||
use crate::permissions::PermissionOverride;
|
||||
|
||||
pub type HookPermissionDecision = PermissionOverride;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum HookEvent {
|
||||
PreToolUse,
|
||||
PostToolUse,
|
||||
PostToolUseFailure,
|
||||
}
|
||||
|
||||
impl HookEvent {
|
||||
#[must_use]
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::PreToolUse => "PreToolUse",
|
||||
Self::PostToolUse => "PostToolUse",
|
||||
Self::PostToolUseFailure => "PostToolUseFailure",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum HookProgressEvent {
|
||||
Started {
|
||||
event: HookEvent,
|
||||
tool_name: String,
|
||||
command: String,
|
||||
},
|
||||
Completed {
|
||||
event: HookEvent,
|
||||
tool_name: String,
|
||||
command: String,
|
||||
},
|
||||
Cancelled {
|
||||
event: HookEvent,
|
||||
tool_name: String,
|
||||
command: String,
|
||||
},
|
||||
}
|
||||
|
||||
pub trait HookProgressReporter: Send {
|
||||
fn on_event(&mut self, event: &HookProgressEvent);
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct HookAbortSignal {
|
||||
aborted: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl HookAbortSignal {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn abort(&self) {
|
||||
self.aborted.store(true, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_aborted(&self) -> bool {
|
||||
self.aborted.load(Ordering::SeqCst)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct HookRunResult {
|
||||
denied: bool,
|
||||
failed: bool,
|
||||
cancelled: bool,
|
||||
messages: Vec<String>,
|
||||
permission_override: Option<PermissionOverride>,
|
||||
permission_reason: Option<String>,
|
||||
updated_input: Option<String>,
|
||||
}
|
||||
|
||||
impl HookRunResult {
|
||||
#[must_use]
|
||||
pub fn allow(messages: Vec<String>) -> Self {
|
||||
Self {
|
||||
denied: false,
|
||||
failed: false,
|
||||
cancelled: false,
|
||||
messages,
|
||||
permission_override: None,
|
||||
permission_reason: None,
|
||||
updated_input: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_denied(&self) -> bool {
|
||||
self.denied
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_failed(&self) -> bool {
|
||||
self.failed
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_cancelled(&self) -> bool {
|
||||
self.cancelled
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn messages(&self) -> &[String] {
|
||||
&self.messages
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn permission_override(&self) -> Option<PermissionOverride> {
|
||||
self.permission_override
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn permission_decision(&self) -> Option<HookPermissionDecision> {
|
||||
self.permission_override
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn permission_reason(&self) -> Option<&str> {
|
||||
self.permission_reason.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn updated_input(&self) -> Option<&str> {
|
||||
self.updated_input.as_deref()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn updated_input_json(&self) -> Option<&str> {
|
||||
self.updated_input()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct HookRunner {
|
||||
config: RuntimeHookConfig,
|
||||
}
|
||||
|
||||
impl HookRunner {
|
||||
#[must_use]
|
||||
pub fn new(config: RuntimeHookConfig) -> Self {
|
||||
Self { config }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn from_feature_config(feature_config: &RuntimeFeatureConfig) -> Self {
|
||||
Self::new(feature_config.hooks().clone())
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_pre_tool_use(&self, tool_name: &str, tool_input: &str) -> HookRunResult {
|
||||
self.run_pre_tool_use_with_context(tool_name, tool_input, None, None)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_pre_tool_use_with_context(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
abort_signal: Option<&HookAbortSignal>,
|
||||
reporter: Option<&mut dyn HookProgressReporter>,
|
||||
) -> HookRunResult {
|
||||
Self::run_commands(
|
||||
HookEvent::PreToolUse,
|
||||
self.config.pre_tool_use(),
|
||||
tool_name,
|
||||
tool_input,
|
||||
None,
|
||||
false,
|
||||
abort_signal,
|
||||
reporter,
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_pre_tool_use_with_signal(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
abort_signal: Option<&HookAbortSignal>,
|
||||
) -> HookRunResult {
|
||||
self.run_pre_tool_use_with_context(tool_name, tool_input, abort_signal, None)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_post_tool_use(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: &str,
|
||||
is_error: bool,
|
||||
) -> HookRunResult {
|
||||
self.run_post_tool_use_with_context(
|
||||
tool_name,
|
||||
tool_input,
|
||||
tool_output,
|
||||
is_error,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_post_tool_use_with_context(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: &str,
|
||||
is_error: bool,
|
||||
abort_signal: Option<&HookAbortSignal>,
|
||||
reporter: Option<&mut dyn HookProgressReporter>,
|
||||
) -> HookRunResult {
|
||||
Self::run_commands(
|
||||
HookEvent::PostToolUse,
|
||||
self.config.post_tool_use(),
|
||||
tool_name,
|
||||
tool_input,
|
||||
Some(tool_output),
|
||||
is_error,
|
||||
abort_signal,
|
||||
reporter,
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_post_tool_use_with_signal(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: &str,
|
||||
is_error: bool,
|
||||
abort_signal: Option<&HookAbortSignal>,
|
||||
) -> HookRunResult {
|
||||
self.run_post_tool_use_with_context(
|
||||
tool_name,
|
||||
tool_input,
|
||||
tool_output,
|
||||
is_error,
|
||||
abort_signal,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_post_tool_use_failure(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_error: &str,
|
||||
) -> HookRunResult {
|
||||
self.run_post_tool_use_failure_with_context(tool_name, tool_input, tool_error, None, None)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_post_tool_use_failure_with_context(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_error: &str,
|
||||
abort_signal: Option<&HookAbortSignal>,
|
||||
reporter: Option<&mut dyn HookProgressReporter>,
|
||||
) -> HookRunResult {
|
||||
Self::run_commands(
|
||||
HookEvent::PostToolUseFailure,
|
||||
self.config.post_tool_use_failure(),
|
||||
tool_name,
|
||||
tool_input,
|
||||
Some(tool_error),
|
||||
true,
|
||||
abort_signal,
|
||||
reporter,
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn run_post_tool_use_failure_with_signal(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_error: &str,
|
||||
abort_signal: Option<&HookAbortSignal>,
|
||||
) -> HookRunResult {
|
||||
self.run_post_tool_use_failure_with_context(
|
||||
tool_name,
|
||||
tool_input,
|
||||
tool_error,
|
||||
abort_signal,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn run_commands(
|
||||
event: HookEvent,
|
||||
commands: &[String],
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: Option<&str>,
|
||||
is_error: bool,
|
||||
abort_signal: Option<&HookAbortSignal>,
|
||||
mut reporter: Option<&mut dyn HookProgressReporter>,
|
||||
) -> HookRunResult {
|
||||
if commands.is_empty() {
|
||||
return HookRunResult::allow(Vec::new());
|
||||
}
|
||||
|
||||
if abort_signal.is_some_and(HookAbortSignal::is_aborted) {
|
||||
return HookRunResult {
|
||||
denied: false,
|
||||
failed: false,
|
||||
cancelled: true,
|
||||
messages: vec![format!(
|
||||
"{} hook cancelled before execution",
|
||||
event.as_str()
|
||||
)],
|
||||
permission_override: None,
|
||||
permission_reason: None,
|
||||
updated_input: None,
|
||||
};
|
||||
}
|
||||
|
||||
let payload = hook_payload(event, tool_name, tool_input, tool_output, is_error).to_string();
|
||||
let mut result = HookRunResult::allow(Vec::new());
|
||||
|
||||
for command in commands {
|
||||
if let Some(reporter) = reporter.as_deref_mut() {
|
||||
reporter.on_event(&HookProgressEvent::Started {
|
||||
event,
|
||||
tool_name: tool_name.to_string(),
|
||||
command: command.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
match Self::run_command(
|
||||
command,
|
||||
event,
|
||||
tool_name,
|
||||
tool_input,
|
||||
tool_output,
|
||||
is_error,
|
||||
&payload,
|
||||
abort_signal,
|
||||
) {
|
||||
HookCommandOutcome::Allow { parsed } => {
|
||||
if let Some(reporter) = reporter.as_deref_mut() {
|
||||
reporter.on_event(&HookProgressEvent::Completed {
|
||||
event,
|
||||
tool_name: tool_name.to_string(),
|
||||
command: command.clone(),
|
||||
});
|
||||
}
|
||||
merge_parsed_hook_output(&mut result, parsed);
|
||||
}
|
||||
HookCommandOutcome::Deny { parsed } => {
|
||||
if let Some(reporter) = reporter.as_deref_mut() {
|
||||
reporter.on_event(&HookProgressEvent::Completed {
|
||||
event,
|
||||
tool_name: tool_name.to_string(),
|
||||
command: command.clone(),
|
||||
});
|
||||
}
|
||||
merge_parsed_hook_output(&mut result, parsed);
|
||||
result.denied = true;
|
||||
return result;
|
||||
}
|
||||
HookCommandOutcome::Failed { parsed } => {
|
||||
if let Some(reporter) = reporter.as_deref_mut() {
|
||||
reporter.on_event(&HookProgressEvent::Completed {
|
||||
event,
|
||||
tool_name: tool_name.to_string(),
|
||||
command: command.clone(),
|
||||
});
|
||||
}
|
||||
merge_parsed_hook_output(&mut result, parsed);
|
||||
result.failed = true;
|
||||
return result;
|
||||
}
|
||||
HookCommandOutcome::Cancelled { message } => {
|
||||
if let Some(reporter) = reporter.as_deref_mut() {
|
||||
reporter.on_event(&HookProgressEvent::Cancelled {
|
||||
event,
|
||||
tool_name: tool_name.to_string(),
|
||||
command: command.clone(),
|
||||
});
|
||||
}
|
||||
result.cancelled = true;
|
||||
result.messages.push(message);
|
||||
return result;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn run_command(
|
||||
command: &str,
|
||||
event: HookEvent,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: Option<&str>,
|
||||
is_error: bool,
|
||||
payload: &str,
|
||||
abort_signal: Option<&HookAbortSignal>,
|
||||
) -> HookCommandOutcome {
|
||||
let mut child = shell_command(command);
|
||||
child.stdin(Stdio::piped());
|
||||
child.stdout(Stdio::piped());
|
||||
child.stderr(Stdio::piped());
|
||||
child.env("HOOK_EVENT", event.as_str());
|
||||
child.env("HOOK_TOOL_NAME", tool_name);
|
||||
child.env("HOOK_TOOL_INPUT", tool_input);
|
||||
child.env("HOOK_TOOL_IS_ERROR", if is_error { "1" } else { "0" });
|
||||
if let Some(tool_output) = tool_output {
|
||||
child.env("HOOK_TOOL_OUTPUT", tool_output);
|
||||
}
|
||||
|
||||
match child.output_with_stdin(payload.as_bytes(), abort_signal) {
|
||||
Ok(CommandExecution::Finished(output)) => {
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
let stderr = String::from_utf8_lossy(&output.stderr).trim().to_string();
|
||||
let parsed = parse_hook_output(&stdout);
|
||||
let primary_message = parsed.primary_message().map(ToOwned::to_owned);
|
||||
match output.status.code() {
|
||||
Some(0) => {
|
||||
if parsed.deny {
|
||||
HookCommandOutcome::Deny { parsed }
|
||||
} else {
|
||||
HookCommandOutcome::Allow { parsed }
|
||||
}
|
||||
}
|
||||
Some(2) => HookCommandOutcome::Deny {
|
||||
parsed: parsed.with_fallback_message(format!(
|
||||
"{} hook denied tool `{tool_name}`",
|
||||
event.as_str()
|
||||
)),
|
||||
},
|
||||
Some(code) => HookCommandOutcome::Failed {
|
||||
parsed: parsed.with_fallback_message(format_hook_failure(
|
||||
command,
|
||||
code,
|
||||
primary_message.as_deref(),
|
||||
stderr.as_str(),
|
||||
)),
|
||||
},
|
||||
None => HookCommandOutcome::Failed {
|
||||
parsed: parsed.with_fallback_message(format!(
|
||||
"{} hook `{command}` terminated by signal while handling `{}`",
|
||||
event.as_str(),
|
||||
tool_name
|
||||
)),
|
||||
},
|
||||
}
|
||||
}
|
||||
Ok(CommandExecution::Cancelled) => HookCommandOutcome::Cancelled {
|
||||
message: format!(
|
||||
"{} hook `{command}` cancelled while handling `{tool_name}`",
|
||||
event.as_str()
|
||||
),
|
||||
},
|
||||
Err(error) => HookCommandOutcome::Failed {
|
||||
parsed: ParsedHookOutput {
|
||||
messages: vec![format!(
|
||||
"{} hook `{command}` failed to start for `{}`: {error}",
|
||||
event.as_str(),
|
||||
tool_name
|
||||
)],
|
||||
..ParsedHookOutput::default()
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum HookCommandOutcome {
|
||||
Allow { parsed: ParsedHookOutput },
|
||||
Deny { parsed: ParsedHookOutput },
|
||||
Failed { parsed: ParsedHookOutput },
|
||||
Cancelled { message: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
struct ParsedHookOutput {
|
||||
messages: Vec<String>,
|
||||
deny: bool,
|
||||
permission_override: Option<PermissionOverride>,
|
||||
permission_reason: Option<String>,
|
||||
updated_input: Option<String>,
|
||||
}
|
||||
|
||||
impl ParsedHookOutput {
|
||||
fn with_fallback_message(mut self, fallback: String) -> Self {
|
||||
if self.messages.is_empty() {
|
||||
self.messages.push(fallback);
|
||||
}
|
||||
self
|
||||
}
|
||||
|
||||
fn primary_message(&self) -> Option<&str> {
|
||||
self.messages.first().map(String::as_str)
|
||||
}
|
||||
}
|
||||
|
||||
fn merge_parsed_hook_output(target: &mut HookRunResult, parsed: ParsedHookOutput) {
|
||||
target.messages.extend(parsed.messages);
|
||||
if parsed.permission_override.is_some() {
|
||||
target.permission_override = parsed.permission_override;
|
||||
}
|
||||
if parsed.permission_reason.is_some() {
|
||||
target.permission_reason = parsed.permission_reason;
|
||||
}
|
||||
if parsed.updated_input.is_some() {
|
||||
target.updated_input = parsed.updated_input;
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_hook_output(stdout: &str) -> ParsedHookOutput {
|
||||
if stdout.is_empty() {
|
||||
return ParsedHookOutput::default();
|
||||
}
|
||||
|
||||
let Ok(Value::Object(root)) = serde_json::from_str::<Value>(stdout) else {
|
||||
return ParsedHookOutput {
|
||||
messages: vec![stdout.to_string()],
|
||||
..ParsedHookOutput::default()
|
||||
};
|
||||
};
|
||||
|
||||
let mut parsed = ParsedHookOutput::default();
|
||||
|
||||
if let Some(message) = root.get("systemMessage").and_then(Value::as_str) {
|
||||
parsed.messages.push(message.to_string());
|
||||
}
|
||||
if let Some(message) = root.get("reason").and_then(Value::as_str) {
|
||||
parsed.messages.push(message.to_string());
|
||||
}
|
||||
if root.get("continue").and_then(Value::as_bool) == Some(false)
|
||||
|| root.get("decision").and_then(Value::as_str) == Some("block")
|
||||
{
|
||||
parsed.deny = true;
|
||||
}
|
||||
|
||||
if let Some(Value::Object(specific)) = root.get("hookSpecificOutput") {
|
||||
if let Some(Value::String(additional_context)) = specific.get("additionalContext") {
|
||||
parsed.messages.push(additional_context.clone());
|
||||
}
|
||||
if let Some(decision) = specific.get("permissionDecision").and_then(Value::as_str) {
|
||||
parsed.permission_override = match decision {
|
||||
"allow" => Some(PermissionOverride::Allow),
|
||||
"deny" => Some(PermissionOverride::Deny),
|
||||
"ask" => Some(PermissionOverride::Ask),
|
||||
_ => None,
|
||||
};
|
||||
}
|
||||
if let Some(reason) = specific
|
||||
.get("permissionDecisionReason")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
parsed.permission_reason = Some(reason.to_string());
|
||||
}
|
||||
if let Some(updated_input) = specific.get("updatedInput") {
|
||||
parsed.updated_input = serde_json::to_string(updated_input).ok();
|
||||
}
|
||||
}
|
||||
|
||||
if parsed.messages.is_empty() {
|
||||
parsed.messages.push(stdout.to_string());
|
||||
}
|
||||
|
||||
parsed
|
||||
}
|
||||
|
||||
fn hook_payload(
|
||||
event: HookEvent,
|
||||
tool_name: &str,
|
||||
tool_input: &str,
|
||||
tool_output: Option<&str>,
|
||||
is_error: bool,
|
||||
) -> Value {
|
||||
match event {
|
||||
HookEvent::PostToolUseFailure => json!({
|
||||
"hook_event_name": event.as_str(),
|
||||
"tool_name": tool_name,
|
||||
"tool_input": parse_tool_input(tool_input),
|
||||
"tool_input_json": tool_input,
|
||||
"tool_error": tool_output,
|
||||
"tool_result_is_error": true,
|
||||
}),
|
||||
_ => json!({
|
||||
"hook_event_name": event.as_str(),
|
||||
"tool_name": tool_name,
|
||||
"tool_input": parse_tool_input(tool_input),
|
||||
"tool_input_json": tool_input,
|
||||
"tool_output": tool_output,
|
||||
"tool_result_is_error": is_error,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_tool_input(tool_input: &str) -> Value {
|
||||
serde_json::from_str(tool_input).unwrap_or_else(|_| json!({ "raw": tool_input }))
|
||||
}
|
||||
|
||||
fn format_hook_failure(command: &str, code: i32, stdout: Option<&str>, stderr: &str) -> String {
|
||||
let mut message = format!("Hook `{command}` exited with status {code}");
|
||||
if let Some(stdout) = stdout.filter(|stdout| !stdout.is_empty()) {
|
||||
message.push_str(": ");
|
||||
message.push_str(stdout);
|
||||
} else if !stderr.is_empty() {
|
||||
message.push_str(": ");
|
||||
message.push_str(stderr);
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
fn shell_command(command: &str) -> CommandWithStdin {
|
||||
#[cfg(windows)]
|
||||
#[allow(unused_mut)]
|
||||
let mut command_builder = {
|
||||
let mut command_builder = Command::new("cmd");
|
||||
command_builder.arg("/C").arg(command);
|
||||
CommandWithStdin::new(command_builder)
|
||||
};
|
||||
|
||||
#[cfg(not(windows))]
|
||||
let command_builder = {
|
||||
let mut command_builder = Command::new("sh");
|
||||
command_builder.arg("-lc").arg(command);
|
||||
CommandWithStdin::new(command_builder)
|
||||
};
|
||||
|
||||
command_builder
|
||||
}
|
||||
|
||||
struct CommandWithStdin {
|
||||
command: Command,
|
||||
}
|
||||
|
||||
impl CommandWithStdin {
|
||||
fn new(command: Command) -> Self {
|
||||
Self { command }
|
||||
}
|
||||
|
||||
fn stdin(&mut self, cfg: Stdio) -> &mut Self {
|
||||
self.command.stdin(cfg);
|
||||
self
|
||||
}
|
||||
|
||||
fn stdout(&mut self, cfg: Stdio) -> &mut Self {
|
||||
self.command.stdout(cfg);
|
||||
self
|
||||
}
|
||||
|
||||
fn stderr(&mut self, cfg: Stdio) -> &mut Self {
|
||||
self.command.stderr(cfg);
|
||||
self
|
||||
}
|
||||
|
||||
fn env<K, V>(&mut self, key: K, value: V) -> &mut Self
|
||||
where
|
||||
K: AsRef<OsStr>,
|
||||
V: AsRef<OsStr>,
|
||||
{
|
||||
self.command.env(key, value);
|
||||
self
|
||||
}
|
||||
|
||||
fn output_with_stdin(
|
||||
&mut self,
|
||||
stdin: &[u8],
|
||||
abort_signal: Option<&HookAbortSignal>,
|
||||
) -> std::io::Result<CommandExecution> {
|
||||
let mut child = self.command.spawn()?;
|
||||
if let Some(mut child_stdin) = child.stdin.take() {
|
||||
child_stdin.write_all(stdin)?;
|
||||
}
|
||||
|
||||
loop {
|
||||
if abort_signal.is_some_and(HookAbortSignal::is_aborted) {
|
||||
let _ = child.kill();
|
||||
let _ = child.wait_with_output();
|
||||
return Ok(CommandExecution::Cancelled);
|
||||
}
|
||||
|
||||
match child.try_wait()? {
|
||||
Some(_) => return child.wait_with_output().map(CommandExecution::Finished),
|
||||
None => thread::sleep(Duration::from_millis(20)),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum CommandExecution {
|
||||
Finished(std::process::Output),
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::thread;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{
|
||||
HookAbortSignal, HookEvent, HookProgressEvent, HookProgressReporter, HookRunResult,
|
||||
HookRunner,
|
||||
};
|
||||
use crate::config::{RuntimeFeatureConfig, RuntimeHookConfig};
|
||||
use crate::permissions::PermissionOverride;
|
||||
|
||||
struct RecordingReporter {
|
||||
events: Vec<HookProgressEvent>,
|
||||
}
|
||||
|
||||
impl HookProgressReporter for RecordingReporter {
|
||||
fn on_event(&mut self, event: &HookProgressEvent) {
|
||||
self.events.push(event.clone());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allows_exit_code_zero_and_captures_stdout() {
|
||||
let runner = HookRunner::new(RuntimeHookConfig::new(
|
||||
vec![shell_snippet("printf 'pre ok'")],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
|
||||
let result = runner.run_pre_tool_use("Read", r#"{"path":"README.md"}"#);
|
||||
|
||||
assert_eq!(result, HookRunResult::allow(vec!["pre ok".to_string()]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn denies_exit_code_two() {
|
||||
let runner = HookRunner::new(RuntimeHookConfig::new(
|
||||
vec![shell_snippet("printf 'blocked by hook'; exit 2")],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
|
||||
let result = runner.run_pre_tool_use("Bash", r#"{"command":"pwd"}"#);
|
||||
|
||||
assert!(result.is_denied());
|
||||
assert_eq!(result.messages(), &["blocked by hook".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn propagates_other_non_zero_statuses_as_failures() {
|
||||
let runner = HookRunner::from_feature_config(&RuntimeFeatureConfig::default().with_hooks(
|
||||
RuntimeHookConfig::new(
|
||||
vec![shell_snippet("printf 'warning hook'; exit 1")],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
),
|
||||
));
|
||||
|
||||
// given
|
||||
// when
|
||||
let result = runner.run_pre_tool_use("Edit", r#"{"file":"src/lib.rs"}"#);
|
||||
|
||||
// then
|
||||
assert!(result.is_failed());
|
||||
assert!(result
|
||||
.messages()
|
||||
.iter()
|
||||
.any(|message| message.contains("warning hook")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_pre_hook_permission_override_and_updated_input() {
|
||||
let runner = HookRunner::new(RuntimeHookConfig::new(
|
||||
vec![shell_snippet(
|
||||
r#"printf '%s' '{"systemMessage":"updated","hookSpecificOutput":{"permissionDecision":"allow","permissionDecisionReason":"hook ok","updatedInput":{"command":"git status"}}}'"#,
|
||||
)],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
|
||||
let result = runner.run_pre_tool_use("bash", r#"{"command":"pwd"}"#);
|
||||
|
||||
assert_eq!(
|
||||
result.permission_override(),
|
||||
Some(PermissionOverride::Allow)
|
||||
);
|
||||
assert_eq!(result.permission_reason(), Some("hook ok"));
|
||||
assert_eq!(result.updated_input(), Some(r#"{"command":"git status"}"#));
|
||||
assert!(result.messages().iter().any(|message| message == "updated"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runs_post_tool_use_failure_hooks() {
|
||||
// given
|
||||
let runner = HookRunner::new(RuntimeHookConfig::new(
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
vec![shell_snippet("printf 'failure hook ran'")],
|
||||
));
|
||||
|
||||
// when
|
||||
let result =
|
||||
runner.run_post_tool_use_failure("bash", r#"{"command":"false"}"#, "command failed");
|
||||
|
||||
// then
|
||||
assert!(!result.is_denied());
|
||||
assert_eq!(result.messages(), &["failure hook ran".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stops_running_failure_hooks_after_failure() {
|
||||
// given
|
||||
let runner = HookRunner::new(RuntimeHookConfig::new(
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
vec![
|
||||
shell_snippet("printf 'broken failure hook'; exit 1"),
|
||||
shell_snippet("printf 'later failure hook'"),
|
||||
],
|
||||
));
|
||||
|
||||
// when
|
||||
let result =
|
||||
runner.run_post_tool_use_failure("bash", r#"{"command":"false"}"#, "command failed");
|
||||
|
||||
// then
|
||||
assert!(result.is_failed());
|
||||
assert!(result
|
||||
.messages()
|
||||
.iter()
|
||||
.any(|message| message.contains("broken failure hook")));
|
||||
assert!(!result
|
||||
.messages()
|
||||
.iter()
|
||||
.any(|message| message == "later failure hook"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn executes_hooks_in_configured_order() {
|
||||
// given
|
||||
let runner = HookRunner::new(RuntimeHookConfig::new(
|
||||
vec![
|
||||
shell_snippet("printf 'first'"),
|
||||
shell_snippet("printf 'second'"),
|
||||
],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let mut reporter = RecordingReporter { events: Vec::new() };
|
||||
|
||||
// when
|
||||
let result = runner.run_pre_tool_use_with_context(
|
||||
"Read",
|
||||
r#"{"path":"README.md"}"#,
|
||||
None,
|
||||
Some(&mut reporter),
|
||||
);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
result,
|
||||
HookRunResult::allow(vec!["first".to_string(), "second".to_string()])
|
||||
);
|
||||
assert_eq!(reporter.events.len(), 4);
|
||||
assert!(matches!(
|
||||
&reporter.events[0],
|
||||
HookProgressEvent::Started {
|
||||
event: HookEvent::PreToolUse,
|
||||
command,
|
||||
..
|
||||
} if command == "printf 'first'"
|
||||
));
|
||||
assert!(matches!(
|
||||
&reporter.events[1],
|
||||
HookProgressEvent::Completed {
|
||||
event: HookEvent::PreToolUse,
|
||||
command,
|
||||
..
|
||||
} if command == "printf 'first'"
|
||||
));
|
||||
assert!(matches!(
|
||||
&reporter.events[2],
|
||||
HookProgressEvent::Started {
|
||||
event: HookEvent::PreToolUse,
|
||||
command,
|
||||
..
|
||||
} if command == "printf 'second'"
|
||||
));
|
||||
assert!(matches!(
|
||||
&reporter.events[3],
|
||||
HookProgressEvent::Completed {
|
||||
event: HookEvent::PreToolUse,
|
||||
command,
|
||||
..
|
||||
} if command == "printf 'second'"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stops_running_hooks_after_failure() {
|
||||
// given
|
||||
let runner = HookRunner::new(RuntimeHookConfig::new(
|
||||
vec![
|
||||
shell_snippet("printf 'broken'; exit 1"),
|
||||
shell_snippet("printf 'later'"),
|
||||
],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
|
||||
// when
|
||||
let result = runner.run_pre_tool_use("Edit", r#"{"file":"src/lib.rs"}"#);
|
||||
|
||||
// then
|
||||
assert!(result.is_failed());
|
||||
assert!(result
|
||||
.messages()
|
||||
.iter()
|
||||
.any(|message| message.contains("broken")));
|
||||
assert!(!result.messages().iter().any(|message| message == "later"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn abort_signal_cancels_long_running_hook_and_reports_progress() {
|
||||
let runner = HookRunner::new(RuntimeHookConfig::new(
|
||||
vec![shell_snippet("sleep 5")],
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
));
|
||||
let abort_signal = HookAbortSignal::new();
|
||||
let abort_signal_for_thread = abort_signal.clone();
|
||||
let mut reporter = RecordingReporter { events: Vec::new() };
|
||||
|
||||
thread::spawn(move || {
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
abort_signal_for_thread.abort();
|
||||
});
|
||||
|
||||
let result = runner.run_pre_tool_use_with_context(
|
||||
"bash",
|
||||
r#"{"command":"sleep 5"}"#,
|
||||
Some(&abort_signal),
|
||||
Some(&mut reporter),
|
||||
);
|
||||
|
||||
assert!(result.is_cancelled());
|
||||
assert!(reporter.events.iter().any(|event| matches!(
|
||||
event,
|
||||
HookProgressEvent::Started {
|
||||
event: HookEvent::PreToolUse,
|
||||
..
|
||||
}
|
||||
)));
|
||||
assert!(reporter.events.iter().any(|event| matches!(
|
||||
event,
|
||||
HookProgressEvent::Cancelled {
|
||||
event: HookEvent::PreToolUse,
|
||||
..
|
||||
}
|
||||
)));
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn shell_snippet(script: &str) -> String {
|
||||
script.replace('\'', "\"")
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
fn shell_snippet(script: &str) -> String {
|
||||
script.to_string()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,383 @@
|
||||
#![allow(clippy::similar_names)]
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum LaneEventName {
|
||||
#[serde(rename = "lane.started")]
|
||||
Started,
|
||||
#[serde(rename = "lane.ready")]
|
||||
Ready,
|
||||
#[serde(rename = "lane.prompt_misdelivery")]
|
||||
PromptMisdelivery,
|
||||
#[serde(rename = "lane.blocked")]
|
||||
Blocked,
|
||||
#[serde(rename = "lane.red")]
|
||||
Red,
|
||||
#[serde(rename = "lane.green")]
|
||||
Green,
|
||||
#[serde(rename = "lane.commit.created")]
|
||||
CommitCreated,
|
||||
#[serde(rename = "lane.pr.opened")]
|
||||
PrOpened,
|
||||
#[serde(rename = "lane.merge.ready")]
|
||||
MergeReady,
|
||||
#[serde(rename = "lane.finished")]
|
||||
Finished,
|
||||
#[serde(rename = "lane.failed")]
|
||||
Failed,
|
||||
#[serde(rename = "lane.reconciled")]
|
||||
Reconciled,
|
||||
#[serde(rename = "lane.merged")]
|
||||
Merged,
|
||||
#[serde(rename = "lane.superseded")]
|
||||
Superseded,
|
||||
#[serde(rename = "lane.closed")]
|
||||
Closed,
|
||||
#[serde(rename = "branch.stale_against_main")]
|
||||
BranchStaleAgainstMain,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum LaneEventStatus {
|
||||
Running,
|
||||
Ready,
|
||||
Blocked,
|
||||
Red,
|
||||
Green,
|
||||
Completed,
|
||||
Failed,
|
||||
Reconciled,
|
||||
Merged,
|
||||
Superseded,
|
||||
Closed,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum LaneFailureClass {
|
||||
PromptDelivery,
|
||||
TrustGate,
|
||||
BranchDivergence,
|
||||
Compile,
|
||||
Test,
|
||||
PluginStartup,
|
||||
McpStartup,
|
||||
McpHandshake,
|
||||
GatewayRouting,
|
||||
ToolRuntime,
|
||||
Infra,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LaneEventBlocker {
|
||||
#[serde(rename = "failureClass")]
|
||||
pub failure_class: LaneFailureClass,
|
||||
pub detail: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LaneCommitProvenance {
|
||||
pub commit: String,
|
||||
pub branch: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub worktree: Option<String>,
|
||||
#[serde(rename = "canonicalCommit", skip_serializing_if = "Option::is_none")]
|
||||
pub canonical_commit: Option<String>,
|
||||
#[serde(rename = "supersededBy", skip_serializing_if = "Option::is_none")]
|
||||
pub superseded_by: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub lineage: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct LaneEvent {
|
||||
pub event: LaneEventName,
|
||||
pub status: LaneEventStatus,
|
||||
#[serde(rename = "emittedAt")]
|
||||
pub emitted_at: String,
|
||||
#[serde(rename = "failureClass", skip_serializing_if = "Option::is_none")]
|
||||
pub failure_class: Option<LaneFailureClass>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub detail: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub data: Option<Value>,
|
||||
}
|
||||
|
||||
impl LaneEvent {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
event: LaneEventName,
|
||||
status: LaneEventStatus,
|
||||
emitted_at: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
event,
|
||||
status,
|
||||
emitted_at: emitted_at.into(),
|
||||
failure_class: None,
|
||||
detail: None,
|
||||
data: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn started(emitted_at: impl Into<String>) -> Self {
|
||||
Self::new(LaneEventName::Started, LaneEventStatus::Running, emitted_at)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn finished(emitted_at: impl Into<String>, detail: Option<String>) -> Self {
|
||||
Self::new(
|
||||
LaneEventName::Finished,
|
||||
LaneEventStatus::Completed,
|
||||
emitted_at,
|
||||
)
|
||||
.with_optional_detail(detail)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn commit_created(
|
||||
emitted_at: impl Into<String>,
|
||||
detail: Option<String>,
|
||||
provenance: LaneCommitProvenance,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
LaneEventName::CommitCreated,
|
||||
LaneEventStatus::Completed,
|
||||
emitted_at,
|
||||
)
|
||||
.with_optional_detail(detail)
|
||||
.with_data(serde_json::to_value(provenance).expect("commit provenance should serialize"))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn superseded(
|
||||
emitted_at: impl Into<String>,
|
||||
detail: Option<String>,
|
||||
provenance: LaneCommitProvenance,
|
||||
) -> Self {
|
||||
Self::new(
|
||||
LaneEventName::Superseded,
|
||||
LaneEventStatus::Superseded,
|
||||
emitted_at,
|
||||
)
|
||||
.with_optional_detail(detail)
|
||||
.with_data(serde_json::to_value(provenance).expect("commit provenance should serialize"))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn blocked(emitted_at: impl Into<String>, blocker: &LaneEventBlocker) -> Self {
|
||||
Self::new(LaneEventName::Blocked, LaneEventStatus::Blocked, emitted_at)
|
||||
.with_failure_class(blocker.failure_class)
|
||||
.with_detail(blocker.detail.clone())
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn failed(emitted_at: impl Into<String>, blocker: &LaneEventBlocker) -> Self {
|
||||
Self::new(LaneEventName::Failed, LaneEventStatus::Failed, emitted_at)
|
||||
.with_failure_class(blocker.failure_class)
|
||||
.with_detail(blocker.detail.clone())
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_failure_class(mut self, failure_class: LaneFailureClass) -> Self {
|
||||
self.failure_class = Some(failure_class);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_detail(mut self, detail: impl Into<String>) -> Self {
|
||||
self.detail = Some(detail.into());
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_optional_detail(mut self, detail: Option<String>) -> Self {
|
||||
self.detail = detail;
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_data(mut self, data: Value) -> Self {
|
||||
self.data = Some(data);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn dedupe_superseded_commit_events(events: &[LaneEvent]) -> Vec<LaneEvent> {
|
||||
let mut keep = vec![true; events.len()];
|
||||
let mut latest_by_key = std::collections::BTreeMap::<String, usize>::new();
|
||||
|
||||
for (index, event) in events.iter().enumerate() {
|
||||
if event.event != LaneEventName::CommitCreated {
|
||||
continue;
|
||||
}
|
||||
let Some(data) = event.data.as_ref() else {
|
||||
continue;
|
||||
};
|
||||
let key = data
|
||||
.get("canonicalCommit")
|
||||
.or_else(|| data.get("commit"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::to_string);
|
||||
let superseded = data
|
||||
.get("supersededBy")
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.is_some();
|
||||
if superseded {
|
||||
keep[index] = false;
|
||||
continue;
|
||||
}
|
||||
if let Some(key) = key {
|
||||
if let Some(previous) = latest_by_key.insert(key, index) {
|
||||
keep[previous] = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
events
|
||||
.iter()
|
||||
.cloned()
|
||||
.zip(keep)
|
||||
.filter_map(|(event, retain)| retain.then_some(event))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::{
|
||||
dedupe_superseded_commit_events, LaneCommitProvenance, LaneEvent, LaneEventBlocker,
|
||||
LaneEventName, LaneEventStatus, LaneFailureClass,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn canonical_lane_event_names_serialize_to_expected_wire_values() {
|
||||
let cases = [
|
||||
(LaneEventName::Started, "lane.started"),
|
||||
(LaneEventName::Ready, "lane.ready"),
|
||||
(LaneEventName::PromptMisdelivery, "lane.prompt_misdelivery"),
|
||||
(LaneEventName::Blocked, "lane.blocked"),
|
||||
(LaneEventName::Red, "lane.red"),
|
||||
(LaneEventName::Green, "lane.green"),
|
||||
(LaneEventName::CommitCreated, "lane.commit.created"),
|
||||
(LaneEventName::PrOpened, "lane.pr.opened"),
|
||||
(LaneEventName::MergeReady, "lane.merge.ready"),
|
||||
(LaneEventName::Finished, "lane.finished"),
|
||||
(LaneEventName::Failed, "lane.failed"),
|
||||
(LaneEventName::Reconciled, "lane.reconciled"),
|
||||
(LaneEventName::Merged, "lane.merged"),
|
||||
(LaneEventName::Superseded, "lane.superseded"),
|
||||
(LaneEventName::Closed, "lane.closed"),
|
||||
(
|
||||
LaneEventName::BranchStaleAgainstMain,
|
||||
"branch.stale_against_main",
|
||||
),
|
||||
];
|
||||
|
||||
for (event, expected) in cases {
|
||||
assert_eq!(
|
||||
serde_json::to_value(event).expect("serialize event"),
|
||||
json!(expected)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failure_classes_cover_canonical_taxonomy_wire_values() {
|
||||
let cases = [
|
||||
(LaneFailureClass::PromptDelivery, "prompt_delivery"),
|
||||
(LaneFailureClass::TrustGate, "trust_gate"),
|
||||
(LaneFailureClass::BranchDivergence, "branch_divergence"),
|
||||
(LaneFailureClass::Compile, "compile"),
|
||||
(LaneFailureClass::Test, "test"),
|
||||
(LaneFailureClass::PluginStartup, "plugin_startup"),
|
||||
(LaneFailureClass::McpStartup, "mcp_startup"),
|
||||
(LaneFailureClass::McpHandshake, "mcp_handshake"),
|
||||
(LaneFailureClass::GatewayRouting, "gateway_routing"),
|
||||
(LaneFailureClass::ToolRuntime, "tool_runtime"),
|
||||
(LaneFailureClass::Infra, "infra"),
|
||||
];
|
||||
|
||||
for (failure_class, expected) in cases {
|
||||
assert_eq!(
|
||||
serde_json::to_value(failure_class).expect("serialize failure class"),
|
||||
json!(expected)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn blocked_and_failed_events_reuse_blocker_details() {
|
||||
let blocker = LaneEventBlocker {
|
||||
failure_class: LaneFailureClass::McpStartup,
|
||||
detail: "broken server".to_string(),
|
||||
};
|
||||
|
||||
let blocked = LaneEvent::blocked("2026-04-04T00:00:00Z", &blocker);
|
||||
let failed = LaneEvent::failed("2026-04-04T00:00:01Z", &blocker);
|
||||
|
||||
assert_eq!(blocked.event, LaneEventName::Blocked);
|
||||
assert_eq!(blocked.status, LaneEventStatus::Blocked);
|
||||
assert_eq!(blocked.failure_class, Some(LaneFailureClass::McpStartup));
|
||||
assert_eq!(failed.event, LaneEventName::Failed);
|
||||
assert_eq!(failed.status, LaneEventStatus::Failed);
|
||||
assert_eq!(failed.detail.as_deref(), Some("broken server"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn commit_events_can_carry_worktree_and_supersession_metadata() {
|
||||
let event = LaneEvent::commit_created(
|
||||
"2026-04-04T00:00:00Z",
|
||||
Some("commit created".to_string()),
|
||||
LaneCommitProvenance {
|
||||
commit: "abc123".to_string(),
|
||||
branch: "feature/provenance".to_string(),
|
||||
worktree: Some("wt-a".to_string()),
|
||||
canonical_commit: Some("abc123".to_string()),
|
||||
superseded_by: None,
|
||||
lineage: vec!["abc123".to_string()],
|
||||
},
|
||||
);
|
||||
let event_json = serde_json::to_value(&event).expect("lane event should serialize");
|
||||
assert_eq!(event_json["event"], "lane.commit.created");
|
||||
assert_eq!(event_json["data"]["branch"], "feature/provenance");
|
||||
assert_eq!(event_json["data"]["worktree"], "wt-a");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedupes_superseded_commit_events_by_canonical_commit() {
|
||||
let retained = dedupe_superseded_commit_events(&[
|
||||
LaneEvent::commit_created(
|
||||
"2026-04-04T00:00:00Z",
|
||||
Some("old".to_string()),
|
||||
LaneCommitProvenance {
|
||||
commit: "old123".to_string(),
|
||||
branch: "feature/provenance".to_string(),
|
||||
worktree: Some("wt-a".to_string()),
|
||||
canonical_commit: Some("canon123".to_string()),
|
||||
superseded_by: Some("new123".to_string()),
|
||||
lineage: vec!["old123".to_string(), "new123".to_string()],
|
||||
},
|
||||
),
|
||||
LaneEvent::commit_created(
|
||||
"2026-04-04T00:00:01Z",
|
||||
Some("new".to_string()),
|
||||
LaneCommitProvenance {
|
||||
commit: "new123".to_string(),
|
||||
branch: "feature/provenance".to_string(),
|
||||
worktree: Some("wt-b".to_string()),
|
||||
canonical_commit: Some("canon123".to_string()),
|
||||
superseded_by: None,
|
||||
lineage: vec!["old123".to_string(), "new123".to_string()],
|
||||
},
|
||||
),
|
||||
]);
|
||||
assert_eq!(retained.len(), 1);
|
||||
assert_eq!(retained[0].detail.as_deref(), Some("new"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
//! Core runtime primitives for the `claw` CLI and supporting crates.
|
||||
//!
|
||||
//! This crate owns session persistence, permission evaluation, prompt assembly,
|
||||
//! MCP plumbing, tool-facing file operations, and the core conversation loop
|
||||
//! that drives interactive and one-shot turns.
|
||||
|
||||
mod bash;
|
||||
pub mod bash_validation;
|
||||
mod bootstrap;
|
||||
pub mod branch_lock;
|
||||
mod compact;
|
||||
mod config;
|
||||
pub mod config_validate;
|
||||
mod conversation;
|
||||
mod file_ops;
|
||||
mod git_context;
|
||||
pub mod green_contract;
|
||||
mod hooks;
|
||||
mod json;
|
||||
mod lane_events;
|
||||
pub mod lsp_client;
|
||||
mod mcp;
|
||||
mod mcp_client;
|
||||
pub mod mcp_lifecycle_hardened;
|
||||
pub mod mcp_server;
|
||||
mod mcp_stdio;
|
||||
pub mod mcp_tool_bridge;
|
||||
mod oauth;
|
||||
pub mod permission_enforcer;
|
||||
mod permissions;
|
||||
pub mod plugin_lifecycle;
|
||||
mod policy_engine;
|
||||
mod prompt;
|
||||
pub mod recovery_recipes;
|
||||
mod remote;
|
||||
pub mod sandbox;
|
||||
mod session;
|
||||
pub mod session_control;
|
||||
pub use session_control::SessionStore;
|
||||
mod sse;
|
||||
pub mod stale_base;
|
||||
pub mod stale_branch;
|
||||
pub mod summary_compression;
|
||||
pub mod task_packet;
|
||||
pub mod task_registry;
|
||||
pub mod team_cron_registry;
|
||||
#[cfg(test)]
|
||||
mod trust_resolver;
|
||||
mod usage;
|
||||
pub mod worker_boot;
|
||||
|
||||
pub use bash::{execute_bash, BashCommandInput, BashCommandOutput};
|
||||
pub use bootstrap::{BootstrapPhase, BootstrapPlan};
|
||||
pub use branch_lock::{detect_branch_lock_collisions, BranchLockCollision, BranchLockIntent};
|
||||
pub use compact::{
|
||||
compact_session, estimate_session_tokens, format_compact_summary,
|
||||
get_compact_continuation_message, should_compact, CompactionConfig, CompactionResult,
|
||||
};
|
||||
pub use config::{
|
||||
ConfigEntry, ConfigError, ConfigLoader, ConfigSource, McpConfigCollection,
|
||||
McpManagedProxyServerConfig, McpOAuthConfig, McpRemoteServerConfig, McpSdkServerConfig,
|
||||
McpServerConfig, McpStdioServerConfig, McpTransport, McpWebSocketServerConfig, OAuthConfig,
|
||||
ProviderFallbackConfig, ResolvedPermissionMode, RuntimeConfig, RuntimeFeatureConfig,
|
||||
RuntimeHookConfig, RuntimePermissionRuleConfig, RuntimePluginConfig, ScopedMcpServerConfig,
|
||||
CLAW_SETTINGS_SCHEMA_NAME,
|
||||
};
|
||||
pub use config_validate::{
|
||||
check_unsupported_format, format_diagnostics, validate_config_file, ConfigDiagnostic,
|
||||
DiagnosticKind, ValidationResult,
|
||||
};
|
||||
pub use conversation::{
|
||||
auto_compaction_threshold_from_env, ApiClient, ApiRequest, AssistantEvent, AutoCompactionEvent,
|
||||
ConversationRuntime, PromptCacheEvent, RuntimeError, StaticToolExecutor, ToolError,
|
||||
ToolExecutor, TurnSummary,
|
||||
};
|
||||
pub use file_ops::{
|
||||
edit_file, glob_search, grep_search, read_file, write_file, EditFileOutput, GlobSearchOutput,
|
||||
GrepSearchInput, GrepSearchOutput, ReadFileOutput, StructuredPatchHunk, TextFilePayload,
|
||||
WriteFileOutput,
|
||||
};
|
||||
pub use git_context::{GitCommitEntry, GitContext};
|
||||
pub use hooks::{
|
||||
HookAbortSignal, HookEvent, HookProgressEvent, HookProgressReporter, HookRunResult, HookRunner,
|
||||
};
|
||||
pub use lane_events::{
|
||||
dedupe_superseded_commit_events, LaneCommitProvenance, LaneEvent, LaneEventBlocker,
|
||||
LaneEventName, LaneEventStatus, LaneFailureClass,
|
||||
};
|
||||
pub use mcp::{
|
||||
mcp_server_signature, mcp_tool_name, mcp_tool_prefix, normalize_name_for_mcp,
|
||||
scoped_mcp_config_hash, unwrap_ccr_proxy_url,
|
||||
};
|
||||
pub use mcp_client::{
|
||||
McpClientAuth, McpClientBootstrap, McpClientTransport, McpManagedProxyTransport,
|
||||
McpRemoteTransport, McpSdkTransport, McpStdioTransport,
|
||||
};
|
||||
pub use mcp_lifecycle_hardened::{
|
||||
McpDegradedReport, McpErrorSurface, McpFailedServer, McpLifecyclePhase, McpLifecycleState,
|
||||
McpLifecycleValidator, McpPhaseResult,
|
||||
};
|
||||
pub use mcp_server::{McpServer, McpServerSpec, ToolCallHandler, MCP_SERVER_PROTOCOL_VERSION};
|
||||
pub use mcp_stdio::{
|
||||
spawn_mcp_stdio_process, JsonRpcError, JsonRpcId, JsonRpcRequest, JsonRpcResponse,
|
||||
ManagedMcpTool, McpDiscoveryFailure, McpInitializeClientInfo, McpInitializeParams,
|
||||
McpInitializeResult, McpInitializeServerInfo, McpListResourcesParams, McpListResourcesResult,
|
||||
McpListToolsParams, McpListToolsResult, McpReadResourceParams, McpReadResourceResult,
|
||||
McpResource, McpResourceContents, McpServerManager, McpServerManagerError, McpStdioProcess,
|
||||
McpTool, McpToolCallContent, McpToolCallParams, McpToolCallResult, McpToolDiscoveryReport,
|
||||
UnsupportedMcpServer,
|
||||
};
|
||||
pub use oauth::{
|
||||
clear_oauth_credentials, code_challenge_s256, credentials_path, generate_pkce_pair,
|
||||
generate_state, load_oauth_credentials, loopback_redirect_uri, parse_oauth_callback_query,
|
||||
parse_oauth_callback_request_target, save_oauth_credentials, OAuthAuthorizationRequest,
|
||||
OAuthCallbackParams, OAuthRefreshRequest, OAuthTokenExchangeRequest, OAuthTokenSet,
|
||||
PkceChallengeMethod, PkceCodePair,
|
||||
};
|
||||
pub use permissions::{
|
||||
PermissionContext, PermissionMode, PermissionOutcome, PermissionOverride, PermissionPolicy,
|
||||
PermissionPromptDecision, PermissionPrompter, PermissionRequest,
|
||||
};
|
||||
pub use plugin_lifecycle::{
|
||||
DegradedMode, DiscoveryResult, PluginHealthcheck, PluginLifecycle, PluginLifecycleEvent,
|
||||
PluginState, ResourceInfo, ServerHealth, ServerStatus, ToolInfo,
|
||||
};
|
||||
pub use policy_engine::{
|
||||
evaluate, DiffScope, GreenLevel, LaneBlocker, LaneContext, PolicyAction, PolicyCondition,
|
||||
PolicyEngine, PolicyRule, ReconcileReason, ReviewStatus,
|
||||
};
|
||||
pub use prompt::{
|
||||
load_system_prompt, prepend_bullets, ContextFile, ProjectContext, PromptBuildError,
|
||||
SystemPromptBuilder, FRONTIER_MODEL_NAME, SYSTEM_PROMPT_DYNAMIC_BOUNDARY,
|
||||
};
|
||||
pub use recovery_recipes::{
|
||||
attempt_recovery, recipe_for, EscalationPolicy, FailureScenario, RecoveryContext,
|
||||
RecoveryEvent, RecoveryRecipe, RecoveryResult, RecoveryStep,
|
||||
};
|
||||
pub use remote::{
|
||||
inherited_upstream_proxy_env, no_proxy_list, read_token, upstream_proxy_ws_url,
|
||||
RemoteSessionContext, UpstreamProxyBootstrap, UpstreamProxyState, DEFAULT_REMOTE_BASE_URL,
|
||||
DEFAULT_SESSION_TOKEN_PATH, DEFAULT_SYSTEM_CA_BUNDLE, NO_PROXY_HOSTS, UPSTREAM_PROXY_ENV_KEYS,
|
||||
};
|
||||
pub use sandbox::{
|
||||
build_linux_sandbox_command, detect_container_environment, detect_container_environment_from,
|
||||
resolve_sandbox_status, resolve_sandbox_status_for_request, ContainerEnvironment,
|
||||
FilesystemIsolationMode, LinuxSandboxCommand, SandboxConfig, SandboxDetectionInputs,
|
||||
SandboxRequest, SandboxStatus,
|
||||
};
|
||||
pub use session::{
|
||||
ContentBlock, ConversationMessage, MessageRole, Session, SessionCompaction, SessionError,
|
||||
SessionFork, SessionPromptEntry,
|
||||
};
|
||||
pub use sse::{IncrementalSseParser, SseEvent};
|
||||
pub use stale_base::{
|
||||
check_base_commit, format_stale_base_warning, read_claw_base_file, resolve_expected_base,
|
||||
BaseCommitSource, BaseCommitState,
|
||||
};
|
||||
pub use stale_branch::{
|
||||
apply_policy, check_freshness, BranchFreshness, StaleBranchAction, StaleBranchEvent,
|
||||
StaleBranchPolicy,
|
||||
};
|
||||
pub use task_packet::{validate_packet, TaskPacket, TaskPacketValidationError, ValidatedPacket};
|
||||
#[cfg(test)]
|
||||
pub use trust_resolver::{TrustConfig, TrustDecision, TrustEvent, TrustPolicy, TrustResolver};
|
||||
pub use usage::{
|
||||
format_usd, pricing_for_model, ModelPricing, TokenUsage, UsageCostEstimate, UsageTracker,
|
||||
};
|
||||
pub use worker_boot::{
|
||||
Worker, WorkerEvent, WorkerEventKind, WorkerEventPayload, WorkerFailure, WorkerFailureKind,
|
||||
WorkerPromptTarget, WorkerReadySnapshot, WorkerRegistry, WorkerStatus, WorkerTrustResolution,
|
||||
};
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn test_env_lock() -> std::sync::MutexGuard<'static, ()> {
|
||||
static LOCK: std::sync::OnceLock<std::sync::Mutex<()>> = std::sync::OnceLock::new();
|
||||
LOCK.get_or_init(|| std::sync::Mutex::new(()))
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner)
|
||||
}
|
||||
@@ -0,0 +1,747 @@
|
||||
#![allow(clippy::should_implement_trait, clippy::must_use_candidate)]
|
||||
//! LSP (Language Server Protocol) client registry for tool dispatch.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Supported LSP actions.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum LspAction {
|
||||
Diagnostics,
|
||||
Hover,
|
||||
Definition,
|
||||
References,
|
||||
Completion,
|
||||
Symbols,
|
||||
Format,
|
||||
}
|
||||
|
||||
impl LspAction {
|
||||
pub fn from_str(s: &str) -> Option<Self> {
|
||||
match s {
|
||||
"diagnostics" => Some(Self::Diagnostics),
|
||||
"hover" => Some(Self::Hover),
|
||||
"definition" | "goto_definition" => Some(Self::Definition),
|
||||
"references" | "find_references" => Some(Self::References),
|
||||
"completion" | "completions" => Some(Self::Completion),
|
||||
"symbols" | "document_symbols" => Some(Self::Symbols),
|
||||
"format" | "formatting" => Some(Self::Format),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LspDiagnostic {
|
||||
pub path: String,
|
||||
pub line: u32,
|
||||
pub character: u32,
|
||||
pub severity: String,
|
||||
pub message: String,
|
||||
pub source: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LspLocation {
|
||||
pub path: String,
|
||||
pub line: u32,
|
||||
pub character: u32,
|
||||
pub end_line: Option<u32>,
|
||||
pub end_character: Option<u32>,
|
||||
pub preview: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LspHoverResult {
|
||||
pub content: String,
|
||||
pub language: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LspCompletionItem {
|
||||
pub label: String,
|
||||
pub kind: Option<String>,
|
||||
pub detail: Option<String>,
|
||||
pub insert_text: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LspSymbol {
|
||||
pub name: String,
|
||||
pub kind: String,
|
||||
pub path: String,
|
||||
pub line: u32,
|
||||
pub character: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum LspServerStatus {
|
||||
Connected,
|
||||
Disconnected,
|
||||
Starting,
|
||||
Error,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for LspServerStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Connected => write!(f, "connected"),
|
||||
Self::Disconnected => write!(f, "disconnected"),
|
||||
Self::Starting => write!(f, "starting"),
|
||||
Self::Error => write!(f, "error"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct LspServerState {
|
||||
pub language: String,
|
||||
pub status: LspServerStatus,
|
||||
pub root_path: Option<String>,
|
||||
pub capabilities: Vec<String>,
|
||||
pub diagnostics: Vec<LspDiagnostic>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct LspRegistry {
|
||||
inner: Arc<Mutex<RegistryInner>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct RegistryInner {
|
||||
servers: HashMap<String, LspServerState>,
|
||||
}
|
||||
|
||||
impl LspRegistry {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn register(
|
||||
&self,
|
||||
language: &str,
|
||||
status: LspServerStatus,
|
||||
root_path: Option<&str>,
|
||||
capabilities: Vec<String>,
|
||||
) {
|
||||
let mut inner = self.inner.lock().expect("lsp registry lock poisoned");
|
||||
inner.servers.insert(
|
||||
language.to_owned(),
|
||||
LspServerState {
|
||||
language: language.to_owned(),
|
||||
status,
|
||||
root_path: root_path.map(str::to_owned),
|
||||
capabilities,
|
||||
diagnostics: Vec::new(),
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub fn get(&self, language: &str) -> Option<LspServerState> {
|
||||
let inner = self.inner.lock().expect("lsp registry lock poisoned");
|
||||
inner.servers.get(language).cloned()
|
||||
}
|
||||
|
||||
/// Find the appropriate server for a file path based on extension.
|
||||
pub fn find_server_for_path(&self, path: &str) -> Option<LspServerState> {
|
||||
let ext = std::path::Path::new(path)
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.unwrap_or("");
|
||||
|
||||
let language = match ext {
|
||||
"rs" => "rust",
|
||||
"ts" | "tsx" => "typescript",
|
||||
"js" | "jsx" => "javascript",
|
||||
"py" => "python",
|
||||
"go" => "go",
|
||||
"java" => "java",
|
||||
"c" | "h" => "c",
|
||||
"cpp" | "hpp" | "cc" => "cpp",
|
||||
"rb" => "ruby",
|
||||
"lua" => "lua",
|
||||
_ => return None,
|
||||
};
|
||||
|
||||
self.get(language)
|
||||
}
|
||||
|
||||
/// List all registered servers.
|
||||
pub fn list_servers(&self) -> Vec<LspServerState> {
|
||||
let inner = self.inner.lock().expect("lsp registry lock poisoned");
|
||||
inner.servers.values().cloned().collect()
|
||||
}
|
||||
|
||||
/// Add diagnostics to a server.
|
||||
pub fn add_diagnostics(
|
||||
&self,
|
||||
language: &str,
|
||||
diagnostics: Vec<LspDiagnostic>,
|
||||
) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("lsp registry lock poisoned");
|
||||
let server = inner
|
||||
.servers
|
||||
.get_mut(language)
|
||||
.ok_or_else(|| format!("LSP server not found for language: {language}"))?;
|
||||
server.diagnostics.extend(diagnostics);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get diagnostics for a specific file path.
|
||||
pub fn get_diagnostics(&self, path: &str) -> Vec<LspDiagnostic> {
|
||||
let inner = self.inner.lock().expect("lsp registry lock poisoned");
|
||||
inner
|
||||
.servers
|
||||
.values()
|
||||
.flat_map(|s| &s.diagnostics)
|
||||
.filter(|d| d.path == path)
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Clear diagnostics for a language server.
|
||||
pub fn clear_diagnostics(&self, language: &str) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("lsp registry lock poisoned");
|
||||
let server = inner
|
||||
.servers
|
||||
.get_mut(language)
|
||||
.ok_or_else(|| format!("LSP server not found for language: {language}"))?;
|
||||
server.diagnostics.clear();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Disconnect a server.
|
||||
pub fn disconnect(&self, language: &str) -> Option<LspServerState> {
|
||||
let mut inner = self.inner.lock().expect("lsp registry lock poisoned");
|
||||
inner.servers.remove(language)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn len(&self) -> usize {
|
||||
let inner = self.inner.lock().expect("lsp registry lock poisoned");
|
||||
inner.servers.len()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
|
||||
/// Dispatch an LSP action and return a structured result.
|
||||
pub fn dispatch(
|
||||
&self,
|
||||
action: &str,
|
||||
path: Option<&str>,
|
||||
line: Option<u32>,
|
||||
character: Option<u32>,
|
||||
_query: Option<&str>,
|
||||
) -> Result<serde_json::Value, String> {
|
||||
let lsp_action =
|
||||
LspAction::from_str(action).ok_or_else(|| format!("unknown LSP action: {action}"))?;
|
||||
|
||||
// For diagnostics, we can check existing cached diagnostics
|
||||
if lsp_action == LspAction::Diagnostics {
|
||||
if let Some(path) = path {
|
||||
let diags = self.get_diagnostics(path);
|
||||
return Ok(serde_json::json!({
|
||||
"action": "diagnostics",
|
||||
"path": path,
|
||||
"diagnostics": diags,
|
||||
"count": diags.len()
|
||||
}));
|
||||
}
|
||||
// All diagnostics across all servers
|
||||
let inner = self.inner.lock().expect("lsp registry lock poisoned");
|
||||
let all_diags: Vec<_> = inner
|
||||
.servers
|
||||
.values()
|
||||
.flat_map(|s| &s.diagnostics)
|
||||
.collect();
|
||||
return Ok(serde_json::json!({
|
||||
"action": "diagnostics",
|
||||
"diagnostics": all_diags,
|
||||
"count": all_diags.len()
|
||||
}));
|
||||
}
|
||||
|
||||
// For other actions, we need a connected server for the given file
|
||||
let path = path.ok_or("path is required for this LSP action")?;
|
||||
let server = self
|
||||
.find_server_for_path(path)
|
||||
.ok_or_else(|| format!("no LSP server available for path: {path}"))?;
|
||||
|
||||
if server.status != LspServerStatus::Connected {
|
||||
return Err(format!(
|
||||
"LSP server for '{}' is not connected (status: {})",
|
||||
server.language, server.status
|
||||
));
|
||||
}
|
||||
|
||||
// Return structured placeholder — actual LSP JSON-RPC calls would
|
||||
// go through the real LSP process here.
|
||||
Ok(serde_json::json!({
|
||||
"action": action,
|
||||
"path": path,
|
||||
"line": line,
|
||||
"character": character,
|
||||
"language": server.language,
|
||||
"status": "dispatched",
|
||||
"message": format!("LSP {} dispatched to {} server", action, server.language)
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn registers_and_retrieves_server() {
|
||||
let registry = LspRegistry::new();
|
||||
registry.register(
|
||||
"rust",
|
||||
LspServerStatus::Connected,
|
||||
Some("/workspace"),
|
||||
vec!["hover".into(), "completion".into()],
|
||||
);
|
||||
|
||||
let server = registry.get("rust").expect("should exist");
|
||||
assert_eq!(server.language, "rust");
|
||||
assert_eq!(server.status, LspServerStatus::Connected);
|
||||
assert_eq!(server.capabilities.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finds_server_by_file_extension() {
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("rust", LspServerStatus::Connected, None, vec![]);
|
||||
registry.register("typescript", LspServerStatus::Connected, None, vec![]);
|
||||
|
||||
let rs_server = registry.find_server_for_path("src/main.rs").unwrap();
|
||||
assert_eq!(rs_server.language, "rust");
|
||||
|
||||
let ts_server = registry.find_server_for_path("src/index.ts").unwrap();
|
||||
assert_eq!(ts_server.language, "typescript");
|
||||
|
||||
assert!(registry.find_server_for_path("data.csv").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manages_diagnostics() {
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("rust", LspServerStatus::Connected, None, vec![]);
|
||||
|
||||
registry
|
||||
.add_diagnostics(
|
||||
"rust",
|
||||
vec![LspDiagnostic {
|
||||
path: "src/main.rs".into(),
|
||||
line: 10,
|
||||
character: 5,
|
||||
severity: "error".into(),
|
||||
message: "mismatched types".into(),
|
||||
source: Some("rust-analyzer".into()),
|
||||
}],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let diags = registry.get_diagnostics("src/main.rs");
|
||||
assert_eq!(diags.len(), 1);
|
||||
assert_eq!(diags[0].message, "mismatched types");
|
||||
|
||||
registry.clear_diagnostics("rust").unwrap();
|
||||
assert!(registry.get_diagnostics("src/main.rs").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatches_diagnostics_action() {
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("rust", LspServerStatus::Connected, None, vec![]);
|
||||
registry
|
||||
.add_diagnostics(
|
||||
"rust",
|
||||
vec![LspDiagnostic {
|
||||
path: "src/lib.rs".into(),
|
||||
line: 1,
|
||||
character: 0,
|
||||
severity: "warning".into(),
|
||||
message: "unused import".into(),
|
||||
source: None,
|
||||
}],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let result = registry
|
||||
.dispatch("diagnostics", Some("src/lib.rs"), None, None, None)
|
||||
.unwrap();
|
||||
assert_eq!(result["count"], 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatches_hover_action() {
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("rust", LspServerStatus::Connected, None, vec![]);
|
||||
|
||||
let result = registry
|
||||
.dispatch("hover", Some("src/main.rs"), Some(10), Some(5), None)
|
||||
.unwrap();
|
||||
assert_eq!(result["action"], "hover");
|
||||
assert_eq!(result["language"], "rust");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_action_on_disconnected_server() {
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("rust", LspServerStatus::Disconnected, None, vec![]);
|
||||
|
||||
assert!(registry
|
||||
.dispatch("hover", Some("src/main.rs"), Some(1), Some(0), None)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_unknown_action() {
|
||||
let registry = LspRegistry::new();
|
||||
assert!(registry
|
||||
.dispatch("unknown_action", Some("file.rs"), None, None, None)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disconnects_server() {
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("rust", LspServerStatus::Connected, None, vec![]);
|
||||
assert_eq!(registry.len(), 1);
|
||||
|
||||
let removed = registry.disconnect("rust");
|
||||
assert!(removed.is_some());
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lsp_action_from_str_all_aliases() {
|
||||
// given
|
||||
let cases = [
|
||||
("diagnostics", Some(LspAction::Diagnostics)),
|
||||
("hover", Some(LspAction::Hover)),
|
||||
("definition", Some(LspAction::Definition)),
|
||||
("goto_definition", Some(LspAction::Definition)),
|
||||
("references", Some(LspAction::References)),
|
||||
("find_references", Some(LspAction::References)),
|
||||
("completion", Some(LspAction::Completion)),
|
||||
("completions", Some(LspAction::Completion)),
|
||||
("symbols", Some(LspAction::Symbols)),
|
||||
("document_symbols", Some(LspAction::Symbols)),
|
||||
("format", Some(LspAction::Format)),
|
||||
("formatting", Some(LspAction::Format)),
|
||||
("unknown", None),
|
||||
];
|
||||
|
||||
// when
|
||||
let resolved: Vec<_> = cases
|
||||
.into_iter()
|
||||
.map(|(input, expected)| (input, LspAction::from_str(input), expected))
|
||||
.collect();
|
||||
|
||||
// then
|
||||
for (input, actual, expected) in resolved {
|
||||
assert_eq!(actual, expected, "unexpected action resolution for {input}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lsp_server_status_display_all_variants() {
|
||||
// given
|
||||
let cases = [
|
||||
(LspServerStatus::Connected, "connected"),
|
||||
(LspServerStatus::Disconnected, "disconnected"),
|
||||
(LspServerStatus::Starting, "starting"),
|
||||
(LspServerStatus::Error, "error"),
|
||||
];
|
||||
|
||||
// when
|
||||
let rendered: Vec<_> = cases
|
||||
.into_iter()
|
||||
.map(|(status, expected)| (status.to_string(), expected))
|
||||
.collect();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
rendered,
|
||||
vec![
|
||||
("connected".to_string(), "connected"),
|
||||
("disconnected".to_string(), "disconnected"),
|
||||
("starting".to_string(), "starting"),
|
||||
("error".to_string(), "error"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatch_diagnostics_without_path_aggregates() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("rust", LspServerStatus::Connected, None, vec![]);
|
||||
registry.register("python", LspServerStatus::Connected, None, vec![]);
|
||||
registry
|
||||
.add_diagnostics(
|
||||
"rust",
|
||||
vec![LspDiagnostic {
|
||||
path: "src/lib.rs".into(),
|
||||
line: 1,
|
||||
character: 0,
|
||||
severity: "warning".into(),
|
||||
message: "unused import".into(),
|
||||
source: Some("rust-analyzer".into()),
|
||||
}],
|
||||
)
|
||||
.expect("rust diagnostics should add");
|
||||
registry
|
||||
.add_diagnostics(
|
||||
"python",
|
||||
vec![LspDiagnostic {
|
||||
path: "script.py".into(),
|
||||
line: 2,
|
||||
character: 4,
|
||||
severity: "error".into(),
|
||||
message: "undefined name".into(),
|
||||
source: Some("pyright".into()),
|
||||
}],
|
||||
)
|
||||
.expect("python diagnostics should add");
|
||||
|
||||
// when
|
||||
let result = registry
|
||||
.dispatch("diagnostics", None, None, None, None)
|
||||
.expect("aggregate diagnostics should work");
|
||||
|
||||
// then
|
||||
assert_eq!(result["action"], "diagnostics");
|
||||
assert_eq!(result["count"], 2);
|
||||
assert_eq!(result["diagnostics"].as_array().map(Vec::len), Some(2));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatch_non_diagnostics_requires_path() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
|
||||
// when
|
||||
let result = registry.dispatch("hover", None, Some(1), Some(0), None);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
result.expect_err("path should be required"),
|
||||
"path is required for this LSP action"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatch_no_server_for_path_errors() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
|
||||
// when
|
||||
let result = registry.dispatch("hover", Some("notes.md"), Some(1), Some(0), None);
|
||||
|
||||
// then
|
||||
let error = result.expect_err("missing server should fail");
|
||||
assert!(error.contains("no LSP server available for path: notes.md"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatch_disconnected_server_error_payload() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("typescript", LspServerStatus::Disconnected, None, vec![]);
|
||||
|
||||
// when
|
||||
let result = registry.dispatch("hover", Some("src/index.ts"), Some(3), Some(2), None);
|
||||
|
||||
// then
|
||||
let error = result.expect_err("disconnected server should fail");
|
||||
assert!(error.contains("typescript"));
|
||||
assert!(error.contains("disconnected"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn find_server_for_all_extensions() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
for language in [
|
||||
"rust",
|
||||
"typescript",
|
||||
"javascript",
|
||||
"python",
|
||||
"go",
|
||||
"java",
|
||||
"c",
|
||||
"cpp",
|
||||
"ruby",
|
||||
"lua",
|
||||
] {
|
||||
registry.register(language, LspServerStatus::Connected, None, vec![]);
|
||||
}
|
||||
let cases = [
|
||||
("src/main.rs", "rust"),
|
||||
("src/index.ts", "typescript"),
|
||||
("src/view.tsx", "typescript"),
|
||||
("src/app.js", "javascript"),
|
||||
("src/app.jsx", "javascript"),
|
||||
("script.py", "python"),
|
||||
("main.go", "go"),
|
||||
("Main.java", "java"),
|
||||
("native.c", "c"),
|
||||
("native.h", "c"),
|
||||
("native.cpp", "cpp"),
|
||||
("native.hpp", "cpp"),
|
||||
("native.cc", "cpp"),
|
||||
("script.rb", "ruby"),
|
||||
("script.lua", "lua"),
|
||||
];
|
||||
|
||||
// when
|
||||
let resolved: Vec<_> = cases
|
||||
.into_iter()
|
||||
.map(|(path, expected)| {
|
||||
(
|
||||
path,
|
||||
registry
|
||||
.find_server_for_path(path)
|
||||
.map(|server| server.language),
|
||||
expected,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
// then
|
||||
for (path, actual, expected) in resolved {
|
||||
assert_eq!(
|
||||
actual.as_deref(),
|
||||
Some(expected),
|
||||
"unexpected mapping for {path}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn find_server_for_path_no_extension() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("rust", LspServerStatus::Connected, None, vec![]);
|
||||
|
||||
// when
|
||||
let result = registry.find_server_for_path("Makefile");
|
||||
|
||||
// then
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_servers_with_multiple() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
registry.register("rust", LspServerStatus::Connected, None, vec![]);
|
||||
registry.register("typescript", LspServerStatus::Starting, None, vec![]);
|
||||
registry.register("python", LspServerStatus::Error, None, vec![]);
|
||||
|
||||
// when
|
||||
let servers = registry.list_servers();
|
||||
|
||||
// then
|
||||
assert_eq!(servers.len(), 3);
|
||||
assert!(servers.iter().any(|server| server.language == "rust"));
|
||||
assert!(servers.iter().any(|server| server.language == "typescript"));
|
||||
assert!(servers.iter().any(|server| server.language == "python"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_missing_server_returns_none() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
|
||||
// when
|
||||
let server = registry.get("missing");
|
||||
|
||||
// then
|
||||
assert!(server.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn add_diagnostics_missing_language_errors() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
|
||||
// when
|
||||
let result = registry.add_diagnostics("missing", vec![]);
|
||||
|
||||
// then
|
||||
let error = result.expect_err("missing language should fail");
|
||||
assert!(error.contains("LSP server not found for language: missing"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_diagnostics_across_servers() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
let shared_path = "shared/file.txt";
|
||||
registry.register("rust", LspServerStatus::Connected, None, vec![]);
|
||||
registry.register("python", LspServerStatus::Connected, None, vec![]);
|
||||
registry
|
||||
.add_diagnostics(
|
||||
"rust",
|
||||
vec![LspDiagnostic {
|
||||
path: shared_path.into(),
|
||||
line: 4,
|
||||
character: 1,
|
||||
severity: "warning".into(),
|
||||
message: "warn".into(),
|
||||
source: None,
|
||||
}],
|
||||
)
|
||||
.expect("rust diagnostics should add");
|
||||
registry
|
||||
.add_diagnostics(
|
||||
"python",
|
||||
vec![LspDiagnostic {
|
||||
path: shared_path.into(),
|
||||
line: 8,
|
||||
character: 3,
|
||||
severity: "error".into(),
|
||||
message: "err".into(),
|
||||
source: None,
|
||||
}],
|
||||
)
|
||||
.expect("python diagnostics should add");
|
||||
|
||||
// when
|
||||
let diagnostics = registry.get_diagnostics(shared_path);
|
||||
|
||||
// then
|
||||
assert_eq!(diagnostics.len(), 2);
|
||||
assert!(diagnostics
|
||||
.iter()
|
||||
.any(|diagnostic| diagnostic.message == "warn"));
|
||||
assert!(diagnostics
|
||||
.iter()
|
||||
.any(|diagnostic| diagnostic.message == "err"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn clear_diagnostics_missing_language_errors() {
|
||||
// given
|
||||
let registry = LspRegistry::new();
|
||||
|
||||
// when
|
||||
let result = registry.clear_diagnostics("missing");
|
||||
|
||||
// then
|
||||
let error = result.expect_err("missing language should fail");
|
||||
assert!(error.contains("LSP server not found for language: missing"));
|
||||
}
|
||||
}
|
||||
@@ -73,7 +73,7 @@ pub fn mcp_server_signature(config: &McpServerConfig) -> Option<String> {
|
||||
Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url)))
|
||||
}
|
||||
McpServerConfig::Ws(config) => Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url))),
|
||||
McpServerConfig::ClaudeAiProxy(config) => {
|
||||
McpServerConfig::ManagedProxy(config) => {
|
||||
Some(format!("url:{}", unwrap_ccr_proxy_url(&config.url)))
|
||||
}
|
||||
McpServerConfig::Sdk(_) => None,
|
||||
@@ -84,10 +84,13 @@ pub fn mcp_server_signature(config: &McpServerConfig) -> Option<String> {
|
||||
pub fn scoped_mcp_config_hash(config: &ScopedMcpServerConfig) -> String {
|
||||
let rendered = match &config.config {
|
||||
McpServerConfig::Stdio(stdio) => format!(
|
||||
"stdio|{}|{}|{}",
|
||||
"stdio|{}|{}|{}|{}",
|
||||
stdio.command,
|
||||
render_command_signature(&stdio.args),
|
||||
render_env_signature(&stdio.env)
|
||||
render_env_signature(&stdio.env),
|
||||
stdio
|
||||
.tool_call_timeout_ms
|
||||
.map_or_else(String::new, |timeout_ms| timeout_ms.to_string())
|
||||
),
|
||||
McpServerConfig::Sse(remote) => format!(
|
||||
"sse|{}|{}|{}|{}",
|
||||
@@ -110,7 +113,7 @@ pub fn scoped_mcp_config_hash(config: &ScopedMcpServerConfig) -> String {
|
||||
ws.headers_helper.as_deref().unwrap_or("")
|
||||
),
|
||||
McpServerConfig::Sdk(sdk) => format!("sdk|{}", sdk.name),
|
||||
McpServerConfig::ClaudeAiProxy(proxy) => {
|
||||
McpServerConfig::ManagedProxy(proxy) => {
|
||||
format!("claudeai-proxy|{}|{}", proxy.url, proxy.id)
|
||||
}
|
||||
};
|
||||
@@ -245,6 +248,7 @@ mod tests {
|
||||
command: "uvx".to_string(),
|
||||
args: vec!["mcp-server".to_string()],
|
||||
env: BTreeMap::from([("TOKEN".to_string(), "secret".to_string())]),
|
||||
tool_call_timeout_ms: None,
|
||||
});
|
||||
assert_eq!(
|
||||
mcp_server_signature(&stdio),
|
||||
@@ -3,6 +3,8 @@ use std::collections::BTreeMap;
|
||||
use crate::config::{McpOAuthConfig, McpServerConfig, ScopedMcpServerConfig};
|
||||
use crate::mcp::{mcp_server_signature, mcp_tool_prefix, normalize_name_for_mcp};
|
||||
|
||||
pub const DEFAULT_MCP_TOOL_CALL_TIMEOUT_MS: u64 = 60_000;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum McpClientTransport {
|
||||
Stdio(McpStdioTransport),
|
||||
@@ -10,7 +12,7 @@ pub enum McpClientTransport {
|
||||
Http(McpRemoteTransport),
|
||||
WebSocket(McpRemoteTransport),
|
||||
Sdk(McpSdkTransport),
|
||||
ClaudeAiProxy(McpClaudeAiProxyTransport),
|
||||
ManagedProxy(McpManagedProxyTransport),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -18,6 +20,7 @@ pub struct McpStdioTransport {
|
||||
pub command: String,
|
||||
pub args: Vec<String>,
|
||||
pub env: BTreeMap<String, String>,
|
||||
pub tool_call_timeout_ms: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -34,7 +37,7 @@ pub struct McpSdkTransport {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct McpClaudeAiProxyTransport {
|
||||
pub struct McpManagedProxyTransport {
|
||||
pub url: String,
|
||||
pub id: String,
|
||||
}
|
||||
@@ -75,6 +78,7 @@ impl McpClientTransport {
|
||||
command: config.command.clone(),
|
||||
args: config.args.clone(),
|
||||
env: config.env.clone(),
|
||||
tool_call_timeout_ms: config.tool_call_timeout_ms,
|
||||
}),
|
||||
McpServerConfig::Sse(config) => Self::Sse(McpRemoteTransport {
|
||||
url: config.url.clone(),
|
||||
@@ -97,14 +101,20 @@ impl McpClientTransport {
|
||||
McpServerConfig::Sdk(config) => Self::Sdk(McpSdkTransport {
|
||||
name: config.name.clone(),
|
||||
}),
|
||||
McpServerConfig::ClaudeAiProxy(config) => {
|
||||
Self::ClaudeAiProxy(McpClaudeAiProxyTransport {
|
||||
McpServerConfig::ManagedProxy(config) => Self::ManagedProxy(McpManagedProxyTransport {
|
||||
url: config.url.clone(),
|
||||
id: config.id.clone(),
|
||||
})
|
||||
}),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl McpStdioTransport {
|
||||
#[must_use]
|
||||
pub fn resolved_tool_call_timeout_ms(&self) -> u64 {
|
||||
self.tool_call_timeout_ms
|
||||
.unwrap_or(DEFAULT_MCP_TOOL_CALL_TIMEOUT_MS)
|
||||
}
|
||||
}
|
||||
|
||||
impl McpClientAuth {
|
||||
@@ -138,6 +148,7 @@ mod tests {
|
||||
command: "uvx".to_string(),
|
||||
args: vec!["mcp-server".to_string()],
|
||||
env: BTreeMap::from([("TOKEN".to_string(), "secret".to_string())]),
|
||||
tool_call_timeout_ms: Some(15_000),
|
||||
}),
|
||||
};
|
||||
|
||||
@@ -156,6 +167,7 @@ mod tests {
|
||||
transport.env.get("TOKEN").map(String::as_str),
|
||||
Some("secret")
|
||||
);
|
||||
assert_eq!(transport.tool_call_timeout_ms, Some(15_000));
|
||||
}
|
||||
other => panic!("expected stdio transport, got {other:?}"),
|
||||
}
|
||||
@@ -0,0 +1,843 @@
|
||||
#![allow(clippy::unnested_or_patterns, clippy::map_unwrap_or)]
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum McpLifecyclePhase {
|
||||
ConfigLoad,
|
||||
ServerRegistration,
|
||||
SpawnConnect,
|
||||
InitializeHandshake,
|
||||
ToolDiscovery,
|
||||
ResourceDiscovery,
|
||||
Ready,
|
||||
Invocation,
|
||||
ErrorSurfacing,
|
||||
Shutdown,
|
||||
Cleanup,
|
||||
}
|
||||
|
||||
impl McpLifecyclePhase {
|
||||
#[must_use]
|
||||
pub fn all() -> [Self; 11] {
|
||||
[
|
||||
Self::ConfigLoad,
|
||||
Self::ServerRegistration,
|
||||
Self::SpawnConnect,
|
||||
Self::InitializeHandshake,
|
||||
Self::ToolDiscovery,
|
||||
Self::ResourceDiscovery,
|
||||
Self::Ready,
|
||||
Self::Invocation,
|
||||
Self::ErrorSurfacing,
|
||||
Self::Shutdown,
|
||||
Self::Cleanup,
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for McpLifecyclePhase {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::ConfigLoad => write!(f, "config_load"),
|
||||
Self::ServerRegistration => write!(f, "server_registration"),
|
||||
Self::SpawnConnect => write!(f, "spawn_connect"),
|
||||
Self::InitializeHandshake => write!(f, "initialize_handshake"),
|
||||
Self::ToolDiscovery => write!(f, "tool_discovery"),
|
||||
Self::ResourceDiscovery => write!(f, "resource_discovery"),
|
||||
Self::Ready => write!(f, "ready"),
|
||||
Self::Invocation => write!(f, "invocation"),
|
||||
Self::ErrorSurfacing => write!(f, "error_surfacing"),
|
||||
Self::Shutdown => write!(f, "shutdown"),
|
||||
Self::Cleanup => write!(f, "cleanup"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct McpErrorSurface {
|
||||
pub phase: McpLifecyclePhase,
|
||||
pub server_name: Option<String>,
|
||||
pub message: String,
|
||||
pub context: BTreeMap<String, String>,
|
||||
pub recoverable: bool,
|
||||
pub timestamp: u64,
|
||||
}
|
||||
|
||||
impl McpErrorSurface {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
phase: McpLifecyclePhase,
|
||||
server_name: Option<String>,
|
||||
message: impl Into<String>,
|
||||
context: BTreeMap<String, String>,
|
||||
recoverable: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
phase,
|
||||
server_name,
|
||||
message: message.into(),
|
||||
context,
|
||||
recoverable,
|
||||
timestamp: now_secs(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for McpErrorSurface {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
write!(
|
||||
f,
|
||||
"MCP lifecycle error during {}: {}",
|
||||
self.phase, self.message
|
||||
)?;
|
||||
if let Some(server_name) = &self.server_name {
|
||||
write!(f, " (server: {server_name})")?;
|
||||
}
|
||||
if !self.context.is_empty() {
|
||||
write!(f, " with context {:?}", self.context)?;
|
||||
}
|
||||
if self.recoverable {
|
||||
write!(f, " [recoverable]")?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for McpErrorSurface {}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum McpPhaseResult {
|
||||
Success {
|
||||
phase: McpLifecyclePhase,
|
||||
duration: Duration,
|
||||
},
|
||||
Failure {
|
||||
phase: McpLifecyclePhase,
|
||||
error: McpErrorSurface,
|
||||
},
|
||||
Timeout {
|
||||
phase: McpLifecyclePhase,
|
||||
waited: Duration,
|
||||
error: McpErrorSurface,
|
||||
},
|
||||
}
|
||||
|
||||
impl McpPhaseResult {
|
||||
#[must_use]
|
||||
pub fn phase(&self) -> McpLifecyclePhase {
|
||||
match self {
|
||||
Self::Success { phase, .. }
|
||||
| Self::Failure { phase, .. }
|
||||
| Self::Timeout { phase, .. } => *phase,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct McpLifecycleState {
|
||||
current_phase: Option<McpLifecyclePhase>,
|
||||
phase_errors: BTreeMap<McpLifecyclePhase, Vec<McpErrorSurface>>,
|
||||
phase_timestamps: BTreeMap<McpLifecyclePhase, u64>,
|
||||
phase_results: Vec<McpPhaseResult>,
|
||||
}
|
||||
|
||||
impl McpLifecycleState {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn current_phase(&self) -> Option<McpLifecyclePhase> {
|
||||
self.current_phase
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn errors_for_phase(&self, phase: McpLifecyclePhase) -> &[McpErrorSurface] {
|
||||
self.phase_errors
|
||||
.get(&phase)
|
||||
.map(Vec::as_slice)
|
||||
.unwrap_or(&[])
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn results(&self) -> &[McpPhaseResult] {
|
||||
&self.phase_results
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn phase_timestamps(&self) -> &BTreeMap<McpLifecyclePhase, u64> {
|
||||
&self.phase_timestamps
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn phase_timestamp(&self, phase: McpLifecyclePhase) -> Option<u64> {
|
||||
self.phase_timestamps.get(&phase).copied()
|
||||
}
|
||||
|
||||
fn record_phase(&mut self, phase: McpLifecyclePhase) {
|
||||
self.current_phase = Some(phase);
|
||||
self.phase_timestamps.insert(phase, now_secs());
|
||||
}
|
||||
|
||||
fn record_error(&mut self, error: McpErrorSurface) {
|
||||
self.phase_errors
|
||||
.entry(error.phase)
|
||||
.or_default()
|
||||
.push(error);
|
||||
}
|
||||
|
||||
fn record_result(&mut self, result: McpPhaseResult) {
|
||||
self.phase_results.push(result);
|
||||
}
|
||||
|
||||
fn can_resume_after_error(&self) -> bool {
|
||||
match self.phase_results.last() {
|
||||
Some(McpPhaseResult::Failure { error, .. } | McpPhaseResult::Timeout { error, .. }) => {
|
||||
error.recoverable
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct McpFailedServer {
|
||||
pub server_name: String,
|
||||
pub phase: McpLifecyclePhase,
|
||||
pub error: McpErrorSurface,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct McpDegradedReport {
|
||||
pub working_servers: Vec<String>,
|
||||
pub failed_servers: Vec<McpFailedServer>,
|
||||
pub available_tools: Vec<String>,
|
||||
pub missing_tools: Vec<String>,
|
||||
}
|
||||
|
||||
impl McpDegradedReport {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
working_servers: Vec<String>,
|
||||
failed_servers: Vec<McpFailedServer>,
|
||||
available_tools: Vec<String>,
|
||||
expected_tools: Vec<String>,
|
||||
) -> Self {
|
||||
let working_servers = dedupe_sorted(working_servers);
|
||||
let available_tools = dedupe_sorted(available_tools);
|
||||
let available_tool_set: BTreeSet<_> = available_tools.iter().cloned().collect();
|
||||
let expected_tools = dedupe_sorted(expected_tools);
|
||||
let missing_tools = expected_tools
|
||||
.into_iter()
|
||||
.filter(|tool| !available_tool_set.contains(tool))
|
||||
.collect();
|
||||
|
||||
Self {
|
||||
working_servers,
|
||||
failed_servers,
|
||||
available_tools,
|
||||
missing_tools,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct McpLifecycleValidator {
|
||||
state: McpLifecycleState,
|
||||
}
|
||||
|
||||
impl McpLifecycleValidator {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn state(&self) -> &McpLifecycleState {
|
||||
&self.state
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn validate_phase_transition(from: McpLifecyclePhase, to: McpLifecyclePhase) -> bool {
|
||||
match (from, to) {
|
||||
(McpLifecyclePhase::ConfigLoad, McpLifecyclePhase::ServerRegistration)
|
||||
| (McpLifecyclePhase::ServerRegistration, McpLifecyclePhase::SpawnConnect)
|
||||
| (McpLifecyclePhase::SpawnConnect, McpLifecyclePhase::InitializeHandshake)
|
||||
| (McpLifecyclePhase::InitializeHandshake, McpLifecyclePhase::ToolDiscovery)
|
||||
| (McpLifecyclePhase::ToolDiscovery, McpLifecyclePhase::ResourceDiscovery)
|
||||
| (McpLifecyclePhase::ToolDiscovery, McpLifecyclePhase::Ready)
|
||||
| (McpLifecyclePhase::ResourceDiscovery, McpLifecyclePhase::Ready)
|
||||
| (McpLifecyclePhase::Ready, McpLifecyclePhase::Invocation)
|
||||
| (McpLifecyclePhase::Invocation, McpLifecyclePhase::Ready)
|
||||
| (McpLifecyclePhase::ErrorSurfacing, McpLifecyclePhase::Ready)
|
||||
| (McpLifecyclePhase::ErrorSurfacing, McpLifecyclePhase::Shutdown)
|
||||
| (McpLifecyclePhase::Shutdown, McpLifecyclePhase::Cleanup) => true,
|
||||
(_, McpLifecyclePhase::Shutdown) => from != McpLifecyclePhase::Cleanup,
|
||||
(_, McpLifecyclePhase::ErrorSurfacing) => {
|
||||
from != McpLifecyclePhase::Cleanup && from != McpLifecyclePhase::Shutdown
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn run_phase(&mut self, phase: McpLifecyclePhase) -> McpPhaseResult {
|
||||
let started = Instant::now();
|
||||
|
||||
if let Some(current_phase) = self.state.current_phase() {
|
||||
if current_phase == McpLifecyclePhase::ErrorSurfacing
|
||||
&& phase == McpLifecyclePhase::Ready
|
||||
&& !self.state.can_resume_after_error()
|
||||
{
|
||||
return self.record_failure(McpErrorSurface::new(
|
||||
phase,
|
||||
None,
|
||||
"cannot return to ready after a non-recoverable MCP lifecycle failure",
|
||||
BTreeMap::from([
|
||||
("from".to_string(), current_phase.to_string()),
|
||||
("to".to_string(), phase.to_string()),
|
||||
]),
|
||||
false,
|
||||
));
|
||||
}
|
||||
|
||||
if !Self::validate_phase_transition(current_phase, phase) {
|
||||
return self.record_failure(McpErrorSurface::new(
|
||||
phase,
|
||||
None,
|
||||
format!("invalid MCP lifecycle transition from {current_phase} to {phase}"),
|
||||
BTreeMap::from([
|
||||
("from".to_string(), current_phase.to_string()),
|
||||
("to".to_string(), phase.to_string()),
|
||||
]),
|
||||
false,
|
||||
));
|
||||
}
|
||||
} else if phase != McpLifecyclePhase::ConfigLoad {
|
||||
return self.record_failure(McpErrorSurface::new(
|
||||
phase,
|
||||
None,
|
||||
format!("invalid initial MCP lifecycle phase {phase}"),
|
||||
BTreeMap::from([("phase".to_string(), phase.to_string())]),
|
||||
false,
|
||||
));
|
||||
}
|
||||
|
||||
self.state.record_phase(phase);
|
||||
let result = McpPhaseResult::Success {
|
||||
phase,
|
||||
duration: started.elapsed(),
|
||||
};
|
||||
self.state.record_result(result.clone());
|
||||
result
|
||||
}
|
||||
|
||||
pub fn record_failure(&mut self, error: McpErrorSurface) -> McpPhaseResult {
|
||||
let phase = error.phase;
|
||||
self.state.record_error(error.clone());
|
||||
self.state.record_phase(McpLifecyclePhase::ErrorSurfacing);
|
||||
let result = McpPhaseResult::Failure { phase, error };
|
||||
self.state.record_result(result.clone());
|
||||
result
|
||||
}
|
||||
|
||||
pub fn record_timeout(
|
||||
&mut self,
|
||||
phase: McpLifecyclePhase,
|
||||
waited: Duration,
|
||||
server_name: Option<String>,
|
||||
mut context: BTreeMap<String, String>,
|
||||
) -> McpPhaseResult {
|
||||
context.insert("waited_ms".to_string(), waited.as_millis().to_string());
|
||||
let error = McpErrorSurface::new(
|
||||
phase,
|
||||
server_name,
|
||||
format!(
|
||||
"MCP lifecycle phase {phase} timed out after {} ms",
|
||||
waited.as_millis()
|
||||
),
|
||||
context,
|
||||
true,
|
||||
);
|
||||
self.state.record_error(error.clone());
|
||||
self.state.record_phase(McpLifecyclePhase::ErrorSurfacing);
|
||||
let result = McpPhaseResult::Timeout {
|
||||
phase,
|
||||
waited,
|
||||
error,
|
||||
};
|
||||
self.state.record_result(result.clone());
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
fn dedupe_sorted(mut values: Vec<String>) -> Vec<String> {
|
||||
values.sort();
|
||||
values.dedup();
|
||||
values
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn phase_display_matches_serde_name() {
|
||||
// given
|
||||
let phases = McpLifecyclePhase::all();
|
||||
|
||||
// when
|
||||
let serialized = phases
|
||||
.into_iter()
|
||||
.map(|phase| {
|
||||
(
|
||||
phase.to_string(),
|
||||
serde_json::to_value(phase).expect("serialize phase"),
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
// then
|
||||
for (display, json_value) in serialized {
|
||||
assert_eq!(json_value, json!(display));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_startup_path_when_running_to_cleanup_then_each_control_transition_succeeds() {
|
||||
// given
|
||||
let mut validator = McpLifecycleValidator::new();
|
||||
let phases = [
|
||||
McpLifecyclePhase::ConfigLoad,
|
||||
McpLifecyclePhase::ServerRegistration,
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
McpLifecyclePhase::InitializeHandshake,
|
||||
McpLifecyclePhase::ToolDiscovery,
|
||||
McpLifecyclePhase::ResourceDiscovery,
|
||||
McpLifecyclePhase::Ready,
|
||||
McpLifecyclePhase::Invocation,
|
||||
McpLifecyclePhase::Ready,
|
||||
McpLifecyclePhase::Shutdown,
|
||||
McpLifecyclePhase::Cleanup,
|
||||
];
|
||||
|
||||
// when
|
||||
let results = phases
|
||||
.into_iter()
|
||||
.map(|phase| validator.run_phase(phase))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
// then
|
||||
assert!(results
|
||||
.iter()
|
||||
.all(|result| matches!(result, McpPhaseResult::Success { .. })));
|
||||
assert_eq!(
|
||||
validator.state().current_phase(),
|
||||
Some(McpLifecyclePhase::Cleanup)
|
||||
);
|
||||
for phase in [
|
||||
McpLifecyclePhase::ConfigLoad,
|
||||
McpLifecyclePhase::ServerRegistration,
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
McpLifecyclePhase::InitializeHandshake,
|
||||
McpLifecyclePhase::ToolDiscovery,
|
||||
McpLifecyclePhase::ResourceDiscovery,
|
||||
McpLifecyclePhase::Ready,
|
||||
McpLifecyclePhase::Invocation,
|
||||
McpLifecyclePhase::Shutdown,
|
||||
McpLifecyclePhase::Cleanup,
|
||||
] {
|
||||
assert!(validator.state().phase_timestamp(phase).is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_tool_discovery_when_resource_discovery_is_skipped_then_ready_is_still_allowed() {
|
||||
// given
|
||||
let mut validator = McpLifecycleValidator::new();
|
||||
for phase in [
|
||||
McpLifecyclePhase::ConfigLoad,
|
||||
McpLifecyclePhase::ServerRegistration,
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
McpLifecyclePhase::InitializeHandshake,
|
||||
McpLifecyclePhase::ToolDiscovery,
|
||||
] {
|
||||
let result = validator.run_phase(phase);
|
||||
assert!(matches!(result, McpPhaseResult::Success { .. }));
|
||||
}
|
||||
|
||||
// when
|
||||
let result = validator.run_phase(McpLifecyclePhase::Ready);
|
||||
|
||||
// then
|
||||
assert!(matches!(result, McpPhaseResult::Success { .. }));
|
||||
assert_eq!(
|
||||
validator.state().current_phase(),
|
||||
Some(McpLifecyclePhase::Ready)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validates_expected_phase_transitions() {
|
||||
// given
|
||||
let valid_transitions = [
|
||||
(
|
||||
McpLifecyclePhase::ConfigLoad,
|
||||
McpLifecyclePhase::ServerRegistration,
|
||||
),
|
||||
(
|
||||
McpLifecyclePhase::ServerRegistration,
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
),
|
||||
(
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
McpLifecyclePhase::InitializeHandshake,
|
||||
),
|
||||
(
|
||||
McpLifecyclePhase::InitializeHandshake,
|
||||
McpLifecyclePhase::ToolDiscovery,
|
||||
),
|
||||
(
|
||||
McpLifecyclePhase::ToolDiscovery,
|
||||
McpLifecyclePhase::ResourceDiscovery,
|
||||
),
|
||||
(McpLifecyclePhase::ToolDiscovery, McpLifecyclePhase::Ready),
|
||||
(
|
||||
McpLifecyclePhase::ResourceDiscovery,
|
||||
McpLifecyclePhase::Ready,
|
||||
),
|
||||
(McpLifecyclePhase::Ready, McpLifecyclePhase::Invocation),
|
||||
(McpLifecyclePhase::Invocation, McpLifecyclePhase::Ready),
|
||||
(McpLifecyclePhase::Ready, McpLifecyclePhase::Shutdown),
|
||||
(
|
||||
McpLifecyclePhase::Invocation,
|
||||
McpLifecyclePhase::ErrorSurfacing,
|
||||
),
|
||||
(
|
||||
McpLifecyclePhase::ErrorSurfacing,
|
||||
McpLifecyclePhase::Shutdown,
|
||||
),
|
||||
(McpLifecyclePhase::Shutdown, McpLifecyclePhase::Cleanup),
|
||||
];
|
||||
|
||||
// when / then
|
||||
for (from, to) in valid_transitions {
|
||||
assert!(McpLifecycleValidator::validate_phase_transition(from, to));
|
||||
}
|
||||
assert!(!McpLifecycleValidator::validate_phase_transition(
|
||||
McpLifecyclePhase::Ready,
|
||||
McpLifecyclePhase::ConfigLoad,
|
||||
));
|
||||
assert!(!McpLifecycleValidator::validate_phase_transition(
|
||||
McpLifecyclePhase::Cleanup,
|
||||
McpLifecyclePhase::Ready,
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_invalid_transition_when_running_phase_then_structured_failure_is_recorded() {
|
||||
// given
|
||||
let mut validator = McpLifecycleValidator::new();
|
||||
let _ = validator.run_phase(McpLifecyclePhase::ConfigLoad);
|
||||
let _ = validator.run_phase(McpLifecyclePhase::ServerRegistration);
|
||||
|
||||
// when
|
||||
let result = validator.run_phase(McpLifecyclePhase::Ready);
|
||||
|
||||
// then
|
||||
match result {
|
||||
McpPhaseResult::Failure { phase, error } => {
|
||||
assert_eq!(phase, McpLifecyclePhase::Ready);
|
||||
assert!(!error.recoverable);
|
||||
assert_eq!(error.phase, McpLifecyclePhase::Ready);
|
||||
assert_eq!(
|
||||
error.context.get("from").map(String::as_str),
|
||||
Some("server_registration")
|
||||
);
|
||||
assert_eq!(error.context.get("to").map(String::as_str), Some("ready"));
|
||||
}
|
||||
other => panic!("expected failure result, got {other:?}"),
|
||||
}
|
||||
assert_eq!(
|
||||
validator.state().current_phase(),
|
||||
Some(McpLifecyclePhase::ErrorSurfacing)
|
||||
);
|
||||
assert_eq!(
|
||||
validator
|
||||
.state()
|
||||
.errors_for_phase(McpLifecyclePhase::Ready)
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_each_phase_when_failure_is_recorded_then_error_is_tracked_per_phase() {
|
||||
// given
|
||||
let mut validator = McpLifecycleValidator::new();
|
||||
|
||||
// when / then
|
||||
for phase in McpLifecyclePhase::all() {
|
||||
let result = validator.record_failure(McpErrorSurface::new(
|
||||
phase,
|
||||
Some("alpha".to_string()),
|
||||
format!("failure at {phase}"),
|
||||
BTreeMap::from([("server".to_string(), "alpha".to_string())]),
|
||||
phase == McpLifecyclePhase::ResourceDiscovery,
|
||||
));
|
||||
|
||||
match result {
|
||||
McpPhaseResult::Failure {
|
||||
phase: failed_phase,
|
||||
error,
|
||||
} => {
|
||||
assert_eq!(failed_phase, phase);
|
||||
assert_eq!(error.phase, phase);
|
||||
assert_eq!(
|
||||
error.recoverable,
|
||||
phase == McpLifecyclePhase::ResourceDiscovery
|
||||
);
|
||||
}
|
||||
other => panic!("expected failure result, got {other:?}"),
|
||||
}
|
||||
assert_eq!(validator.state().errors_for_phase(phase).len(), 1);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_spawn_connect_timeout_when_recorded_then_waited_duration_is_preserved() {
|
||||
// given
|
||||
let mut validator = McpLifecycleValidator::new();
|
||||
let waited = Duration::from_millis(250);
|
||||
|
||||
// when
|
||||
let result = validator.record_timeout(
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
waited,
|
||||
Some("alpha".to_string()),
|
||||
BTreeMap::from([("attempt".to_string(), "1".to_string())]),
|
||||
);
|
||||
|
||||
// then
|
||||
match result {
|
||||
McpPhaseResult::Timeout {
|
||||
phase,
|
||||
waited: actual,
|
||||
error,
|
||||
} => {
|
||||
assert_eq!(phase, McpLifecyclePhase::SpawnConnect);
|
||||
assert_eq!(actual, waited);
|
||||
assert!(error.recoverable);
|
||||
assert_eq!(error.server_name.as_deref(), Some("alpha"));
|
||||
}
|
||||
other => panic!("expected timeout result, got {other:?}"),
|
||||
}
|
||||
let errors = validator
|
||||
.state()
|
||||
.errors_for_phase(McpLifecyclePhase::SpawnConnect);
|
||||
assert_eq!(errors.len(), 1);
|
||||
assert_eq!(
|
||||
errors[0].context.get("waited_ms").map(String::as_str),
|
||||
Some("250")
|
||||
);
|
||||
assert_eq!(
|
||||
validator.state().current_phase(),
|
||||
Some(McpLifecyclePhase::ErrorSurfacing)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_partial_server_health_when_building_degraded_report_then_missing_tools_are_reported() {
|
||||
// given
|
||||
let failed = vec![McpFailedServer {
|
||||
server_name: "broken".to_string(),
|
||||
phase: McpLifecyclePhase::InitializeHandshake,
|
||||
error: McpErrorSurface::new(
|
||||
McpLifecyclePhase::InitializeHandshake,
|
||||
Some("broken".to_string()),
|
||||
"initialize failed",
|
||||
BTreeMap::from([("reason".to_string(), "broken pipe".to_string())]),
|
||||
false,
|
||||
),
|
||||
}];
|
||||
|
||||
// when
|
||||
let report = McpDegradedReport::new(
|
||||
vec!["alpha".to_string(), "beta".to_string(), "alpha".to_string()],
|
||||
failed,
|
||||
vec![
|
||||
"alpha.echo".to_string(),
|
||||
"beta.search".to_string(),
|
||||
"alpha.echo".to_string(),
|
||||
],
|
||||
vec![
|
||||
"alpha.echo".to_string(),
|
||||
"beta.search".to_string(),
|
||||
"broken.fetch".to_string(),
|
||||
],
|
||||
);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
report.working_servers,
|
||||
vec!["alpha".to_string(), "beta".to_string()]
|
||||
);
|
||||
assert_eq!(report.failed_servers.len(), 1);
|
||||
assert_eq!(report.failed_servers[0].server_name, "broken");
|
||||
assert_eq!(
|
||||
report.available_tools,
|
||||
vec!["alpha.echo".to_string(), "beta.search".to_string()]
|
||||
);
|
||||
assert_eq!(report.missing_tools, vec!["broken.fetch".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_failure_during_resource_discovery_when_shutting_down_then_cleanup_still_succeeds() {
|
||||
// given
|
||||
let mut validator = McpLifecycleValidator::new();
|
||||
for phase in [
|
||||
McpLifecyclePhase::ConfigLoad,
|
||||
McpLifecyclePhase::ServerRegistration,
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
McpLifecyclePhase::InitializeHandshake,
|
||||
McpLifecyclePhase::ToolDiscovery,
|
||||
] {
|
||||
let result = validator.run_phase(phase);
|
||||
assert!(matches!(result, McpPhaseResult::Success { .. }));
|
||||
}
|
||||
let _ = validator.record_failure(McpErrorSurface::new(
|
||||
McpLifecyclePhase::ResourceDiscovery,
|
||||
Some("alpha".to_string()),
|
||||
"resource listing failed",
|
||||
BTreeMap::from([("reason".to_string(), "timeout".to_string())]),
|
||||
true,
|
||||
));
|
||||
|
||||
// when
|
||||
let shutdown = validator.run_phase(McpLifecyclePhase::Shutdown);
|
||||
let cleanup = validator.run_phase(McpLifecyclePhase::Cleanup);
|
||||
|
||||
// then
|
||||
assert!(matches!(shutdown, McpPhaseResult::Success { .. }));
|
||||
assert!(matches!(cleanup, McpPhaseResult::Success { .. }));
|
||||
assert_eq!(
|
||||
validator.state().current_phase(),
|
||||
Some(McpLifecyclePhase::Cleanup)
|
||||
);
|
||||
assert!(validator
|
||||
.state()
|
||||
.phase_timestamp(McpLifecyclePhase::ErrorSurfacing)
|
||||
.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn error_surface_display_includes_phase_server_and_recoverable_flag() {
|
||||
// given
|
||||
let error = McpErrorSurface::new(
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
Some("alpha".to_string()),
|
||||
"process exited early",
|
||||
BTreeMap::from([("exit_code".to_string(), "1".to_string())]),
|
||||
true,
|
||||
);
|
||||
|
||||
// when
|
||||
let rendered = error.to_string();
|
||||
|
||||
// then
|
||||
assert!(rendered.contains("spawn_connect"));
|
||||
assert!(rendered.contains("process exited early"));
|
||||
assert!(rendered.contains("server: alpha"));
|
||||
assert!(rendered.contains("recoverable"));
|
||||
let trait_object: &dyn std::error::Error = &error;
|
||||
assert_eq!(trait_object.to_string(), rendered);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_nonrecoverable_failure_when_returning_to_ready_then_validator_rejects_resume() {
|
||||
// given
|
||||
let mut validator = McpLifecycleValidator::new();
|
||||
for phase in [
|
||||
McpLifecyclePhase::ConfigLoad,
|
||||
McpLifecyclePhase::ServerRegistration,
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
McpLifecyclePhase::InitializeHandshake,
|
||||
McpLifecyclePhase::ToolDiscovery,
|
||||
McpLifecyclePhase::Ready,
|
||||
] {
|
||||
let result = validator.run_phase(phase);
|
||||
assert!(matches!(result, McpPhaseResult::Success { .. }));
|
||||
}
|
||||
let _ = validator.record_failure(McpErrorSurface::new(
|
||||
McpLifecyclePhase::Invocation,
|
||||
Some("alpha".to_string()),
|
||||
"tool call corrupted the session",
|
||||
BTreeMap::from([("reason".to_string(), "invalid frame".to_string())]),
|
||||
false,
|
||||
));
|
||||
|
||||
// when
|
||||
let result = validator.run_phase(McpLifecyclePhase::Ready);
|
||||
|
||||
// then
|
||||
match result {
|
||||
McpPhaseResult::Failure { phase, error } => {
|
||||
assert_eq!(phase, McpLifecyclePhase::Ready);
|
||||
assert!(!error.recoverable);
|
||||
assert!(error.message.contains("non-recoverable"));
|
||||
}
|
||||
other => panic!("expected failure result, got {other:?}"),
|
||||
}
|
||||
assert_eq!(
|
||||
validator.state().current_phase(),
|
||||
Some(McpLifecyclePhase::ErrorSurfacing)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_recoverable_failure_when_returning_to_ready_then_validator_allows_resume() {
|
||||
// given
|
||||
let mut validator = McpLifecycleValidator::new();
|
||||
for phase in [
|
||||
McpLifecyclePhase::ConfigLoad,
|
||||
McpLifecyclePhase::ServerRegistration,
|
||||
McpLifecyclePhase::SpawnConnect,
|
||||
McpLifecyclePhase::InitializeHandshake,
|
||||
McpLifecyclePhase::ToolDiscovery,
|
||||
McpLifecyclePhase::Ready,
|
||||
] {
|
||||
let result = validator.run_phase(phase);
|
||||
assert!(matches!(result, McpPhaseResult::Success { .. }));
|
||||
}
|
||||
let _ = validator.record_failure(McpErrorSurface::new(
|
||||
McpLifecyclePhase::Invocation,
|
||||
Some("alpha".to_string()),
|
||||
"tool call failed but can be retried",
|
||||
BTreeMap::from([("reason".to_string(), "upstream timeout".to_string())]),
|
||||
true,
|
||||
));
|
||||
|
||||
// when
|
||||
let result = validator.run_phase(McpLifecyclePhase::Ready);
|
||||
|
||||
// then
|
||||
assert!(matches!(result, McpPhaseResult::Success { .. }));
|
||||
assert_eq!(
|
||||
validator.state().current_phase(),
|
||||
Some(McpLifecyclePhase::Ready)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,440 @@
|
||||
//! Minimal Model Context Protocol (MCP) server.
|
||||
//!
|
||||
//! Implements a newline-safe, LSP-framed JSON-RPC server over stdio that
|
||||
//! answers `initialize`, `tools/list`, and `tools/call` requests. The framing
|
||||
//! matches the client transport implemented in [`crate::mcp_stdio`] so this
|
||||
//! server can be driven by either an external MCP client (e.g. Claude
|
||||
//! Desktop) or `claw`'s own [`McpServerManager`](crate::McpServerManager).
|
||||
//!
|
||||
//! The server is intentionally small: it exposes a list of pre-built
|
||||
//! [`McpTool`] descriptors and delegates `tools/call` to a caller-supplied
|
||||
//! handler. Tool execution itself lives in the `tools` crate; this module is
|
||||
//! purely the transport + dispatch loop.
|
||||
//!
|
||||
//! [`McpTool`]: crate::mcp_stdio::McpTool
|
||||
|
||||
use std::io;
|
||||
|
||||
use serde_json::{json, Value as JsonValue};
|
||||
use tokio::io::{
|
||||
stdin, stdout, AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, Stdin, Stdout,
|
||||
};
|
||||
|
||||
use crate::mcp_stdio::{
|
||||
JsonRpcError, JsonRpcId, JsonRpcRequest, JsonRpcResponse, McpInitializeResult,
|
||||
McpInitializeServerInfo, McpListToolsResult, McpTool, McpToolCallContent, McpToolCallParams,
|
||||
McpToolCallResult,
|
||||
};
|
||||
|
||||
/// Protocol version the server advertises during `initialize`.
|
||||
///
|
||||
/// Matches the version used by the built-in client in
|
||||
/// [`crate::mcp_stdio`], so the two stay in lockstep.
|
||||
pub const MCP_SERVER_PROTOCOL_VERSION: &str = "2025-03-26";
|
||||
|
||||
/// Synchronous handler invoked for every `tools/call` request.
|
||||
///
|
||||
/// Returning `Ok(text)` yields a single `text` content block and
|
||||
/// `isError: false`. Returning `Err(message)` yields a `text` block with the
|
||||
/// error and `isError: true`, mirroring the error-surfacing convention used
|
||||
/// elsewhere in claw.
|
||||
pub type ToolCallHandler =
|
||||
Box<dyn Fn(&str, &JsonValue) -> Result<String, String> + Send + Sync + 'static>;
|
||||
|
||||
/// Configuration for an [`McpServer`] instance.
|
||||
///
|
||||
/// Named `McpServerSpec` rather than `McpServerConfig` to avoid colliding
|
||||
/// with the existing client-side [`crate::config::McpServerConfig`] that
|
||||
/// describes *remote* MCP servers the runtime connects to.
|
||||
pub struct McpServerSpec {
|
||||
/// Name advertised in the `serverInfo` field of the `initialize` response.
|
||||
pub server_name: String,
|
||||
/// Version advertised in the `serverInfo` field of the `initialize`
|
||||
/// response.
|
||||
pub server_version: String,
|
||||
/// Tool descriptors returned for `tools/list`.
|
||||
pub tools: Vec<McpTool>,
|
||||
/// Handler invoked for `tools/call`.
|
||||
pub tool_handler: ToolCallHandler,
|
||||
}
|
||||
|
||||
/// Minimal MCP stdio server.
|
||||
///
|
||||
/// The server runs a blocking read/dispatch/write loop over the current
|
||||
/// process's stdin/stdout, terminating cleanly when the peer closes the
|
||||
/// stream.
|
||||
pub struct McpServer {
|
||||
spec: McpServerSpec,
|
||||
stdin: BufReader<Stdin>,
|
||||
stdout: Stdout,
|
||||
}
|
||||
|
||||
impl McpServer {
|
||||
#[must_use]
|
||||
pub fn new(spec: McpServerSpec) -> Self {
|
||||
Self {
|
||||
spec,
|
||||
stdin: BufReader::new(stdin()),
|
||||
stdout: stdout(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Runs the server until the client closes stdin.
|
||||
///
|
||||
/// Returns `Ok(())` on clean EOF; any other I/O error is propagated so
|
||||
/// callers can log and exit non-zero.
|
||||
pub async fn run(&mut self) -> io::Result<()> {
|
||||
loop {
|
||||
let Some(payload) = read_frame(&mut self.stdin).await? else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
// Requests and notifications share a wire format; the absence of
|
||||
// `id` distinguishes notifications, which must never receive a
|
||||
// response.
|
||||
let message: JsonValue = match serde_json::from_slice(&payload) {
|
||||
Ok(value) => value,
|
||||
Err(error) => {
|
||||
// Parse error with null id per JSON-RPC 2.0 §4.2.
|
||||
let response = JsonRpcResponse::<JsonValue> {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: JsonRpcId::Null,
|
||||
result: None,
|
||||
error: Some(JsonRpcError {
|
||||
code: -32700,
|
||||
message: format!("parse error: {error}"),
|
||||
data: None,
|
||||
}),
|
||||
};
|
||||
write_response(&mut self.stdout, &response).await?;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if message.get("id").is_none() {
|
||||
// Notification: dispatch for side effects only (e.g. log),
|
||||
// but send no reply.
|
||||
continue;
|
||||
}
|
||||
|
||||
let request: JsonRpcRequest<JsonValue> = match serde_json::from_value(message) {
|
||||
Ok(request) => request,
|
||||
Err(error) => {
|
||||
let response = JsonRpcResponse::<JsonValue> {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: JsonRpcId::Null,
|
||||
result: None,
|
||||
error: Some(JsonRpcError {
|
||||
code: -32600,
|
||||
message: format!("invalid request: {error}"),
|
||||
data: None,
|
||||
}),
|
||||
};
|
||||
write_response(&mut self.stdout, &response).await?;
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
let response = self.dispatch(request);
|
||||
write_response(&mut self.stdout, &response).await?;
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch(&self, request: JsonRpcRequest<JsonValue>) -> JsonRpcResponse<JsonValue> {
|
||||
let id = request.id.clone();
|
||||
match request.method.as_str() {
|
||||
"initialize" => self.handle_initialize(id),
|
||||
"tools/list" => self.handle_tools_list(id),
|
||||
"tools/call" => self.handle_tools_call(id, request.params),
|
||||
other => JsonRpcResponse {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id,
|
||||
result: None,
|
||||
error: Some(JsonRpcError {
|
||||
code: -32601,
|
||||
message: format!("method not found: {other}"),
|
||||
data: None,
|
||||
}),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_initialize(&self, id: JsonRpcId) -> JsonRpcResponse<JsonValue> {
|
||||
let result = McpInitializeResult {
|
||||
protocol_version: MCP_SERVER_PROTOCOL_VERSION.to_string(),
|
||||
capabilities: json!({ "tools": {} }),
|
||||
server_info: McpInitializeServerInfo {
|
||||
name: self.spec.server_name.clone(),
|
||||
version: self.spec.server_version.clone(),
|
||||
},
|
||||
};
|
||||
JsonRpcResponse {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id,
|
||||
result: serde_json::to_value(result).ok(),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_tools_list(&self, id: JsonRpcId) -> JsonRpcResponse<JsonValue> {
|
||||
let result = McpListToolsResult {
|
||||
tools: self.spec.tools.clone(),
|
||||
next_cursor: None,
|
||||
};
|
||||
JsonRpcResponse {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id,
|
||||
result: serde_json::to_value(result).ok(),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_tools_call(
|
||||
&self,
|
||||
id: JsonRpcId,
|
||||
params: Option<JsonValue>,
|
||||
) -> JsonRpcResponse<JsonValue> {
|
||||
let Some(params) = params else {
|
||||
return invalid_params_response(id, "missing params for tools/call");
|
||||
};
|
||||
let call: McpToolCallParams = match serde_json::from_value(params) {
|
||||
Ok(value) => value,
|
||||
Err(error) => {
|
||||
return invalid_params_response(id, &format!("invalid tools/call params: {error}"));
|
||||
}
|
||||
};
|
||||
let arguments = call.arguments.unwrap_or_else(|| json!({}));
|
||||
let tool_result = (self.spec.tool_handler)(&call.name, &arguments);
|
||||
let (text, is_error) = match tool_result {
|
||||
Ok(text) => (text, false),
|
||||
Err(message) => (message, true),
|
||||
};
|
||||
let mut data = std::collections::BTreeMap::new();
|
||||
data.insert("text".to_string(), JsonValue::String(text));
|
||||
let call_result = McpToolCallResult {
|
||||
content: vec![McpToolCallContent {
|
||||
kind: "text".to_string(),
|
||||
data,
|
||||
}],
|
||||
structured_content: None,
|
||||
is_error: Some(is_error),
|
||||
meta: None,
|
||||
};
|
||||
JsonRpcResponse {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id,
|
||||
result: serde_json::to_value(call_result).ok(),
|
||||
error: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn invalid_params_response(id: JsonRpcId, message: &str) -> JsonRpcResponse<JsonValue> {
|
||||
JsonRpcResponse {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id,
|
||||
result: None,
|
||||
error: Some(JsonRpcError {
|
||||
code: -32602,
|
||||
message: message.to_string(),
|
||||
data: None,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reads a single LSP-framed JSON-RPC payload from `reader`.
|
||||
///
|
||||
/// Returns `Ok(None)` on clean EOF before any header bytes have been read,
|
||||
/// matching how [`crate::mcp_stdio::McpStdioProcess`] treats stream closure.
|
||||
async fn read_frame(reader: &mut BufReader<Stdin>) -> io::Result<Option<Vec<u8>>> {
|
||||
let mut content_length: Option<usize> = None;
|
||||
let mut first_header = true;
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
let bytes_read = reader.read_line(&mut line).await?;
|
||||
if bytes_read == 0 {
|
||||
if first_header {
|
||||
return Ok(None);
|
||||
}
|
||||
return Err(io::Error::new(
|
||||
io::ErrorKind::UnexpectedEof,
|
||||
"MCP stdio stream closed while reading headers",
|
||||
));
|
||||
}
|
||||
first_header = false;
|
||||
if line == "\r\n" || line == "\n" {
|
||||
break;
|
||||
}
|
||||
let header = line.trim_end_matches(['\r', '\n']);
|
||||
if let Some((name, value)) = header.split_once(':') {
|
||||
if name.trim().eq_ignore_ascii_case("Content-Length") {
|
||||
let parsed = value
|
||||
.trim()
|
||||
.parse::<usize>()
|
||||
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
|
||||
content_length = Some(parsed);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let content_length = content_length.ok_or_else(|| {
|
||||
io::Error::new(io::ErrorKind::InvalidData, "missing Content-Length header")
|
||||
})?;
|
||||
let mut payload = vec![0_u8; content_length];
|
||||
reader.read_exact(&mut payload).await?;
|
||||
Ok(Some(payload))
|
||||
}
|
||||
|
||||
async fn write_response(
|
||||
stdout: &mut Stdout,
|
||||
response: &JsonRpcResponse<JsonValue>,
|
||||
) -> io::Result<()> {
|
||||
let body = serde_json::to_vec(response)
|
||||
.map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
|
||||
let header = format!("Content-Length: {}\r\n\r\n", body.len());
|
||||
stdout.write_all(header.as_bytes()).await?;
|
||||
stdout.write_all(&body).await?;
|
||||
stdout.flush().await
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn dispatch_initialize_returns_server_info() {
|
||||
let server = McpServer {
|
||||
spec: McpServerSpec {
|
||||
server_name: "test".to_string(),
|
||||
server_version: "9.9.9".to_string(),
|
||||
tools: Vec::new(),
|
||||
tool_handler: Box::new(|_, _| Ok(String::new())),
|
||||
},
|
||||
stdin: BufReader::new(stdin()),
|
||||
stdout: stdout(),
|
||||
};
|
||||
let request = JsonRpcRequest::<JsonValue> {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: JsonRpcId::Number(1),
|
||||
method: "initialize".to_string(),
|
||||
params: None,
|
||||
};
|
||||
let response = server.dispatch(request);
|
||||
assert_eq!(response.id, JsonRpcId::Number(1));
|
||||
assert!(response.error.is_none());
|
||||
let result = response.result.expect("initialize result");
|
||||
assert_eq!(result["protocolVersion"], MCP_SERVER_PROTOCOL_VERSION);
|
||||
assert_eq!(result["serverInfo"]["name"], "test");
|
||||
assert_eq!(result["serverInfo"]["version"], "9.9.9");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatch_tools_list_returns_registered_tools() {
|
||||
let tool = McpTool {
|
||||
name: "echo".to_string(),
|
||||
description: Some("Echo".to_string()),
|
||||
input_schema: Some(json!({"type": "object"})),
|
||||
annotations: None,
|
||||
meta: None,
|
||||
};
|
||||
let server = McpServer {
|
||||
spec: McpServerSpec {
|
||||
server_name: "test".to_string(),
|
||||
server_version: "0.0.0".to_string(),
|
||||
tools: vec![tool.clone()],
|
||||
tool_handler: Box::new(|_, _| Ok(String::new())),
|
||||
},
|
||||
stdin: BufReader::new(stdin()),
|
||||
stdout: stdout(),
|
||||
};
|
||||
let request = JsonRpcRequest::<JsonValue> {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: JsonRpcId::Number(2),
|
||||
method: "tools/list".to_string(),
|
||||
params: None,
|
||||
};
|
||||
let response = server.dispatch(request);
|
||||
assert!(response.error.is_none());
|
||||
let result = response.result.expect("tools/list result");
|
||||
assert_eq!(result["tools"][0]["name"], "echo");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatch_tools_call_wraps_handler_output() {
|
||||
let server = McpServer {
|
||||
spec: McpServerSpec {
|
||||
server_name: "test".to_string(),
|
||||
server_version: "0.0.0".to_string(),
|
||||
tools: Vec::new(),
|
||||
tool_handler: Box::new(|name, args| Ok(format!("called {name} with {args}"))),
|
||||
},
|
||||
stdin: BufReader::new(stdin()),
|
||||
stdout: stdout(),
|
||||
};
|
||||
let request = JsonRpcRequest::<JsonValue> {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: JsonRpcId::Number(3),
|
||||
method: "tools/call".to_string(),
|
||||
params: Some(json!({
|
||||
"name": "echo",
|
||||
"arguments": {"text": "hi"}
|
||||
})),
|
||||
};
|
||||
let response = server.dispatch(request);
|
||||
assert!(response.error.is_none());
|
||||
let result = response.result.expect("tools/call result");
|
||||
assert_eq!(result["isError"], false);
|
||||
assert_eq!(result["content"][0]["type"], "text");
|
||||
assert!(result["content"][0]["text"]
|
||||
.as_str()
|
||||
.unwrap()
|
||||
.starts_with("called echo"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatch_tools_call_surfaces_handler_error() {
|
||||
let server = McpServer {
|
||||
spec: McpServerSpec {
|
||||
server_name: "test".to_string(),
|
||||
server_version: "0.0.0".to_string(),
|
||||
tools: Vec::new(),
|
||||
tool_handler: Box::new(|_, _| Err("boom".to_string())),
|
||||
},
|
||||
stdin: BufReader::new(stdin()),
|
||||
stdout: stdout(),
|
||||
};
|
||||
let request = JsonRpcRequest::<JsonValue> {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: JsonRpcId::Number(4),
|
||||
method: "tools/call".to_string(),
|
||||
params: Some(json!({"name": "broken"})),
|
||||
};
|
||||
let response = server.dispatch(request);
|
||||
let result = response.result.expect("tools/call result");
|
||||
assert_eq!(result["isError"], true);
|
||||
assert_eq!(result["content"][0]["text"], "boom");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dispatch_unknown_method_returns_method_not_found() {
|
||||
let server = McpServer {
|
||||
spec: McpServerSpec {
|
||||
server_name: "test".to_string(),
|
||||
server_version: "0.0.0".to_string(),
|
||||
tools: Vec::new(),
|
||||
tool_handler: Box::new(|_, _| Ok(String::new())),
|
||||
},
|
||||
stdin: BufReader::new(stdin()),
|
||||
stdout: stdout(),
|
||||
};
|
||||
let request = JsonRpcRequest::<JsonValue> {
|
||||
jsonrpc: "2.0".to_string(),
|
||||
id: JsonRpcId::Number(5),
|
||||
method: "nonsense".to_string(),
|
||||
params: None,
|
||||
};
|
||||
let response = server.dispatch(request);
|
||||
let error = response.error.expect("error payload");
|
||||
assert_eq!(error.code, -32601);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,920 @@
|
||||
#![allow(
|
||||
clippy::await_holding_lock,
|
||||
clippy::doc_markdown,
|
||||
clippy::match_same_arms,
|
||||
clippy::must_use_candidate,
|
||||
clippy::uninlined_format_args,
|
||||
clippy::unnested_or_patterns
|
||||
)]
|
||||
//! Bridge between MCP tool surface (ListMcpResources, ReadMcpResource, McpAuth, MCP)
|
||||
//! and the existing McpServerManager runtime.
|
||||
//!
|
||||
//! Provides a stateful client registry that tool handlers can use to
|
||||
//! connect to MCP servers and invoke their capabilities.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex, OnceLock};
|
||||
|
||||
use crate::mcp::mcp_tool_name;
|
||||
use crate::mcp_stdio::McpServerManager;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Status of a managed MCP server connection.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum McpConnectionStatus {
|
||||
Disconnected,
|
||||
Connecting,
|
||||
Connected,
|
||||
AuthRequired,
|
||||
Error,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for McpConnectionStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Disconnected => write!(f, "disconnected"),
|
||||
Self::Connecting => write!(f, "connecting"),
|
||||
Self::Connected => write!(f, "connected"),
|
||||
Self::AuthRequired => write!(f, "auth_required"),
|
||||
Self::Error => write!(f, "error"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Metadata about an MCP resource.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpResourceInfo {
|
||||
pub uri: String,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub mime_type: Option<String>,
|
||||
}
|
||||
|
||||
/// Metadata about an MCP tool exposed by a server.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpToolInfo {
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub input_schema: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Tracked state of an MCP server connection.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpServerState {
|
||||
pub server_name: String,
|
||||
pub status: McpConnectionStatus,
|
||||
pub tools: Vec<McpToolInfo>,
|
||||
pub resources: Vec<McpResourceInfo>,
|
||||
pub server_info: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct McpToolRegistry {
|
||||
inner: Arc<Mutex<HashMap<String, McpServerState>>>,
|
||||
manager: Arc<OnceLock<Arc<Mutex<McpServerManager>>>>,
|
||||
}
|
||||
|
||||
impl McpToolRegistry {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn set_manager(
|
||||
&self,
|
||||
manager: Arc<Mutex<McpServerManager>>,
|
||||
) -> Result<(), Arc<Mutex<McpServerManager>>> {
|
||||
self.manager.set(manager)
|
||||
}
|
||||
|
||||
pub fn register_server(
|
||||
&self,
|
||||
server_name: &str,
|
||||
status: McpConnectionStatus,
|
||||
tools: Vec<McpToolInfo>,
|
||||
resources: Vec<McpResourceInfo>,
|
||||
server_info: Option<String>,
|
||||
) {
|
||||
let mut inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
inner.insert(
|
||||
server_name.to_owned(),
|
||||
McpServerState {
|
||||
server_name: server_name.to_owned(),
|
||||
status,
|
||||
tools,
|
||||
resources,
|
||||
server_info,
|
||||
error_message: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
pub fn get_server(&self, server_name: &str) -> Option<McpServerState> {
|
||||
let inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
inner.get(server_name).cloned()
|
||||
}
|
||||
|
||||
pub fn list_servers(&self) -> Vec<McpServerState> {
|
||||
let inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
inner.values().cloned().collect()
|
||||
}
|
||||
|
||||
pub fn list_resources(&self, server_name: &str) -> Result<Vec<McpResourceInfo>, String> {
|
||||
let inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
match inner.get(server_name) {
|
||||
Some(state) => {
|
||||
if state.status != McpConnectionStatus::Connected {
|
||||
return Err(format!(
|
||||
"server '{}' is not connected (status: {})",
|
||||
server_name, state.status
|
||||
));
|
||||
}
|
||||
Ok(state.resources.clone())
|
||||
}
|
||||
None => Err(format!("server '{}' not found", server_name)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read_resource(&self, server_name: &str, uri: &str) -> Result<McpResourceInfo, String> {
|
||||
let inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
let state = inner
|
||||
.get(server_name)
|
||||
.ok_or_else(|| format!("server '{}' not found", server_name))?;
|
||||
|
||||
if state.status != McpConnectionStatus::Connected {
|
||||
return Err(format!(
|
||||
"server '{}' is not connected (status: {})",
|
||||
server_name, state.status
|
||||
));
|
||||
}
|
||||
|
||||
state
|
||||
.resources
|
||||
.iter()
|
||||
.find(|r| r.uri == uri)
|
||||
.cloned()
|
||||
.ok_or_else(|| format!("resource '{}' not found on server '{}'", uri, server_name))
|
||||
}
|
||||
|
||||
pub fn list_tools(&self, server_name: &str) -> Result<Vec<McpToolInfo>, String> {
|
||||
let inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
match inner.get(server_name) {
|
||||
Some(state) => {
|
||||
if state.status != McpConnectionStatus::Connected {
|
||||
return Err(format!(
|
||||
"server '{}' is not connected (status: {})",
|
||||
server_name, state.status
|
||||
));
|
||||
}
|
||||
Ok(state.tools.clone())
|
||||
}
|
||||
None => Err(format!("server '{}' not found", server_name)),
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_tool_call(
|
||||
manager: Arc<Mutex<McpServerManager>>,
|
||||
qualified_tool_name: String,
|
||||
arguments: Option<serde_json::Value>,
|
||||
) -> Result<serde_json::Value, String> {
|
||||
let join_handle = std::thread::Builder::new()
|
||||
.name(format!("mcp-tool-call-{qualified_tool_name}"))
|
||||
.spawn(move || {
|
||||
let runtime = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.map_err(|error| format!("failed to create MCP tool runtime: {error}"))?;
|
||||
|
||||
runtime.block_on(async move {
|
||||
let response = {
|
||||
let mut manager = manager
|
||||
.lock()
|
||||
.map_err(|_| "mcp server manager lock poisoned".to_string())?;
|
||||
manager
|
||||
.discover_tools()
|
||||
.await
|
||||
.map_err(|error| error.to_string())?;
|
||||
let response = manager
|
||||
.call_tool(&qualified_tool_name, arguments)
|
||||
.await
|
||||
.map_err(|error| error.to_string());
|
||||
let shutdown = manager.shutdown().await.map_err(|error| error.to_string());
|
||||
|
||||
match (response, shutdown) {
|
||||
(Ok(response), Ok(())) => Ok(response),
|
||||
(Err(error), Ok(())) | (Err(error), Err(_)) => Err(error),
|
||||
(Ok(_), Err(error)) => Err(error),
|
||||
}
|
||||
}?;
|
||||
|
||||
if let Some(error) = response.error {
|
||||
return Err(format!(
|
||||
"MCP server returned JSON-RPC error for tools/call: {} ({})",
|
||||
error.message, error.code
|
||||
));
|
||||
}
|
||||
|
||||
let result = response.result.ok_or_else(|| {
|
||||
"MCP server returned no result for tools/call".to_string()
|
||||
})?;
|
||||
|
||||
serde_json::to_value(result)
|
||||
.map_err(|error| format!("failed to serialize MCP tool result: {error}"))
|
||||
})
|
||||
})
|
||||
.map_err(|error| format!("failed to spawn MCP tool call thread: {error}"))?;
|
||||
|
||||
join_handle.join().map_err(|panic_payload| {
|
||||
if let Some(message) = panic_payload.downcast_ref::<&str>() {
|
||||
format!("MCP tool call thread panicked: {message}")
|
||||
} else if let Some(message) = panic_payload.downcast_ref::<String>() {
|
||||
format!("MCP tool call thread panicked: {message}")
|
||||
} else {
|
||||
"MCP tool call thread panicked".to_string()
|
||||
}
|
||||
})?
|
||||
}
|
||||
|
||||
pub fn call_tool(
|
||||
&self,
|
||||
server_name: &str,
|
||||
tool_name: &str,
|
||||
arguments: &serde_json::Value,
|
||||
) -> Result<serde_json::Value, String> {
|
||||
let inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
let state = inner
|
||||
.get(server_name)
|
||||
.ok_or_else(|| format!("server '{}' not found", server_name))?;
|
||||
|
||||
if state.status != McpConnectionStatus::Connected {
|
||||
return Err(format!(
|
||||
"server '{}' is not connected (status: {})",
|
||||
server_name, state.status
|
||||
));
|
||||
}
|
||||
|
||||
if !state.tools.iter().any(|t| t.name == tool_name) {
|
||||
return Err(format!(
|
||||
"tool '{}' not found on server '{}'",
|
||||
tool_name, server_name
|
||||
));
|
||||
}
|
||||
|
||||
drop(inner);
|
||||
|
||||
let manager = self
|
||||
.manager
|
||||
.get()
|
||||
.cloned()
|
||||
.ok_or_else(|| "MCP server manager is not configured".to_string())?;
|
||||
|
||||
Self::spawn_tool_call(
|
||||
manager,
|
||||
mcp_tool_name(server_name, tool_name),
|
||||
(!arguments.is_null()).then(|| arguments.clone()),
|
||||
)
|
||||
}
|
||||
|
||||
/// Set auth status for a server.
|
||||
pub fn set_auth_status(
|
||||
&self,
|
||||
server_name: &str,
|
||||
status: McpConnectionStatus,
|
||||
) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
let state = inner
|
||||
.get_mut(server_name)
|
||||
.ok_or_else(|| format!("server '{}' not found", server_name))?;
|
||||
state.status = status;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Disconnect / remove a server.
|
||||
pub fn disconnect(&self, server_name: &str) -> Option<McpServerState> {
|
||||
let mut inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
inner.remove(server_name)
|
||||
}
|
||||
|
||||
/// Number of registered servers.
|
||||
#[must_use]
|
||||
pub fn len(&self) -> usize {
|
||||
let inner = self.inner.lock().expect("mcp registry lock poisoned");
|
||||
inner.len()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs;
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use super::*;
|
||||
use crate::config::{
|
||||
ConfigSource, McpServerConfig, McpStdioServerConfig, ScopedMcpServerConfig,
|
||||
};
|
||||
|
||||
fn temp_dir() -> PathBuf {
|
||||
static NEXT_TEMP_DIR_ID: AtomicU64 = AtomicU64::new(0);
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_nanos();
|
||||
let unique_id = NEXT_TEMP_DIR_ID.fetch_add(1, Ordering::Relaxed);
|
||||
std::env::temp_dir().join(format!("runtime-mcp-tool-bridge-{nanos}-{unique_id}"))
|
||||
}
|
||||
|
||||
fn cleanup_script(script_path: &Path) {
|
||||
if let Some(root) = script_path.parent() {
|
||||
let _ = fs::remove_dir_all(root);
|
||||
}
|
||||
}
|
||||
|
||||
fn write_bridge_mcp_server_script() -> PathBuf {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("temp dir");
|
||||
let script_path = root.join("bridge-mcp-server.py");
|
||||
let script = [
|
||||
"#!/usr/bin/env python3",
|
||||
"import json, os, sys",
|
||||
"LABEL = os.environ.get('MCP_SERVER_LABEL', 'server')",
|
||||
"LOG_PATH = os.environ.get('MCP_LOG_PATH')",
|
||||
"",
|
||||
"def log(method):",
|
||||
" if LOG_PATH:",
|
||||
" with open(LOG_PATH, 'a', encoding='utf-8') as handle:",
|
||||
" handle.write(f'{method}\\n')",
|
||||
"",
|
||||
"def read_message():",
|
||||
" header = b''",
|
||||
r" while not header.endswith(b'\r\n\r\n'):",
|
||||
" chunk = sys.stdin.buffer.read(1)",
|
||||
" if not chunk:",
|
||||
" return None",
|
||||
" header += chunk",
|
||||
" length = 0",
|
||||
r" for line in header.decode().split('\r\n'):",
|
||||
r" if line.lower().startswith('content-length:'):",
|
||||
r" length = int(line.split(':', 1)[1].strip())",
|
||||
" payload = sys.stdin.buffer.read(length)",
|
||||
" return json.loads(payload.decode())",
|
||||
"",
|
||||
"def send_message(message):",
|
||||
" payload = json.dumps(message).encode()",
|
||||
r" sys.stdout.buffer.write(f'Content-Length: {len(payload)}\r\n\r\n'.encode() + payload)",
|
||||
" sys.stdout.buffer.flush()",
|
||||
"",
|
||||
"while True:",
|
||||
" request = read_message()",
|
||||
" if request is None:",
|
||||
" break",
|
||||
" method = request['method']",
|
||||
" log(method)",
|
||||
" if method == 'initialize':",
|
||||
" send_message({",
|
||||
" 'jsonrpc': '2.0',",
|
||||
" 'id': request['id'],",
|
||||
" 'result': {",
|
||||
" 'protocolVersion': request['params']['protocolVersion'],",
|
||||
" 'capabilities': {'tools': {}},",
|
||||
" 'serverInfo': {'name': LABEL, 'version': '1.0.0'}",
|
||||
" }",
|
||||
" })",
|
||||
" elif method == 'tools/list':",
|
||||
" send_message({",
|
||||
" 'jsonrpc': '2.0',",
|
||||
" 'id': request['id'],",
|
||||
" 'result': {",
|
||||
" 'tools': [",
|
||||
" {",
|
||||
" 'name': 'echo',",
|
||||
" 'description': f'Echo tool for {LABEL}',",
|
||||
" 'inputSchema': {",
|
||||
" 'type': 'object',",
|
||||
" 'properties': {'text': {'type': 'string'}},",
|
||||
" 'required': ['text']",
|
||||
" }",
|
||||
" }",
|
||||
" ]",
|
||||
" }",
|
||||
" })",
|
||||
" elif method == 'tools/call':",
|
||||
" args = request['params'].get('arguments') or {}",
|
||||
" text = args.get('text', '')",
|
||||
" send_message({",
|
||||
" 'jsonrpc': '2.0',",
|
||||
" 'id': request['id'],",
|
||||
" 'result': {",
|
||||
" 'content': [{'type': 'text', 'text': f'{LABEL}:{text}'}],",
|
||||
" 'structuredContent': {'server': LABEL, 'echoed': text},",
|
||||
" 'isError': False",
|
||||
" }",
|
||||
" })",
|
||||
" else:",
|
||||
" send_message({",
|
||||
" 'jsonrpc': '2.0',",
|
||||
" 'id': request['id'],",
|
||||
" 'error': {'code': -32601, 'message': f'unknown method: {method}'},",
|
||||
" })",
|
||||
"",
|
||||
]
|
||||
.join("\n");
|
||||
fs::write(&script_path, script).expect("write script");
|
||||
let mut permissions = fs::metadata(&script_path).expect("metadata").permissions();
|
||||
permissions.set_mode(0o755);
|
||||
fs::set_permissions(&script_path, permissions).expect("chmod");
|
||||
script_path
|
||||
}
|
||||
|
||||
fn manager_server_config(
|
||||
script_path: &Path,
|
||||
server_name: &str,
|
||||
log_path: &Path,
|
||||
) -> ScopedMcpServerConfig {
|
||||
ScopedMcpServerConfig {
|
||||
scope: ConfigSource::Local,
|
||||
config: McpServerConfig::Stdio(McpStdioServerConfig {
|
||||
command: "python3".to_string(),
|
||||
args: vec![script_path.to_string_lossy().into_owned()],
|
||||
env: BTreeMap::from([
|
||||
("MCP_SERVER_LABEL".to_string(), server_name.to_string()),
|
||||
(
|
||||
"MCP_LOG_PATH".to_string(),
|
||||
log_path.to_string_lossy().into_owned(),
|
||||
),
|
||||
]),
|
||||
tool_call_timeout_ms: Some(1_000),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn registers_and_retrieves_server() {
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"test-server",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![McpToolInfo {
|
||||
name: "greet".into(),
|
||||
description: Some("Greet someone".into()),
|
||||
input_schema: None,
|
||||
}],
|
||||
vec![McpResourceInfo {
|
||||
uri: "res://data".into(),
|
||||
name: "Data".into(),
|
||||
description: None,
|
||||
mime_type: Some("application/json".into()),
|
||||
}],
|
||||
Some("TestServer v1.0".into()),
|
||||
);
|
||||
|
||||
let server = registry.get_server("test-server").expect("should exist");
|
||||
assert_eq!(server.status, McpConnectionStatus::Connected);
|
||||
assert_eq!(server.tools.len(), 1);
|
||||
assert_eq!(server.resources.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lists_resources_from_connected_server() {
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![],
|
||||
vec![McpResourceInfo {
|
||||
uri: "res://alpha".into(),
|
||||
name: "Alpha".into(),
|
||||
description: None,
|
||||
mime_type: None,
|
||||
}],
|
||||
None,
|
||||
);
|
||||
|
||||
let resources = registry.list_resources("srv").expect("should succeed");
|
||||
assert_eq!(resources.len(), 1);
|
||||
assert_eq!(resources[0].uri, "res://alpha");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_resource_listing_for_disconnected_server() {
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::Disconnected,
|
||||
vec![],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
assert!(registry.list_resources("srv").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_specific_resource() {
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![],
|
||||
vec![McpResourceInfo {
|
||||
uri: "res://data".into(),
|
||||
name: "Data".into(),
|
||||
description: Some("Test data".into()),
|
||||
mime_type: Some("text/plain".into()),
|
||||
}],
|
||||
None,
|
||||
);
|
||||
|
||||
let resource = registry
|
||||
.read_resource("srv", "res://data")
|
||||
.expect("should find");
|
||||
assert_eq!(resource.name, "Data");
|
||||
|
||||
assert!(registry.read_resource("srv", "res://missing").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_connected_server_without_manager_when_calling_tool_then_it_errors() {
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![McpToolInfo {
|
||||
name: "greet".into(),
|
||||
description: None,
|
||||
input_schema: None,
|
||||
}],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
|
||||
let error = registry
|
||||
.call_tool("srv", "greet", &serde_json::json!({"name": "world"}))
|
||||
.expect_err("should require a configured manager");
|
||||
assert!(error.contains("MCP server manager is not configured"));
|
||||
|
||||
// Unknown tool should fail
|
||||
assert!(registry
|
||||
.call_tool("srv", "missing", &serde_json::json!({}))
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn given_connected_server_with_manager_when_calling_tool_then_it_returns_live_result() {
|
||||
let script_path = write_bridge_mcp_server_script();
|
||||
let root = script_path.parent().expect("script parent");
|
||||
let log_path = root.join("bridge.log");
|
||||
let servers = BTreeMap::from([(
|
||||
"alpha".to_string(),
|
||||
manager_server_config(&script_path, "alpha", &log_path),
|
||||
)]);
|
||||
let manager = Arc::new(Mutex::new(McpServerManager::from_servers(&servers)));
|
||||
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"alpha",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![McpToolInfo {
|
||||
name: "echo".into(),
|
||||
description: Some("Echo tool for alpha".into()),
|
||||
input_schema: Some(serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {"text": {"type": "string"}},
|
||||
"required": ["text"]
|
||||
})),
|
||||
}],
|
||||
vec![],
|
||||
Some("bridge test server".into()),
|
||||
);
|
||||
registry
|
||||
.set_manager(Arc::clone(&manager))
|
||||
.expect("manager should only be set once");
|
||||
|
||||
let result = registry
|
||||
.call_tool("alpha", "echo", &serde_json::json!({"text": "hello"}))
|
||||
.expect("should return live MCP result");
|
||||
|
||||
assert_eq!(
|
||||
result["structuredContent"]["server"],
|
||||
serde_json::json!("alpha")
|
||||
);
|
||||
assert_eq!(
|
||||
result["structuredContent"]["echoed"],
|
||||
serde_json::json!("hello")
|
||||
);
|
||||
assert_eq!(
|
||||
result["content"][0]["text"],
|
||||
serde_json::json!("alpha:hello")
|
||||
);
|
||||
|
||||
let log = fs::read_to_string(&log_path).expect("read log");
|
||||
assert_eq!(
|
||||
log.lines().collect::<Vec<_>>(),
|
||||
vec!["initialize", "tools/list", "tools/call"]
|
||||
);
|
||||
|
||||
cleanup_script(&script_path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_tool_call_on_disconnected_server() {
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::AuthRequired,
|
||||
vec![McpToolInfo {
|
||||
name: "greet".into(),
|
||||
description: None,
|
||||
input_schema: None,
|
||||
}],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(registry
|
||||
.call_tool("srv", "greet", &serde_json::json!({}))
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sets_auth_and_disconnects() {
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::AuthRequired,
|
||||
vec![],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
|
||||
registry
|
||||
.set_auth_status("srv", McpConnectionStatus::Connected)
|
||||
.expect("should succeed");
|
||||
let state = registry.get_server("srv").unwrap();
|
||||
assert_eq!(state.status, McpConnectionStatus::Connected);
|
||||
|
||||
let removed = registry.disconnect("srv");
|
||||
assert!(removed.is_some());
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_operations_on_missing_server() {
|
||||
let registry = McpToolRegistry::new();
|
||||
assert!(registry.list_resources("missing").is_err());
|
||||
assert!(registry.read_resource("missing", "uri").is_err());
|
||||
assert!(registry.list_tools("missing").is_err());
|
||||
assert!(registry
|
||||
.call_tool("missing", "tool", &serde_json::json!({}))
|
||||
.is_err());
|
||||
assert!(registry
|
||||
.set_auth_status("missing", McpConnectionStatus::Connected)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_connection_status_display_all_variants() {
|
||||
// given
|
||||
let cases = [
|
||||
(McpConnectionStatus::Disconnected, "disconnected"),
|
||||
(McpConnectionStatus::Connecting, "connecting"),
|
||||
(McpConnectionStatus::Connected, "connected"),
|
||||
(McpConnectionStatus::AuthRequired, "auth_required"),
|
||||
(McpConnectionStatus::Error, "error"),
|
||||
];
|
||||
|
||||
// when
|
||||
let rendered: Vec<_> = cases
|
||||
.into_iter()
|
||||
.map(|(status, expected)| (status.to_string(), expected))
|
||||
.collect();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
rendered,
|
||||
vec![
|
||||
("disconnected".to_string(), "disconnected"),
|
||||
("connecting".to_string(), "connecting"),
|
||||
("connected".to_string(), "connected"),
|
||||
("auth_required".to_string(), "auth_required"),
|
||||
("error".to_string(), "error"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_servers_returns_all_registered() {
|
||||
// given
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"alpha",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
registry.register_server(
|
||||
"beta",
|
||||
McpConnectionStatus::Connecting,
|
||||
vec![],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let servers = registry.list_servers();
|
||||
|
||||
// then
|
||||
assert_eq!(servers.len(), 2);
|
||||
assert!(servers.iter().any(|server| server.server_name == "alpha"));
|
||||
assert!(servers.iter().any(|server| server.server_name == "beta"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_tools_from_connected_server() {
|
||||
// given
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![McpToolInfo {
|
||||
name: "inspect".into(),
|
||||
description: Some("Inspect data".into()),
|
||||
input_schema: Some(serde_json::json!({"type": "object"})),
|
||||
}],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let tools = registry.list_tools("srv").expect("tools should list");
|
||||
|
||||
// then
|
||||
assert_eq!(tools.len(), 1);
|
||||
assert_eq!(tools[0].name, "inspect");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_tools_rejects_disconnected_server() {
|
||||
// given
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::AuthRequired,
|
||||
vec![],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let result = registry.list_tools("srv");
|
||||
|
||||
// then
|
||||
let error = result.expect_err("non-connected server should fail");
|
||||
assert!(error.contains("not connected"));
|
||||
assert!(error.contains("auth_required"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_tools_rejects_missing_server() {
|
||||
// given
|
||||
let registry = McpToolRegistry::new();
|
||||
|
||||
// when
|
||||
let result = registry.list_tools("missing");
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
result.expect_err("missing server should fail"),
|
||||
"server 'missing' not found"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn get_server_returns_none_for_missing() {
|
||||
// given
|
||||
let registry = McpToolRegistry::new();
|
||||
|
||||
// when
|
||||
let server = registry.get_server("missing");
|
||||
|
||||
// then
|
||||
assert!(server.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn call_tool_payload_structure() {
|
||||
let script_path = write_bridge_mcp_server_script();
|
||||
let root = script_path.parent().expect("script parent");
|
||||
let log_path = root.join("payload.log");
|
||||
let servers = BTreeMap::from([(
|
||||
"srv".to_string(),
|
||||
manager_server_config(&script_path, "srv", &log_path),
|
||||
)]);
|
||||
let registry = McpToolRegistry::new();
|
||||
let arguments = serde_json::json!({"text": "world"});
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![McpToolInfo {
|
||||
name: "echo".into(),
|
||||
description: Some("Echo tool for srv".into()),
|
||||
input_schema: Some(serde_json::json!({
|
||||
"type": "object",
|
||||
"properties": {"text": {"type": "string"}},
|
||||
"required": ["text"]
|
||||
})),
|
||||
}],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
registry
|
||||
.set_manager(Arc::new(Mutex::new(McpServerManager::from_servers(
|
||||
&servers,
|
||||
))))
|
||||
.expect("manager should only be set once");
|
||||
|
||||
let result = registry
|
||||
.call_tool("srv", "echo", &arguments)
|
||||
.expect("tool should return live payload");
|
||||
|
||||
assert_eq!(result["structuredContent"]["server"], "srv");
|
||||
assert_eq!(result["structuredContent"]["echoed"], "world");
|
||||
assert_eq!(result["content"][0]["text"], "srv:world");
|
||||
|
||||
cleanup_script(&script_path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upsert_overwrites_existing_server() {
|
||||
// given
|
||||
let registry = McpToolRegistry::new();
|
||||
registry.register_server("srv", McpConnectionStatus::Connecting, vec![], vec![], None);
|
||||
|
||||
// when
|
||||
registry.register_server(
|
||||
"srv",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![McpToolInfo {
|
||||
name: "inspect".into(),
|
||||
description: None,
|
||||
input_schema: None,
|
||||
}],
|
||||
vec![],
|
||||
Some("Inspector".into()),
|
||||
);
|
||||
let state = registry.get_server("srv").expect("server should exist");
|
||||
|
||||
// then
|
||||
assert_eq!(state.status, McpConnectionStatus::Connected);
|
||||
assert_eq!(state.tools.len(), 1);
|
||||
assert_eq!(state.server_info.as_deref(), Some("Inspector"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disconnect_missing_returns_none() {
|
||||
// given
|
||||
let registry = McpToolRegistry::new();
|
||||
|
||||
// when
|
||||
let removed = registry.disconnect("missing");
|
||||
|
||||
// then
|
||||
assert!(removed.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn len_and_is_empty_transitions() {
|
||||
// given
|
||||
let registry = McpToolRegistry::new();
|
||||
|
||||
// when
|
||||
registry.register_server(
|
||||
"alpha",
|
||||
McpConnectionStatus::Connected,
|
||||
vec![],
|
||||
vec![],
|
||||
None,
|
||||
);
|
||||
registry.register_server("beta", McpConnectionStatus::Connected, vec![], vec![], None);
|
||||
let after_create = registry.len();
|
||||
registry.disconnect("alpha");
|
||||
let after_first_remove = registry.len();
|
||||
registry.disconnect("beta");
|
||||
|
||||
// then
|
||||
assert_eq!(after_create, 2);
|
||||
assert_eq!(after_first_remove, 1);
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,7 @@ use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::config::OAuthConfig;
|
||||
|
||||
/// Persisted OAuth access token bundle used by the CLI.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct OAuthTokenSet {
|
||||
pub access_token: String,
|
||||
@@ -17,6 +18,7 @@ pub struct OAuthTokenSet {
|
||||
pub scopes: Vec<String>,
|
||||
}
|
||||
|
||||
/// PKCE verifier/challenge pair generated for an OAuth authorization flow.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PkceCodePair {
|
||||
pub verifier: String,
|
||||
@@ -24,6 +26,7 @@ pub struct PkceCodePair {
|
||||
pub challenge_method: PkceChallengeMethod,
|
||||
}
|
||||
|
||||
/// Challenge algorithms supported by the local PKCE helpers.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PkceChallengeMethod {
|
||||
S256,
|
||||
@@ -38,6 +41,7 @@ impl PkceChallengeMethod {
|
||||
}
|
||||
}
|
||||
|
||||
/// Parameters needed to build an authorization URL for browser-based login.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct OAuthAuthorizationRequest {
|
||||
pub authorize_url: String,
|
||||
@@ -50,6 +54,7 @@ pub struct OAuthAuthorizationRequest {
|
||||
pub extra_params: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
/// Request body for exchanging an OAuth authorization code for tokens.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct OAuthTokenExchangeRequest {
|
||||
pub grant_type: &'static str,
|
||||
@@ -60,6 +65,7 @@ pub struct OAuthTokenExchangeRequest {
|
||||
pub state: String,
|
||||
}
|
||||
|
||||
/// Request body for refreshing an existing OAuth token set.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct OAuthRefreshRequest {
|
||||
pub grant_type: &'static str,
|
||||
@@ -68,6 +74,7 @@ pub struct OAuthRefreshRequest {
|
||||
pub scopes: Vec<String>,
|
||||
}
|
||||
|
||||
/// Parsed query parameters returned to the local OAuth callback endpoint.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct OAuthCallbackParams {
|
||||
pub code: Option<String>,
|
||||
@@ -324,12 +331,19 @@ fn generate_random_token(bytes: usize) -> io::Result<String> {
|
||||
}
|
||||
|
||||
fn credentials_home_dir() -> io::Result<PathBuf> {
|
||||
if let Some(path) = std::env::var_os("CLAUDE_CONFIG_HOME") {
|
||||
if let Some(path) = std::env::var_os("CLAW_CONFIG_HOME") {
|
||||
return Ok(PathBuf::from(path));
|
||||
}
|
||||
let home = std::env::var_os("HOME")
|
||||
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "HOME is not set"))?;
|
||||
Ok(PathBuf::from(home).join(".claude"))
|
||||
.or_else(|| std::env::var_os("USERPROFILE"))
|
||||
.ok_or_else(|| {
|
||||
io::Error::new(
|
||||
io::ErrorKind::NotFound,
|
||||
"HOME is not set (on Windows, set USERPROFILE or HOME, \
|
||||
or use CLAW_CONFIG_HOME to point directly at the config directory)",
|
||||
)
|
||||
})?;
|
||||
Ok(PathBuf::from(home).join(".claw"))
|
||||
}
|
||||
|
||||
fn read_credentials_root(path: &PathBuf) -> io::Result<Map<String, Value>> {
|
||||
@@ -442,7 +456,7 @@ fn decode_hex(byte: u8) -> Result<u8, String> {
|
||||
b'0'..=b'9' => Ok(byte - b'0'),
|
||||
b'a'..=b'f' => Ok(byte - b'a' + 10),
|
||||
b'A'..=b'F' => Ok(byte - b'A' + 10),
|
||||
_ => Err(format!("invalid percent-encoding byte: {byte}")),
|
||||
_ => Err(format!("invalid percent byte: {byte}")),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -541,7 +555,7 @@ mod tests {
|
||||
fn oauth_credentials_round_trip_and_clear_preserves_other_fields() {
|
||||
let _guard = env_lock();
|
||||
let config_home = temp_config_home();
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", &config_home);
|
||||
std::env::set_var("CLAW_CONFIG_HOME", &config_home);
|
||||
let path = credentials_path().expect("credentials path");
|
||||
std::fs::create_dir_all(path.parent().expect("parent")).expect("create parent");
|
||||
std::fs::write(&path, "{\"other\":\"value\"}\n").expect("seed credentials");
|
||||
@@ -567,7 +581,7 @@ mod tests {
|
||||
assert!(cleared.contains("\"other\": \"value\""));
|
||||
assert!(!cleared.contains("\"oauth\""));
|
||||
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
std::env::remove_var("CLAW_CONFIG_HOME");
|
||||
std::fs::remove_dir_all(config_home).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
@@ -0,0 +1,551 @@
|
||||
#![allow(
|
||||
clippy::match_wildcard_for_single_variants,
|
||||
clippy::must_use_candidate,
|
||||
clippy::uninlined_format_args
|
||||
)]
|
||||
//! Permission enforcement layer that gates tool execution based on the
|
||||
//! active `PermissionPolicy`.
|
||||
|
||||
use crate::permissions::{PermissionMode, PermissionOutcome, PermissionPolicy};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(tag = "outcome")]
|
||||
pub enum EnforcementResult {
|
||||
/// Tool execution is allowed.
|
||||
Allowed,
|
||||
/// Tool execution was denied due to insufficient permissions.
|
||||
Denied {
|
||||
tool: String,
|
||||
active_mode: String,
|
||||
required_mode: String,
|
||||
reason: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct PermissionEnforcer {
|
||||
policy: PermissionPolicy,
|
||||
}
|
||||
|
||||
impl PermissionEnforcer {
|
||||
#[must_use]
|
||||
pub fn new(policy: PermissionPolicy) -> Self {
|
||||
Self { policy }
|
||||
}
|
||||
|
||||
/// Check whether a tool can be executed under the current permission policy.
|
||||
/// Auto-denies when prompting is required but no prompter is provided.
|
||||
pub fn check(&self, tool_name: &str, input: &str) -> EnforcementResult {
|
||||
// When the active mode is Prompt, defer to the caller's interactive
|
||||
// prompt flow rather than hard-denying (the enforcer has no prompter).
|
||||
if self.policy.active_mode() == PermissionMode::Prompt {
|
||||
return EnforcementResult::Allowed;
|
||||
}
|
||||
|
||||
let outcome = self.policy.authorize(tool_name, input, None);
|
||||
|
||||
match outcome {
|
||||
PermissionOutcome::Allow => EnforcementResult::Allowed,
|
||||
PermissionOutcome::Deny { reason } => {
|
||||
let active_mode = self.policy.active_mode();
|
||||
let required_mode = self.policy.required_mode_for(tool_name);
|
||||
EnforcementResult::Denied {
|
||||
tool: tool_name.to_owned(),
|
||||
active_mode: active_mode.as_str().to_owned(),
|
||||
required_mode: required_mode.as_str().to_owned(),
|
||||
reason,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_allowed(&self, tool_name: &str, input: &str) -> bool {
|
||||
matches!(self.check(tool_name, input), EnforcementResult::Allowed)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn active_mode(&self) -> PermissionMode {
|
||||
self.policy.active_mode()
|
||||
}
|
||||
|
||||
/// Classify a file operation against workspace boundaries.
|
||||
pub fn check_file_write(&self, path: &str, workspace_root: &str) -> EnforcementResult {
|
||||
let mode = self.policy.active_mode();
|
||||
|
||||
match mode {
|
||||
PermissionMode::ReadOnly => EnforcementResult::Denied {
|
||||
tool: "write_file".to_owned(),
|
||||
active_mode: mode.as_str().to_owned(),
|
||||
required_mode: PermissionMode::WorkspaceWrite.as_str().to_owned(),
|
||||
reason: format!("file writes are not allowed in '{}' mode", mode.as_str()),
|
||||
},
|
||||
PermissionMode::WorkspaceWrite => {
|
||||
if is_within_workspace(path, workspace_root) {
|
||||
EnforcementResult::Allowed
|
||||
} else {
|
||||
EnforcementResult::Denied {
|
||||
tool: "write_file".to_owned(),
|
||||
active_mode: mode.as_str().to_owned(),
|
||||
required_mode: PermissionMode::DangerFullAccess.as_str().to_owned(),
|
||||
reason: format!(
|
||||
"path '{}' is outside workspace root '{}'",
|
||||
path, workspace_root
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
// Allow and DangerFullAccess permit all writes
|
||||
PermissionMode::Allow | PermissionMode::DangerFullAccess => EnforcementResult::Allowed,
|
||||
PermissionMode::Prompt => EnforcementResult::Denied {
|
||||
tool: "write_file".to_owned(),
|
||||
active_mode: mode.as_str().to_owned(),
|
||||
required_mode: PermissionMode::WorkspaceWrite.as_str().to_owned(),
|
||||
reason: "file write requires confirmation in prompt mode".to_owned(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a bash command should be allowed based on current mode.
|
||||
pub fn check_bash(&self, command: &str) -> EnforcementResult {
|
||||
let mode = self.policy.active_mode();
|
||||
|
||||
match mode {
|
||||
PermissionMode::ReadOnly => {
|
||||
if is_read_only_command(command) {
|
||||
EnforcementResult::Allowed
|
||||
} else {
|
||||
EnforcementResult::Denied {
|
||||
tool: "bash".to_owned(),
|
||||
active_mode: mode.as_str().to_owned(),
|
||||
required_mode: PermissionMode::WorkspaceWrite.as_str().to_owned(),
|
||||
reason: format!(
|
||||
"command may modify state; not allowed in '{}' mode",
|
||||
mode.as_str()
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
PermissionMode::Prompt => EnforcementResult::Denied {
|
||||
tool: "bash".to_owned(),
|
||||
active_mode: mode.as_str().to_owned(),
|
||||
required_mode: PermissionMode::DangerFullAccess.as_str().to_owned(),
|
||||
reason: "bash requires confirmation in prompt mode".to_owned(),
|
||||
},
|
||||
// WorkspaceWrite, Allow, DangerFullAccess: permit bash
|
||||
_ => EnforcementResult::Allowed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple workspace boundary check via string prefix.
|
||||
fn is_within_workspace(path: &str, workspace_root: &str) -> bool {
|
||||
let normalized = if path.starts_with('/') {
|
||||
path.to_owned()
|
||||
} else {
|
||||
format!("{workspace_root}/{path}")
|
||||
};
|
||||
|
||||
let root = if workspace_root.ends_with('/') {
|
||||
workspace_root.to_owned()
|
||||
} else {
|
||||
format!("{workspace_root}/")
|
||||
};
|
||||
|
||||
normalized.starts_with(&root) || normalized == workspace_root.trim_end_matches('/')
|
||||
}
|
||||
|
||||
/// Conservative heuristic: is this bash command read-only?
|
||||
fn is_read_only_command(command: &str) -> bool {
|
||||
let first_token = command
|
||||
.split_whitespace()
|
||||
.next()
|
||||
.unwrap_or("")
|
||||
.rsplit('/')
|
||||
.next()
|
||||
.unwrap_or("");
|
||||
|
||||
matches!(
|
||||
first_token,
|
||||
"cat"
|
||||
| "head"
|
||||
| "tail"
|
||||
| "less"
|
||||
| "more"
|
||||
| "wc"
|
||||
| "ls"
|
||||
| "find"
|
||||
| "grep"
|
||||
| "rg"
|
||||
| "awk"
|
||||
| "sed"
|
||||
| "echo"
|
||||
| "printf"
|
||||
| "which"
|
||||
| "where"
|
||||
| "whoami"
|
||||
| "pwd"
|
||||
| "env"
|
||||
| "printenv"
|
||||
| "date"
|
||||
| "cal"
|
||||
| "df"
|
||||
| "du"
|
||||
| "free"
|
||||
| "uptime"
|
||||
| "uname"
|
||||
| "file"
|
||||
| "stat"
|
||||
| "diff"
|
||||
| "sort"
|
||||
| "uniq"
|
||||
| "tr"
|
||||
| "cut"
|
||||
| "paste"
|
||||
| "tee"
|
||||
| "xargs"
|
||||
| "test"
|
||||
| "true"
|
||||
| "false"
|
||||
| "type"
|
||||
| "readlink"
|
||||
| "realpath"
|
||||
| "basename"
|
||||
| "dirname"
|
||||
| "sha256sum"
|
||||
| "md5sum"
|
||||
| "b3sum"
|
||||
| "xxd"
|
||||
| "hexdump"
|
||||
| "od"
|
||||
| "strings"
|
||||
| "tree"
|
||||
| "jq"
|
||||
| "yq"
|
||||
| "python3"
|
||||
| "python"
|
||||
| "node"
|
||||
| "ruby"
|
||||
| "cargo"
|
||||
| "rustc"
|
||||
| "git"
|
||||
| "gh"
|
||||
) && !command.contains("-i ")
|
||||
&& !command.contains("--in-place")
|
||||
&& !command.contains(" > ")
|
||||
&& !command.contains(" >> ")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_enforcer(mode: PermissionMode) -> PermissionEnforcer {
|
||||
let policy = PermissionPolicy::new(mode);
|
||||
PermissionEnforcer::new(policy)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allow_mode_permits_everything() {
|
||||
let enforcer = make_enforcer(PermissionMode::Allow);
|
||||
assert!(enforcer.is_allowed("bash", ""));
|
||||
assert!(enforcer.is_allowed("write_file", ""));
|
||||
assert!(enforcer.is_allowed("edit_file", ""));
|
||||
assert_eq!(
|
||||
enforcer.check_file_write("/outside/path", "/workspace"),
|
||||
EnforcementResult::Allowed
|
||||
);
|
||||
assert_eq!(enforcer.check_bash("rm -rf /"), EnforcementResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_only_denies_writes() {
|
||||
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
|
||||
.with_tool_requirement("read_file", PermissionMode::ReadOnly)
|
||||
.with_tool_requirement("grep_search", PermissionMode::ReadOnly)
|
||||
.with_tool_requirement("write_file", PermissionMode::WorkspaceWrite);
|
||||
|
||||
let enforcer = PermissionEnforcer::new(policy);
|
||||
assert!(enforcer.is_allowed("read_file", ""));
|
||||
assert!(enforcer.is_allowed("grep_search", ""));
|
||||
|
||||
// write_file requires WorkspaceWrite but we're in ReadOnly
|
||||
let result = enforcer.check("write_file", "");
|
||||
assert!(matches!(result, EnforcementResult::Denied { .. }));
|
||||
|
||||
let result = enforcer.check_file_write("/workspace/file.rs", "/workspace");
|
||||
assert!(matches!(result, EnforcementResult::Denied { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_only_allows_read_commands() {
|
||||
let enforcer = make_enforcer(PermissionMode::ReadOnly);
|
||||
assert_eq!(
|
||||
enforcer.check_bash("cat src/main.rs"),
|
||||
EnforcementResult::Allowed
|
||||
);
|
||||
assert_eq!(
|
||||
enforcer.check_bash("grep -r 'pattern' ."),
|
||||
EnforcementResult::Allowed
|
||||
);
|
||||
assert_eq!(enforcer.check_bash("ls -la"), EnforcementResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_only_denies_write_commands() {
|
||||
let enforcer = make_enforcer(PermissionMode::ReadOnly);
|
||||
let result = enforcer.check_bash("rm file.txt");
|
||||
assert!(matches!(result, EnforcementResult::Denied { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_write_allows_within_workspace() {
|
||||
let enforcer = make_enforcer(PermissionMode::WorkspaceWrite);
|
||||
let result = enforcer.check_file_write("/workspace/src/main.rs", "/workspace");
|
||||
assert_eq!(result, EnforcementResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_write_denies_outside_workspace() {
|
||||
let enforcer = make_enforcer(PermissionMode::WorkspaceWrite);
|
||||
let result = enforcer.check_file_write("/etc/passwd", "/workspace");
|
||||
assert!(matches!(result, EnforcementResult::Denied { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_mode_denies_without_prompter() {
|
||||
let enforcer = make_enforcer(PermissionMode::Prompt);
|
||||
let result = enforcer.check_bash("echo test");
|
||||
assert!(matches!(result, EnforcementResult::Denied { .. }));
|
||||
|
||||
let result = enforcer.check_file_write("/workspace/file.rs", "/workspace");
|
||||
assert!(matches!(result, EnforcementResult::Denied { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_boundary_check() {
|
||||
assert!(is_within_workspace("/workspace/src/main.rs", "/workspace"));
|
||||
assert!(is_within_workspace("/workspace", "/workspace"));
|
||||
assert!(!is_within_workspace("/etc/passwd", "/workspace"));
|
||||
assert!(!is_within_workspace("/workspacex/hack", "/workspace"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_only_command_heuristic() {
|
||||
assert!(is_read_only_command("cat file.txt"));
|
||||
assert!(is_read_only_command("grep pattern file"));
|
||||
assert!(is_read_only_command("git log --oneline"));
|
||||
assert!(!is_read_only_command("rm file.txt"));
|
||||
assert!(!is_read_only_command("echo test > file.txt"));
|
||||
assert!(!is_read_only_command("sed -i 's/a/b/' file"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn active_mode_returns_policy_mode() {
|
||||
// given
|
||||
let modes = [
|
||||
PermissionMode::ReadOnly,
|
||||
PermissionMode::WorkspaceWrite,
|
||||
PermissionMode::DangerFullAccess,
|
||||
PermissionMode::Prompt,
|
||||
PermissionMode::Allow,
|
||||
];
|
||||
|
||||
// when
|
||||
let active_modes: Vec<_> = modes
|
||||
.into_iter()
|
||||
.map(|mode| make_enforcer(mode).active_mode())
|
||||
.collect();
|
||||
|
||||
// then
|
||||
assert_eq!(active_modes, modes);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn danger_full_access_permits_file_writes_and_bash() {
|
||||
// given
|
||||
let enforcer = make_enforcer(PermissionMode::DangerFullAccess);
|
||||
|
||||
// when
|
||||
let file_result = enforcer.check_file_write("/outside/workspace/file.txt", "/workspace");
|
||||
let bash_result = enforcer.check_bash("rm -rf /tmp/scratch");
|
||||
|
||||
// then
|
||||
assert_eq!(file_result, EnforcementResult::Allowed);
|
||||
assert_eq!(bash_result, EnforcementResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn check_denied_payload_contains_tool_and_modes() {
|
||||
// given
|
||||
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
|
||||
.with_tool_requirement("write_file", PermissionMode::WorkspaceWrite);
|
||||
let enforcer = PermissionEnforcer::new(policy);
|
||||
|
||||
// when
|
||||
let result = enforcer.check("write_file", "{}");
|
||||
|
||||
// then
|
||||
match result {
|
||||
EnforcementResult::Denied {
|
||||
tool,
|
||||
active_mode,
|
||||
required_mode,
|
||||
reason,
|
||||
} => {
|
||||
assert_eq!(tool, "write_file");
|
||||
assert_eq!(active_mode, "read-only");
|
||||
assert_eq!(required_mode, "workspace-write");
|
||||
assert!(reason.contains("requires workspace-write permission"));
|
||||
}
|
||||
other => panic!("expected denied result, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_write_relative_path_resolved() {
|
||||
// given
|
||||
let enforcer = make_enforcer(PermissionMode::WorkspaceWrite);
|
||||
|
||||
// when
|
||||
let result = enforcer.check_file_write("src/main.rs", "/workspace");
|
||||
|
||||
// then
|
||||
assert_eq!(result, EnforcementResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_root_with_trailing_slash() {
|
||||
// given
|
||||
let enforcer = make_enforcer(PermissionMode::WorkspaceWrite);
|
||||
|
||||
// when
|
||||
let result = enforcer.check_file_write("/workspace/src/main.rs", "/workspace/");
|
||||
|
||||
// then
|
||||
assert_eq!(result, EnforcementResult::Allowed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_root_equality() {
|
||||
// given
|
||||
let root = "/workspace/";
|
||||
|
||||
// when
|
||||
let equal_to_root = is_within_workspace("/workspace", root);
|
||||
|
||||
// then
|
||||
assert!(equal_to_root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bash_heuristic_full_path_prefix() {
|
||||
// given
|
||||
let full_path_command = "/usr/bin/cat Cargo.toml";
|
||||
let git_path_command = "/usr/local/bin/git status";
|
||||
|
||||
// when
|
||||
let cat_result = is_read_only_command(full_path_command);
|
||||
let git_result = is_read_only_command(git_path_command);
|
||||
|
||||
// then
|
||||
assert!(cat_result);
|
||||
assert!(git_result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bash_heuristic_redirects_block_read_only_commands() {
|
||||
// given
|
||||
let overwrite = "cat Cargo.toml > out.txt";
|
||||
let append = "echo test >> out.txt";
|
||||
|
||||
// when
|
||||
let overwrite_result = is_read_only_command(overwrite);
|
||||
let append_result = is_read_only_command(append);
|
||||
|
||||
// then
|
||||
assert!(!overwrite_result);
|
||||
assert!(!append_result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bash_heuristic_in_place_flag_blocks() {
|
||||
// given
|
||||
let interactive_python = "python -i script.py";
|
||||
let in_place_sed = "sed --in-place 's/a/b/' file.txt";
|
||||
|
||||
// when
|
||||
let interactive_result = is_read_only_command(interactive_python);
|
||||
let in_place_result = is_read_only_command(in_place_sed);
|
||||
|
||||
// then
|
||||
assert!(!interactive_result);
|
||||
assert!(!in_place_result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bash_heuristic_empty_command() {
|
||||
// given
|
||||
let empty = "";
|
||||
let whitespace = " ";
|
||||
|
||||
// when
|
||||
let empty_result = is_read_only_command(empty);
|
||||
let whitespace_result = is_read_only_command(whitespace);
|
||||
|
||||
// then
|
||||
assert!(!empty_result);
|
||||
assert!(!whitespace_result);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompt_mode_check_bash_denied_payload_fields() {
|
||||
// given
|
||||
let enforcer = make_enforcer(PermissionMode::Prompt);
|
||||
|
||||
// when
|
||||
let result = enforcer.check_bash("git status");
|
||||
|
||||
// then
|
||||
match result {
|
||||
EnforcementResult::Denied {
|
||||
tool,
|
||||
active_mode,
|
||||
required_mode,
|
||||
reason,
|
||||
} => {
|
||||
assert_eq!(tool, "bash");
|
||||
assert_eq!(active_mode, "prompt");
|
||||
assert_eq!(required_mode, "danger-full-access");
|
||||
assert_eq!(reason, "bash requires confirmation in prompt mode");
|
||||
}
|
||||
other => panic!("expected denied result, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_only_check_file_write_denied_payload() {
|
||||
// given
|
||||
let enforcer = make_enforcer(PermissionMode::ReadOnly);
|
||||
|
||||
// when
|
||||
let result = enforcer.check_file_write("/workspace/file.txt", "/workspace");
|
||||
|
||||
// then
|
||||
match result {
|
||||
EnforcementResult::Denied {
|
||||
tool,
|
||||
active_mode,
|
||||
required_mode,
|
||||
reason,
|
||||
} => {
|
||||
assert_eq!(tool, "write_file");
|
||||
assert_eq!(active_mode, "read-only");
|
||||
assert_eq!(required_mode, "workspace-write");
|
||||
assert!(reason.contains("file writes are not allowed"));
|
||||
}
|
||||
other => panic!("expected denied result, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,683 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::config::RuntimePermissionRuleConfig;
|
||||
|
||||
/// Permission level assigned to a tool invocation or runtime session.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
pub enum PermissionMode {
|
||||
ReadOnly,
|
||||
WorkspaceWrite,
|
||||
DangerFullAccess,
|
||||
Prompt,
|
||||
Allow,
|
||||
}
|
||||
|
||||
impl PermissionMode {
|
||||
#[must_use]
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::ReadOnly => "read-only",
|
||||
Self::WorkspaceWrite => "workspace-write",
|
||||
Self::DangerFullAccess => "danger-full-access",
|
||||
Self::Prompt => "prompt",
|
||||
Self::Allow => "allow",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Hook-provided override applied before standard permission evaluation.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum PermissionOverride {
|
||||
Allow,
|
||||
Deny,
|
||||
Ask,
|
||||
}
|
||||
|
||||
/// Additional permission context supplied by hooks or higher-level orchestration.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct PermissionContext {
|
||||
override_decision: Option<PermissionOverride>,
|
||||
override_reason: Option<String>,
|
||||
}
|
||||
|
||||
impl PermissionContext {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
override_decision: Option<PermissionOverride>,
|
||||
override_reason: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
override_decision,
|
||||
override_reason,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn override_decision(&self) -> Option<PermissionOverride> {
|
||||
self.override_decision
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn override_reason(&self) -> Option<&str> {
|
||||
self.override_reason.as_deref()
|
||||
}
|
||||
}
|
||||
|
||||
/// Full authorization request presented to a permission prompt.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PermissionRequest {
|
||||
pub tool_name: String,
|
||||
pub input: String,
|
||||
pub current_mode: PermissionMode,
|
||||
pub required_mode: PermissionMode,
|
||||
pub reason: Option<String>,
|
||||
}
|
||||
|
||||
/// User-facing decision returned by a [`PermissionPrompter`].
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PermissionPromptDecision {
|
||||
Allow,
|
||||
Deny { reason: String },
|
||||
}
|
||||
|
||||
/// Prompting interface used when policy requires interactive approval.
|
||||
pub trait PermissionPrompter {
|
||||
fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision;
|
||||
}
|
||||
|
||||
/// Final authorization result after evaluating static rules and prompts.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PermissionOutcome {
|
||||
Allow,
|
||||
Deny { reason: String },
|
||||
}
|
||||
|
||||
/// Evaluates permission mode requirements plus allow/deny/ask rules.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PermissionPolicy {
|
||||
active_mode: PermissionMode,
|
||||
tool_requirements: BTreeMap<String, PermissionMode>,
|
||||
allow_rules: Vec<PermissionRule>,
|
||||
deny_rules: Vec<PermissionRule>,
|
||||
ask_rules: Vec<PermissionRule>,
|
||||
}
|
||||
|
||||
impl PermissionPolicy {
|
||||
#[must_use]
|
||||
pub fn new(active_mode: PermissionMode) -> Self {
|
||||
Self {
|
||||
active_mode,
|
||||
tool_requirements: BTreeMap::new(),
|
||||
allow_rules: Vec::new(),
|
||||
deny_rules: Vec::new(),
|
||||
ask_rules: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_tool_requirement(
|
||||
mut self,
|
||||
tool_name: impl Into<String>,
|
||||
required_mode: PermissionMode,
|
||||
) -> Self {
|
||||
self.tool_requirements
|
||||
.insert(tool_name.into(), required_mode);
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_permission_rules(mut self, config: &RuntimePermissionRuleConfig) -> Self {
|
||||
self.allow_rules = config
|
||||
.allow()
|
||||
.iter()
|
||||
.map(|rule| PermissionRule::parse(rule))
|
||||
.collect();
|
||||
self.deny_rules = config
|
||||
.deny()
|
||||
.iter()
|
||||
.map(|rule| PermissionRule::parse(rule))
|
||||
.collect();
|
||||
self.ask_rules = config
|
||||
.ask()
|
||||
.iter()
|
||||
.map(|rule| PermissionRule::parse(rule))
|
||||
.collect();
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn active_mode(&self) -> PermissionMode {
|
||||
self.active_mode
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn required_mode_for(&self, tool_name: &str) -> PermissionMode {
|
||||
self.tool_requirements
|
||||
.get(tool_name)
|
||||
.copied()
|
||||
.unwrap_or(PermissionMode::DangerFullAccess)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn authorize(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
input: &str,
|
||||
prompter: Option<&mut dyn PermissionPrompter>,
|
||||
) -> PermissionOutcome {
|
||||
self.authorize_with_context(tool_name, input, &PermissionContext::default(), prompter)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
#[allow(clippy::too_many_lines)]
|
||||
pub fn authorize_with_context(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
input: &str,
|
||||
context: &PermissionContext,
|
||||
prompter: Option<&mut dyn PermissionPrompter>,
|
||||
) -> PermissionOutcome {
|
||||
if let Some(rule) = Self::find_matching_rule(&self.deny_rules, tool_name, input) {
|
||||
return PermissionOutcome::Deny {
|
||||
reason: format!(
|
||||
"Permission to use {tool_name} has been denied by rule '{}'",
|
||||
rule.raw
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
let current_mode = self.active_mode();
|
||||
let required_mode = self.required_mode_for(tool_name);
|
||||
let ask_rule = Self::find_matching_rule(&self.ask_rules, tool_name, input);
|
||||
let allow_rule = Self::find_matching_rule(&self.allow_rules, tool_name, input);
|
||||
|
||||
match context.override_decision() {
|
||||
Some(PermissionOverride::Deny) => {
|
||||
return PermissionOutcome::Deny {
|
||||
reason: context.override_reason().map_or_else(
|
||||
|| format!("tool '{tool_name}' denied by hook"),
|
||||
ToOwned::to_owned,
|
||||
),
|
||||
};
|
||||
}
|
||||
Some(PermissionOverride::Ask) => {
|
||||
let reason = context.override_reason().map_or_else(
|
||||
|| format!("tool '{tool_name}' requires approval due to hook guidance"),
|
||||
ToOwned::to_owned,
|
||||
);
|
||||
return Self::prompt_or_deny(
|
||||
tool_name,
|
||||
input,
|
||||
current_mode,
|
||||
required_mode,
|
||||
Some(reason),
|
||||
prompter,
|
||||
);
|
||||
}
|
||||
Some(PermissionOverride::Allow) => {
|
||||
if let Some(rule) = ask_rule {
|
||||
let reason = format!(
|
||||
"tool '{tool_name}' requires approval due to ask rule '{}'",
|
||||
rule.raw
|
||||
);
|
||||
return Self::prompt_or_deny(
|
||||
tool_name,
|
||||
input,
|
||||
current_mode,
|
||||
required_mode,
|
||||
Some(reason),
|
||||
prompter,
|
||||
);
|
||||
}
|
||||
if allow_rule.is_some()
|
||||
|| current_mode == PermissionMode::Allow
|
||||
|| current_mode >= required_mode
|
||||
{
|
||||
return PermissionOutcome::Allow;
|
||||
}
|
||||
}
|
||||
None => {}
|
||||
}
|
||||
|
||||
if let Some(rule) = ask_rule {
|
||||
let reason = format!(
|
||||
"tool '{tool_name}' requires approval due to ask rule '{}'",
|
||||
rule.raw
|
||||
);
|
||||
return Self::prompt_or_deny(
|
||||
tool_name,
|
||||
input,
|
||||
current_mode,
|
||||
required_mode,
|
||||
Some(reason),
|
||||
prompter,
|
||||
);
|
||||
}
|
||||
|
||||
if allow_rule.is_some()
|
||||
|| current_mode == PermissionMode::Allow
|
||||
|| current_mode >= required_mode
|
||||
{
|
||||
return PermissionOutcome::Allow;
|
||||
}
|
||||
|
||||
if current_mode == PermissionMode::Prompt
|
||||
|| (current_mode == PermissionMode::WorkspaceWrite
|
||||
&& required_mode == PermissionMode::DangerFullAccess)
|
||||
{
|
||||
let reason = Some(format!(
|
||||
"tool '{tool_name}' requires approval to escalate from {} to {}",
|
||||
current_mode.as_str(),
|
||||
required_mode.as_str()
|
||||
));
|
||||
return Self::prompt_or_deny(
|
||||
tool_name,
|
||||
input,
|
||||
current_mode,
|
||||
required_mode,
|
||||
reason,
|
||||
prompter,
|
||||
);
|
||||
}
|
||||
|
||||
PermissionOutcome::Deny {
|
||||
reason: format!(
|
||||
"tool '{tool_name}' requires {} permission; current mode is {}",
|
||||
required_mode.as_str(),
|
||||
current_mode.as_str()
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn prompt_or_deny(
|
||||
tool_name: &str,
|
||||
input: &str,
|
||||
current_mode: PermissionMode,
|
||||
required_mode: PermissionMode,
|
||||
reason: Option<String>,
|
||||
mut prompter: Option<&mut dyn PermissionPrompter>,
|
||||
) -> PermissionOutcome {
|
||||
let request = PermissionRequest {
|
||||
tool_name: tool_name.to_string(),
|
||||
input: input.to_string(),
|
||||
current_mode,
|
||||
required_mode,
|
||||
reason: reason.clone(),
|
||||
};
|
||||
|
||||
match prompter.as_mut() {
|
||||
Some(prompter) => match prompter.decide(&request) {
|
||||
PermissionPromptDecision::Allow => PermissionOutcome::Allow,
|
||||
PermissionPromptDecision::Deny { reason } => PermissionOutcome::Deny { reason },
|
||||
},
|
||||
None => PermissionOutcome::Deny {
|
||||
reason: reason.unwrap_or_else(|| {
|
||||
format!(
|
||||
"tool '{tool_name}' requires approval to run while mode is {}",
|
||||
current_mode.as_str()
|
||||
)
|
||||
}),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn find_matching_rule<'a>(
|
||||
rules: &'a [PermissionRule],
|
||||
tool_name: &str,
|
||||
input: &str,
|
||||
) -> Option<&'a PermissionRule> {
|
||||
rules.iter().find(|rule| rule.matches(tool_name, input))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct PermissionRule {
|
||||
raw: String,
|
||||
tool_name: String,
|
||||
matcher: PermissionRuleMatcher,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
enum PermissionRuleMatcher {
|
||||
Any,
|
||||
Exact(String),
|
||||
Prefix(String),
|
||||
}
|
||||
|
||||
impl PermissionRule {
|
||||
fn parse(raw: &str) -> Self {
|
||||
let trimmed = raw.trim();
|
||||
let open = find_first_unescaped(trimmed, '(');
|
||||
let close = find_last_unescaped(trimmed, ')');
|
||||
|
||||
if let (Some(open), Some(close)) = (open, close) {
|
||||
if close == trimmed.len() - 1 && open < close {
|
||||
let tool_name = trimmed[..open].trim();
|
||||
let content = &trimmed[open + 1..close];
|
||||
if !tool_name.is_empty() {
|
||||
let matcher = parse_rule_matcher(content);
|
||||
return Self {
|
||||
raw: trimmed.to_string(),
|
||||
tool_name: tool_name.to_string(),
|
||||
matcher,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Self {
|
||||
raw: trimmed.to_string(),
|
||||
tool_name: trimmed.to_string(),
|
||||
matcher: PermissionRuleMatcher::Any,
|
||||
}
|
||||
}
|
||||
|
||||
fn matches(&self, tool_name: &str, input: &str) -> bool {
|
||||
if self.tool_name != tool_name {
|
||||
return false;
|
||||
}
|
||||
|
||||
match &self.matcher {
|
||||
PermissionRuleMatcher::Any => true,
|
||||
PermissionRuleMatcher::Exact(expected) => {
|
||||
extract_permission_subject(input).is_some_and(|candidate| candidate == *expected)
|
||||
}
|
||||
PermissionRuleMatcher::Prefix(prefix) => extract_permission_subject(input)
|
||||
.is_some_and(|candidate| candidate.starts_with(prefix)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_rule_matcher(content: &str) -> PermissionRuleMatcher {
|
||||
let unescaped = unescape_rule_content(content.trim());
|
||||
if unescaped.is_empty() || unescaped == "*" {
|
||||
PermissionRuleMatcher::Any
|
||||
} else if let Some(prefix) = unescaped.strip_suffix(":*") {
|
||||
PermissionRuleMatcher::Prefix(prefix.to_string())
|
||||
} else {
|
||||
PermissionRuleMatcher::Exact(unescaped)
|
||||
}
|
||||
}
|
||||
|
||||
fn unescape_rule_content(content: &str) -> String {
|
||||
content
|
||||
.replace(r"\(", "(")
|
||||
.replace(r"\)", ")")
|
||||
.replace(r"\\", r"\")
|
||||
}
|
||||
|
||||
fn find_first_unescaped(value: &str, needle: char) -> Option<usize> {
|
||||
let mut escaped = false;
|
||||
for (idx, ch) in value.char_indices() {
|
||||
if ch == '\\' {
|
||||
escaped = !escaped;
|
||||
continue;
|
||||
}
|
||||
if ch == needle && !escaped {
|
||||
return Some(idx);
|
||||
}
|
||||
escaped = false;
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn find_last_unescaped(value: &str, needle: char) -> Option<usize> {
|
||||
let chars = value.char_indices().collect::<Vec<_>>();
|
||||
for (pos, (idx, ch)) in chars.iter().enumerate().rev() {
|
||||
if *ch != needle {
|
||||
continue;
|
||||
}
|
||||
let mut backslashes = 0;
|
||||
for (_, prev) in chars[..pos].iter().rev() {
|
||||
if *prev == '\\' {
|
||||
backslashes += 1;
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
if backslashes % 2 == 0 {
|
||||
return Some(*idx);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn extract_permission_subject(input: &str) -> Option<String> {
|
||||
let parsed = serde_json::from_str::<Value>(input).ok();
|
||||
if let Some(Value::Object(object)) = parsed {
|
||||
for key in [
|
||||
"command",
|
||||
"path",
|
||||
"file_path",
|
||||
"filePath",
|
||||
"notebook_path",
|
||||
"notebookPath",
|
||||
"url",
|
||||
"pattern",
|
||||
"code",
|
||||
"message",
|
||||
] {
|
||||
if let Some(value) = object.get(key).and_then(Value::as_str) {
|
||||
return Some(value.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
(!input.trim().is_empty()).then(|| input.to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
PermissionContext, PermissionMode, PermissionOutcome, PermissionOverride, PermissionPolicy,
|
||||
PermissionPromptDecision, PermissionPrompter, PermissionRequest,
|
||||
};
|
||||
use crate::config::RuntimePermissionRuleConfig;
|
||||
|
||||
struct RecordingPrompter {
|
||||
seen: Vec<PermissionRequest>,
|
||||
allow: bool,
|
||||
}
|
||||
|
||||
impl PermissionPrompter for RecordingPrompter {
|
||||
fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision {
|
||||
self.seen.push(request.clone());
|
||||
if self.allow {
|
||||
PermissionPromptDecision::Allow
|
||||
} else {
|
||||
PermissionPromptDecision::Deny {
|
||||
reason: "not now".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn allows_tools_when_active_mode_meets_requirement() {
|
||||
let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
|
||||
.with_tool_requirement("read_file", PermissionMode::ReadOnly)
|
||||
.with_tool_requirement("write_file", PermissionMode::WorkspaceWrite);
|
||||
|
||||
assert_eq!(
|
||||
policy.authorize("read_file", "{}", None),
|
||||
PermissionOutcome::Allow
|
||||
);
|
||||
assert_eq!(
|
||||
policy.authorize("write_file", "{}", None),
|
||||
PermissionOutcome::Allow
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn denies_read_only_escalations_without_prompt() {
|
||||
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
|
||||
.with_tool_requirement("write_file", PermissionMode::WorkspaceWrite)
|
||||
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
|
||||
|
||||
assert!(matches!(
|
||||
policy.authorize("write_file", "{}", None),
|
||||
PermissionOutcome::Deny { reason } if reason.contains("requires workspace-write permission")
|
||||
));
|
||||
assert!(matches!(
|
||||
policy.authorize("bash", "{}", None),
|
||||
PermissionOutcome::Deny { reason } if reason.contains("requires danger-full-access permission")
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prompts_for_workspace_write_to_danger_full_access_escalation() {
|
||||
let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
|
||||
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
|
||||
let mut prompter = RecordingPrompter {
|
||||
seen: Vec::new(),
|
||||
allow: true,
|
||||
};
|
||||
|
||||
let outcome = policy.authorize("bash", "echo hi", Some(&mut prompter));
|
||||
|
||||
assert_eq!(outcome, PermissionOutcome::Allow);
|
||||
assert_eq!(prompter.seen.len(), 1);
|
||||
assert_eq!(prompter.seen[0].tool_name, "bash");
|
||||
assert_eq!(
|
||||
prompter.seen[0].current_mode,
|
||||
PermissionMode::WorkspaceWrite
|
||||
);
|
||||
assert_eq!(
|
||||
prompter.seen[0].required_mode,
|
||||
PermissionMode::DangerFullAccess
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn honors_prompt_rejection_reason() {
|
||||
let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
|
||||
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
|
||||
let mut prompter = RecordingPrompter {
|
||||
seen: Vec::new(),
|
||||
allow: false,
|
||||
};
|
||||
|
||||
assert!(matches!(
|
||||
policy.authorize("bash", "echo hi", Some(&mut prompter)),
|
||||
PermissionOutcome::Deny { reason } if reason == "not now"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn applies_rule_based_denials_and_allows() {
|
||||
let rules = RuntimePermissionRuleConfig::new(
|
||||
vec!["bash(git:*)".to_string()],
|
||||
vec!["bash(rm -rf:*)".to_string()],
|
||||
Vec::new(),
|
||||
);
|
||||
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
|
||||
.with_tool_requirement("bash", PermissionMode::DangerFullAccess)
|
||||
.with_permission_rules(&rules);
|
||||
|
||||
assert_eq!(
|
||||
policy.authorize("bash", r#"{"command":"git status"}"#, None),
|
||||
PermissionOutcome::Allow
|
||||
);
|
||||
assert!(matches!(
|
||||
policy.authorize("bash", r#"{"command":"rm -rf /tmp/x"}"#, None),
|
||||
PermissionOutcome::Deny { reason } if reason.contains("denied by rule")
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ask_rules_force_prompt_even_when_mode_allows() {
|
||||
let rules = RuntimePermissionRuleConfig::new(
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
vec!["bash(git:*)".to_string()],
|
||||
);
|
||||
let policy = PermissionPolicy::new(PermissionMode::DangerFullAccess)
|
||||
.with_tool_requirement("bash", PermissionMode::DangerFullAccess)
|
||||
.with_permission_rules(&rules);
|
||||
let mut prompter = RecordingPrompter {
|
||||
seen: Vec::new(),
|
||||
allow: true,
|
||||
};
|
||||
|
||||
let outcome = policy.authorize("bash", r#"{"command":"git status"}"#, Some(&mut prompter));
|
||||
|
||||
assert_eq!(outcome, PermissionOutcome::Allow);
|
||||
assert_eq!(prompter.seen.len(), 1);
|
||||
assert!(prompter.seen[0]
|
||||
.reason
|
||||
.as_deref()
|
||||
.is_some_and(|reason| reason.contains("ask rule")));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hook_allow_still_respects_ask_rules() {
|
||||
let rules = RuntimePermissionRuleConfig::new(
|
||||
Vec::new(),
|
||||
Vec::new(),
|
||||
vec!["bash(git:*)".to_string()],
|
||||
);
|
||||
let policy = PermissionPolicy::new(PermissionMode::ReadOnly)
|
||||
.with_tool_requirement("bash", PermissionMode::DangerFullAccess)
|
||||
.with_permission_rules(&rules);
|
||||
let context = PermissionContext::new(
|
||||
Some(PermissionOverride::Allow),
|
||||
Some("hook approved".to_string()),
|
||||
);
|
||||
let mut prompter = RecordingPrompter {
|
||||
seen: Vec::new(),
|
||||
allow: true,
|
||||
};
|
||||
|
||||
let outcome = policy.authorize_with_context(
|
||||
"bash",
|
||||
r#"{"command":"git status"}"#,
|
||||
&context,
|
||||
Some(&mut prompter),
|
||||
);
|
||||
|
||||
assert_eq!(outcome, PermissionOutcome::Allow);
|
||||
assert_eq!(prompter.seen.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hook_deny_short_circuits_permission_flow() {
|
||||
let policy = PermissionPolicy::new(PermissionMode::DangerFullAccess)
|
||||
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
|
||||
let context = PermissionContext::new(
|
||||
Some(PermissionOverride::Deny),
|
||||
Some("blocked by hook".to_string()),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
policy.authorize_with_context("bash", "{}", &context, None),
|
||||
PermissionOutcome::Deny {
|
||||
reason: "blocked by hook".to_string(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hook_ask_forces_prompt() {
|
||||
let policy = PermissionPolicy::new(PermissionMode::DangerFullAccess)
|
||||
.with_tool_requirement("bash", PermissionMode::DangerFullAccess);
|
||||
let context = PermissionContext::new(
|
||||
Some(PermissionOverride::Ask),
|
||||
Some("hook requested confirmation".to_string()),
|
||||
);
|
||||
let mut prompter = RecordingPrompter {
|
||||
seen: Vec::new(),
|
||||
allow: true,
|
||||
};
|
||||
|
||||
let outcome = policy.authorize_with_context("bash", "{}", &context, Some(&mut prompter));
|
||||
|
||||
assert_eq!(outcome, PermissionOutcome::Allow);
|
||||
assert_eq!(prompter.seen.len(), 1);
|
||||
assert_eq!(
|
||||
prompter.seen[0].reason.as_deref(),
|
||||
Some("hook requested confirmation")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,533 @@
|
||||
#![allow(clippy::redundant_closure_for_method_calls)]
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::config::RuntimePluginConfig;
|
||||
use crate::mcp_tool_bridge::{McpResourceInfo, McpToolInfo};
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
pub type ToolInfo = McpToolInfo;
|
||||
pub type ResourceInfo = McpResourceInfo;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum ServerStatus {
|
||||
Healthy,
|
||||
Degraded,
|
||||
Failed,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for ServerStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Healthy => write!(f, "healthy"),
|
||||
Self::Degraded => write!(f, "degraded"),
|
||||
Self::Failed => write!(f, "failed"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ServerHealth {
|
||||
pub server_name: String,
|
||||
pub status: ServerStatus,
|
||||
pub capabilities: Vec<String>,
|
||||
pub last_error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case", tag = "state")]
|
||||
pub enum PluginState {
|
||||
Unconfigured,
|
||||
Validated,
|
||||
Starting,
|
||||
Healthy,
|
||||
Degraded {
|
||||
healthy_servers: Vec<String>,
|
||||
failed_servers: Vec<ServerHealth>,
|
||||
},
|
||||
Failed {
|
||||
reason: String,
|
||||
},
|
||||
ShuttingDown,
|
||||
Stopped,
|
||||
}
|
||||
|
||||
impl PluginState {
|
||||
#[must_use]
|
||||
pub fn from_servers(servers: &[ServerHealth]) -> Self {
|
||||
if servers.is_empty() {
|
||||
return Self::Failed {
|
||||
reason: "no servers available".to_string(),
|
||||
};
|
||||
}
|
||||
|
||||
let healthy_servers = servers
|
||||
.iter()
|
||||
.filter(|server| server.status != ServerStatus::Failed)
|
||||
.map(|server| server.server_name.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let failed_servers = servers
|
||||
.iter()
|
||||
.filter(|server| server.status == ServerStatus::Failed)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let has_degraded_server = servers
|
||||
.iter()
|
||||
.any(|server| server.status == ServerStatus::Degraded);
|
||||
|
||||
if failed_servers.is_empty() && !has_degraded_server {
|
||||
Self::Healthy
|
||||
} else if healthy_servers.is_empty() {
|
||||
Self::Failed {
|
||||
reason: format!("all {} servers failed", failed_servers.len()),
|
||||
}
|
||||
} else {
|
||||
Self::Degraded {
|
||||
healthy_servers,
|
||||
failed_servers,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for PluginState {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Unconfigured => write!(f, "unconfigured"),
|
||||
Self::Validated => write!(f, "validated"),
|
||||
Self::Starting => write!(f, "starting"),
|
||||
Self::Healthy => write!(f, "healthy"),
|
||||
Self::Degraded { .. } => write!(f, "degraded"),
|
||||
Self::Failed { .. } => write!(f, "failed"),
|
||||
Self::ShuttingDown => write!(f, "shutting_down"),
|
||||
Self::Stopped => write!(f, "stopped"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct PluginHealthcheck {
|
||||
pub plugin_name: String,
|
||||
pub state: PluginState,
|
||||
pub servers: Vec<ServerHealth>,
|
||||
pub last_check: u64,
|
||||
}
|
||||
|
||||
impl PluginHealthcheck {
|
||||
#[must_use]
|
||||
pub fn new(plugin_name: impl Into<String>, servers: Vec<ServerHealth>) -> Self {
|
||||
let state = PluginState::from_servers(&servers);
|
||||
Self {
|
||||
plugin_name: plugin_name.into(),
|
||||
state,
|
||||
servers,
|
||||
last_check: now_secs(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn degraded_mode(&self, discovery: &DiscoveryResult) -> Option<DegradedMode> {
|
||||
match &self.state {
|
||||
PluginState::Degraded {
|
||||
healthy_servers,
|
||||
failed_servers,
|
||||
} => Some(DegradedMode {
|
||||
available_tools: discovery
|
||||
.tools
|
||||
.iter()
|
||||
.map(|tool| tool.name.clone())
|
||||
.collect(),
|
||||
unavailable_tools: failed_servers
|
||||
.iter()
|
||||
.flat_map(|server| server.capabilities.iter().cloned())
|
||||
.collect(),
|
||||
reason: format!(
|
||||
"{} servers healthy, {} servers failed",
|
||||
healthy_servers.len(),
|
||||
failed_servers.len()
|
||||
),
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct DiscoveryResult {
|
||||
pub tools: Vec<ToolInfo>,
|
||||
pub resources: Vec<ResourceInfo>,
|
||||
pub partial: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct DegradedMode {
|
||||
pub available_tools: Vec<String>,
|
||||
pub unavailable_tools: Vec<String>,
|
||||
pub reason: String,
|
||||
}
|
||||
|
||||
impl DegradedMode {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
available_tools: Vec<String>,
|
||||
unavailable_tools: Vec<String>,
|
||||
reason: impl Into<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
available_tools,
|
||||
unavailable_tools,
|
||||
reason: reason.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum PluginLifecycleEvent {
|
||||
ConfigValidated,
|
||||
StartupHealthy,
|
||||
StartupDegraded,
|
||||
StartupFailed,
|
||||
Shutdown,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for PluginLifecycleEvent {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::ConfigValidated => write!(f, "config_validated"),
|
||||
Self::StartupHealthy => write!(f, "startup_healthy"),
|
||||
Self::StartupDegraded => write!(f, "startup_degraded"),
|
||||
Self::StartupFailed => write!(f, "startup_failed"),
|
||||
Self::Shutdown => write!(f, "shutdown"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub trait PluginLifecycle {
|
||||
fn validate_config(&self, config: &RuntimePluginConfig) -> Result<(), String>;
|
||||
fn healthcheck(&self) -> PluginHealthcheck;
|
||||
fn discover(&self) -> DiscoveryResult;
|
||||
fn shutdown(&mut self) -> Result<(), String>;
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MockPluginLifecycle {
|
||||
plugin_name: String,
|
||||
valid_config: bool,
|
||||
healthcheck: PluginHealthcheck,
|
||||
discovery: DiscoveryResult,
|
||||
shutdown_error: Option<String>,
|
||||
shutdown_called: bool,
|
||||
}
|
||||
|
||||
impl MockPluginLifecycle {
|
||||
fn new(
|
||||
plugin_name: &str,
|
||||
valid_config: bool,
|
||||
servers: Vec<ServerHealth>,
|
||||
discovery: DiscoveryResult,
|
||||
shutdown_error: Option<String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
plugin_name: plugin_name.to_string(),
|
||||
valid_config,
|
||||
healthcheck: PluginHealthcheck::new(plugin_name, servers),
|
||||
discovery,
|
||||
shutdown_error,
|
||||
shutdown_called: false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PluginLifecycle for MockPluginLifecycle {
|
||||
fn validate_config(&self, _config: &RuntimePluginConfig) -> Result<(), String> {
|
||||
if self.valid_config {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!(
|
||||
"plugin `{}` failed configuration validation",
|
||||
self.plugin_name
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn healthcheck(&self) -> PluginHealthcheck {
|
||||
if self.shutdown_called {
|
||||
PluginHealthcheck {
|
||||
plugin_name: self.plugin_name.clone(),
|
||||
state: PluginState::Stopped,
|
||||
servers: self.healthcheck.servers.clone(),
|
||||
last_check: now_secs(),
|
||||
}
|
||||
} else {
|
||||
self.healthcheck.clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn discover(&self) -> DiscoveryResult {
|
||||
self.discovery.clone()
|
||||
}
|
||||
|
||||
fn shutdown(&mut self) -> Result<(), String> {
|
||||
if let Some(error) = &self.shutdown_error {
|
||||
return Err(error.clone());
|
||||
}
|
||||
|
||||
self.shutdown_called = true;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn healthy_server(name: &str, capabilities: &[&str]) -> ServerHealth {
|
||||
ServerHealth {
|
||||
server_name: name.to_string(),
|
||||
status: ServerStatus::Healthy,
|
||||
capabilities: capabilities
|
||||
.iter()
|
||||
.map(|capability| capability.to_string())
|
||||
.collect(),
|
||||
last_error: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn failed_server(name: &str, capabilities: &[&str], error: &str) -> ServerHealth {
|
||||
ServerHealth {
|
||||
server_name: name.to_string(),
|
||||
status: ServerStatus::Failed,
|
||||
capabilities: capabilities
|
||||
.iter()
|
||||
.map(|capability| capability.to_string())
|
||||
.collect(),
|
||||
last_error: Some(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn degraded_server(name: &str, capabilities: &[&str], error: &str) -> ServerHealth {
|
||||
ServerHealth {
|
||||
server_name: name.to_string(),
|
||||
status: ServerStatus::Degraded,
|
||||
capabilities: capabilities
|
||||
.iter()
|
||||
.map(|capability| capability.to_string())
|
||||
.collect(),
|
||||
last_error: Some(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn tool(name: &str) -> ToolInfo {
|
||||
ToolInfo {
|
||||
name: name.to_string(),
|
||||
description: Some(format!("{name} tool")),
|
||||
input_schema: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn resource(name: &str, uri: &str) -> ResourceInfo {
|
||||
ResourceInfo {
|
||||
uri: uri.to_string(),
|
||||
name: name.to_string(),
|
||||
description: Some(format!("{name} resource")),
|
||||
mime_type: Some("application/json".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn full_lifecycle_happy_path() {
|
||||
// given
|
||||
let mut lifecycle = MockPluginLifecycle::new(
|
||||
"healthy-plugin",
|
||||
true,
|
||||
vec![
|
||||
healthy_server("alpha", &["search", "read"]),
|
||||
healthy_server("beta", &["write"]),
|
||||
],
|
||||
DiscoveryResult {
|
||||
tools: vec![tool("search"), tool("read"), tool("write")],
|
||||
resources: vec![resource("docs", "file:///docs")],
|
||||
partial: false,
|
||||
},
|
||||
None,
|
||||
);
|
||||
let config = RuntimePluginConfig::default();
|
||||
|
||||
// when
|
||||
let validation = lifecycle.validate_config(&config);
|
||||
let healthcheck = lifecycle.healthcheck();
|
||||
let discovery = lifecycle.discover();
|
||||
let shutdown = lifecycle.shutdown();
|
||||
let post_shutdown = lifecycle.healthcheck();
|
||||
|
||||
// then
|
||||
assert_eq!(validation, Ok(()));
|
||||
assert_eq!(healthcheck.state, PluginState::Healthy);
|
||||
assert_eq!(healthcheck.plugin_name, "healthy-plugin");
|
||||
assert_eq!(discovery.tools.len(), 3);
|
||||
assert_eq!(discovery.resources.len(), 1);
|
||||
assert!(!discovery.partial);
|
||||
assert_eq!(shutdown, Ok(()));
|
||||
assert_eq!(post_shutdown.state, PluginState::Stopped);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn degraded_startup_when_one_of_three_servers_fails() {
|
||||
// given
|
||||
let lifecycle = MockPluginLifecycle::new(
|
||||
"degraded-plugin",
|
||||
true,
|
||||
vec![
|
||||
healthy_server("alpha", &["search"]),
|
||||
failed_server("beta", &["write"], "connection refused"),
|
||||
healthy_server("gamma", &["read"]),
|
||||
],
|
||||
DiscoveryResult {
|
||||
tools: vec![tool("search"), tool("read")],
|
||||
resources: vec![resource("alpha-docs", "file:///alpha")],
|
||||
partial: true,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let healthcheck = lifecycle.healthcheck();
|
||||
let discovery = lifecycle.discover();
|
||||
let degraded_mode = healthcheck
|
||||
.degraded_mode(&discovery)
|
||||
.expect("degraded startup should expose degraded mode");
|
||||
|
||||
// then
|
||||
match healthcheck.state {
|
||||
PluginState::Degraded {
|
||||
healthy_servers,
|
||||
failed_servers,
|
||||
} => {
|
||||
assert_eq!(
|
||||
healthy_servers,
|
||||
vec!["alpha".to_string(), "gamma".to_string()]
|
||||
);
|
||||
assert_eq!(failed_servers.len(), 1);
|
||||
assert_eq!(failed_servers[0].server_name, "beta");
|
||||
assert_eq!(
|
||||
failed_servers[0].last_error.as_deref(),
|
||||
Some("connection refused")
|
||||
);
|
||||
}
|
||||
other => panic!("expected degraded state, got {other:?}"),
|
||||
}
|
||||
assert!(discovery.partial);
|
||||
assert_eq!(
|
||||
degraded_mode.available_tools,
|
||||
vec!["search".to_string(), "read".to_string()]
|
||||
);
|
||||
assert_eq!(degraded_mode.unavailable_tools, vec!["write".to_string()]);
|
||||
assert_eq!(degraded_mode.reason, "2 servers healthy, 1 servers failed");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn degraded_server_status_keeps_server_usable() {
|
||||
// given
|
||||
let lifecycle = MockPluginLifecycle::new(
|
||||
"soft-degraded-plugin",
|
||||
true,
|
||||
vec![
|
||||
healthy_server("alpha", &["search"]),
|
||||
degraded_server("beta", &["write"], "high latency"),
|
||||
],
|
||||
DiscoveryResult {
|
||||
tools: vec![tool("search"), tool("write")],
|
||||
resources: Vec::new(),
|
||||
partial: true,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let healthcheck = lifecycle.healthcheck();
|
||||
|
||||
// then
|
||||
match healthcheck.state {
|
||||
PluginState::Degraded {
|
||||
healthy_servers,
|
||||
failed_servers,
|
||||
} => {
|
||||
assert_eq!(
|
||||
healthy_servers,
|
||||
vec!["alpha".to_string(), "beta".to_string()]
|
||||
);
|
||||
assert!(failed_servers.is_empty());
|
||||
}
|
||||
other => panic!("expected degraded state, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn complete_failure_when_all_servers_fail() {
|
||||
// given
|
||||
let lifecycle = MockPluginLifecycle::new(
|
||||
"failed-plugin",
|
||||
true,
|
||||
vec![
|
||||
failed_server("alpha", &["search"], "timeout"),
|
||||
failed_server("beta", &["read"], "handshake failed"),
|
||||
],
|
||||
DiscoveryResult {
|
||||
tools: Vec::new(),
|
||||
resources: Vec::new(),
|
||||
partial: false,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let healthcheck = lifecycle.healthcheck();
|
||||
let discovery = lifecycle.discover();
|
||||
|
||||
// then
|
||||
match &healthcheck.state {
|
||||
PluginState::Failed { reason } => {
|
||||
assert_eq!(reason, "all 2 servers failed");
|
||||
}
|
||||
other => panic!("expected failed state, got {other:?}"),
|
||||
}
|
||||
assert!(!discovery.partial);
|
||||
assert!(discovery.tools.is_empty());
|
||||
assert!(discovery.resources.is_empty());
|
||||
assert!(healthcheck.degraded_mode(&discovery).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn graceful_shutdown() {
|
||||
// given
|
||||
let mut lifecycle = MockPluginLifecycle::new(
|
||||
"shutdown-plugin",
|
||||
true,
|
||||
vec![healthy_server("alpha", &["search"])],
|
||||
DiscoveryResult {
|
||||
tools: vec![tool("search")],
|
||||
resources: Vec::new(),
|
||||
partial: false,
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
// when
|
||||
let shutdown = lifecycle.shutdown();
|
||||
let post_shutdown = lifecycle.healthcheck();
|
||||
|
||||
// then
|
||||
assert_eq!(shutdown, Ok(()));
|
||||
assert_eq!(PluginLifecycleEvent::Shutdown.to_string(), "shutdown");
|
||||
assert_eq!(post_shutdown.state, PluginState::Stopped);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,581 @@
|
||||
use std::time::Duration;
|
||||
|
||||
pub type GreenLevel = u8;
|
||||
|
||||
const STALE_BRANCH_THRESHOLD: Duration = Duration::from_secs(60 * 60);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PolicyRule {
|
||||
pub name: String,
|
||||
pub condition: PolicyCondition,
|
||||
pub action: PolicyAction,
|
||||
pub priority: u32,
|
||||
}
|
||||
|
||||
impl PolicyRule {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
name: impl Into<String>,
|
||||
condition: PolicyCondition,
|
||||
action: PolicyAction,
|
||||
priority: u32,
|
||||
) -> Self {
|
||||
Self {
|
||||
name: name.into(),
|
||||
condition,
|
||||
action,
|
||||
priority,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn matches(&self, context: &LaneContext) -> bool {
|
||||
self.condition.matches(context)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PolicyCondition {
|
||||
And(Vec<PolicyCondition>),
|
||||
Or(Vec<PolicyCondition>),
|
||||
GreenAt { level: GreenLevel },
|
||||
StaleBranch,
|
||||
StartupBlocked,
|
||||
LaneCompleted,
|
||||
LaneReconciled,
|
||||
ReviewPassed,
|
||||
ScopedDiff,
|
||||
TimedOut { duration: Duration },
|
||||
}
|
||||
|
||||
impl PolicyCondition {
|
||||
#[must_use]
|
||||
pub fn matches(&self, context: &LaneContext) -> bool {
|
||||
match self {
|
||||
Self::And(conditions) => conditions
|
||||
.iter()
|
||||
.all(|condition| condition.matches(context)),
|
||||
Self::Or(conditions) => conditions
|
||||
.iter()
|
||||
.any(|condition| condition.matches(context)),
|
||||
Self::GreenAt { level } => context.green_level >= *level,
|
||||
Self::StaleBranch => context.branch_freshness >= STALE_BRANCH_THRESHOLD,
|
||||
Self::StartupBlocked => context.blocker == LaneBlocker::Startup,
|
||||
Self::LaneCompleted => context.completed,
|
||||
Self::LaneReconciled => context.reconciled,
|
||||
Self::ReviewPassed => context.review_status == ReviewStatus::Approved,
|
||||
Self::ScopedDiff => context.diff_scope == DiffScope::Scoped,
|
||||
Self::TimedOut { duration } => context.branch_freshness >= *duration,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum PolicyAction {
|
||||
MergeToDev,
|
||||
MergeForward,
|
||||
RecoverOnce,
|
||||
Escalate { reason: String },
|
||||
CloseoutLane,
|
||||
CleanupSession,
|
||||
Reconcile { reason: ReconcileReason },
|
||||
Notify { channel: String },
|
||||
Block { reason: String },
|
||||
Chain(Vec<PolicyAction>),
|
||||
}
|
||||
|
||||
/// Why a lane was reconciled without further action.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ReconcileReason {
|
||||
/// Branch already merged into main — no PR needed.
|
||||
AlreadyMerged,
|
||||
/// Work superseded by another lane or direct commit.
|
||||
Superseded,
|
||||
/// PR would be empty — all changes already landed.
|
||||
EmptyDiff,
|
||||
/// Lane manually closed by operator.
|
||||
ManualClose,
|
||||
}
|
||||
|
||||
impl PolicyAction {
|
||||
fn flatten_into(&self, actions: &mut Vec<PolicyAction>) {
|
||||
match self {
|
||||
Self::Chain(chained) => {
|
||||
for action in chained {
|
||||
action.flatten_into(actions);
|
||||
}
|
||||
}
|
||||
_ => actions.push(self.clone()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum LaneBlocker {
|
||||
None,
|
||||
Startup,
|
||||
External,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ReviewStatus {
|
||||
Pending,
|
||||
Approved,
|
||||
Rejected,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum DiffScope {
|
||||
Full,
|
||||
Scoped,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LaneContext {
|
||||
pub lane_id: String,
|
||||
pub green_level: GreenLevel,
|
||||
pub branch_freshness: Duration,
|
||||
pub blocker: LaneBlocker,
|
||||
pub review_status: ReviewStatus,
|
||||
pub diff_scope: DiffScope,
|
||||
pub completed: bool,
|
||||
pub reconciled: bool,
|
||||
}
|
||||
|
||||
impl LaneContext {
|
||||
#[must_use]
|
||||
pub fn new(
|
||||
lane_id: impl Into<String>,
|
||||
green_level: GreenLevel,
|
||||
branch_freshness: Duration,
|
||||
blocker: LaneBlocker,
|
||||
review_status: ReviewStatus,
|
||||
diff_scope: DiffScope,
|
||||
completed: bool,
|
||||
) -> Self {
|
||||
Self {
|
||||
lane_id: lane_id.into(),
|
||||
green_level,
|
||||
branch_freshness,
|
||||
blocker,
|
||||
review_status,
|
||||
diff_scope,
|
||||
completed,
|
||||
reconciled: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a lane context that is already reconciled (no further action needed).
|
||||
#[must_use]
|
||||
pub fn reconciled(lane_id: impl Into<String>) -> Self {
|
||||
Self {
|
||||
lane_id: lane_id.into(),
|
||||
green_level: 0,
|
||||
branch_freshness: Duration::from_secs(0),
|
||||
blocker: LaneBlocker::None,
|
||||
review_status: ReviewStatus::Pending,
|
||||
diff_scope: DiffScope::Full,
|
||||
completed: true,
|
||||
reconciled: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PolicyEngine {
|
||||
rules: Vec<PolicyRule>,
|
||||
}
|
||||
|
||||
impl PolicyEngine {
|
||||
#[must_use]
|
||||
pub fn new(mut rules: Vec<PolicyRule>) -> Self {
|
||||
rules.sort_by_key(|rule| rule.priority);
|
||||
Self { rules }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn rules(&self) -> &[PolicyRule] {
|
||||
&self.rules
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn evaluate(&self, context: &LaneContext) -> Vec<PolicyAction> {
|
||||
evaluate(self, context)
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn evaluate(engine: &PolicyEngine, context: &LaneContext) -> Vec<PolicyAction> {
|
||||
let mut actions = Vec::new();
|
||||
for rule in &engine.rules {
|
||||
if rule.matches(context) {
|
||||
rule.action.flatten_into(&mut actions);
|
||||
}
|
||||
}
|
||||
actions
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{
|
||||
evaluate, DiffScope, LaneBlocker, LaneContext, PolicyAction, PolicyCondition, PolicyEngine,
|
||||
PolicyRule, ReconcileReason, ReviewStatus, STALE_BRANCH_THRESHOLD,
|
||||
};
|
||||
|
||||
fn default_context() -> LaneContext {
|
||||
LaneContext::new(
|
||||
"lane-7",
|
||||
0,
|
||||
Duration::from_secs(0),
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_to_dev_rule_fires_for_green_scoped_reviewed_lane() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"merge-to-dev",
|
||||
PolicyCondition::And(vec![
|
||||
PolicyCondition::GreenAt { level: 2 },
|
||||
PolicyCondition::ScopedDiff,
|
||||
PolicyCondition::ReviewPassed,
|
||||
]),
|
||||
PolicyAction::MergeToDev,
|
||||
20,
|
||||
)]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
3,
|
||||
Duration::from_secs(5),
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Approved,
|
||||
DiffScope::Scoped,
|
||||
false,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(actions, vec![PolicyAction::MergeToDev]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_branch_rule_fires_at_threshold() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"merge-forward",
|
||||
PolicyCondition::StaleBranch,
|
||||
PolicyAction::MergeForward,
|
||||
10,
|
||||
)]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
1,
|
||||
STALE_BRANCH_THRESHOLD,
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(actions, vec![PolicyAction::MergeForward]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn startup_blocked_rule_recovers_then_escalates() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"startup-recovery",
|
||||
PolicyCondition::StartupBlocked,
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::RecoverOnce,
|
||||
PolicyAction::Escalate {
|
||||
reason: "startup remained blocked".to_string(),
|
||||
},
|
||||
]),
|
||||
15,
|
||||
)]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
0,
|
||||
Duration::from_secs(0),
|
||||
LaneBlocker::Startup,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::RecoverOnce,
|
||||
PolicyAction::Escalate {
|
||||
reason: "startup remained blocked".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completed_lane_rule_closes_out_and_cleans_up() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"lane-closeout",
|
||||
PolicyCondition::LaneCompleted,
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::CloseoutLane,
|
||||
PolicyAction::CleanupSession,
|
||||
]),
|
||||
30,
|
||||
)]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
0,
|
||||
Duration::from_secs(0),
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
true,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![PolicyAction::CloseoutLane, PolicyAction::CleanupSession]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matching_rules_are_returned_in_priority_order_with_stable_ties() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![
|
||||
PolicyRule::new(
|
||||
"late-cleanup",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::CleanupSession,
|
||||
30,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"first-notify",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::Notify {
|
||||
channel: "ops".to_string(),
|
||||
},
|
||||
10,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"second-notify",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::Notify {
|
||||
channel: "review".to_string(),
|
||||
},
|
||||
10,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"merge",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::MergeToDev,
|
||||
20,
|
||||
),
|
||||
]);
|
||||
let context = default_context();
|
||||
|
||||
// when
|
||||
let actions = evaluate(&engine, &context);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::Notify {
|
||||
channel: "ops".to_string(),
|
||||
},
|
||||
PolicyAction::Notify {
|
||||
channel: "review".to_string(),
|
||||
},
|
||||
PolicyAction::MergeToDev,
|
||||
PolicyAction::CleanupSession,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn combinators_handle_empty_cases_and_nested_chains() {
|
||||
// given
|
||||
let engine = PolicyEngine::new(vec![
|
||||
PolicyRule::new(
|
||||
"empty-and",
|
||||
PolicyCondition::And(vec![]),
|
||||
PolicyAction::Notify {
|
||||
channel: "orchestrator".to_string(),
|
||||
},
|
||||
5,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"empty-or",
|
||||
PolicyCondition::Or(vec![]),
|
||||
PolicyAction::Block {
|
||||
reason: "should not fire".to_string(),
|
||||
},
|
||||
10,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"nested",
|
||||
PolicyCondition::Or(vec![
|
||||
PolicyCondition::StartupBlocked,
|
||||
PolicyCondition::And(vec![
|
||||
PolicyCondition::GreenAt { level: 2 },
|
||||
PolicyCondition::TimedOut {
|
||||
duration: Duration::from_secs(5),
|
||||
},
|
||||
]),
|
||||
]),
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::Notify {
|
||||
channel: "alerts".to_string(),
|
||||
},
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::MergeForward,
|
||||
PolicyAction::CleanupSession,
|
||||
]),
|
||||
]),
|
||||
15,
|
||||
),
|
||||
]);
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
2,
|
||||
Duration::from_secs(10),
|
||||
LaneBlocker::External,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::Notify {
|
||||
channel: "orchestrator".to_string(),
|
||||
},
|
||||
PolicyAction::Notify {
|
||||
channel: "alerts".to_string(),
|
||||
},
|
||||
PolicyAction::MergeForward,
|
||||
PolicyAction::CleanupSession,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconciled_lane_emits_reconcile_and_cleanup() {
|
||||
// given — a lane where branch is already merged, no PR needed, session stale
|
||||
let engine = PolicyEngine::new(vec![
|
||||
PolicyRule::new(
|
||||
"reconcile-closeout",
|
||||
PolicyCondition::LaneReconciled,
|
||||
PolicyAction::Chain(vec![
|
||||
PolicyAction::Reconcile {
|
||||
reason: ReconcileReason::AlreadyMerged,
|
||||
},
|
||||
PolicyAction::CloseoutLane,
|
||||
PolicyAction::CleanupSession,
|
||||
]),
|
||||
5,
|
||||
),
|
||||
// This rule should NOT fire — reconciled lanes are completed but we want
|
||||
// the more specific reconcile rule to handle them
|
||||
PolicyRule::new(
|
||||
"generic-closeout",
|
||||
PolicyCondition::And(vec![
|
||||
PolicyCondition::LaneCompleted,
|
||||
// Only fire if NOT reconciled
|
||||
PolicyCondition::And(vec![]),
|
||||
]),
|
||||
PolicyAction::CloseoutLane,
|
||||
30,
|
||||
),
|
||||
]);
|
||||
let context = LaneContext::reconciled("lane-9411");
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then — reconcile rule fires first (priority 5), then generic closeout also fires
|
||||
// because reconciled context has completed=true
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::Reconcile {
|
||||
reason: ReconcileReason::AlreadyMerged,
|
||||
},
|
||||
PolicyAction::CloseoutLane,
|
||||
PolicyAction::CleanupSession,
|
||||
PolicyAction::CloseoutLane,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconciled_context_has_correct_defaults() {
|
||||
let ctx = LaneContext::reconciled("test-lane");
|
||||
assert_eq!(ctx.lane_id, "test-lane");
|
||||
assert!(ctx.completed);
|
||||
assert!(ctx.reconciled);
|
||||
assert_eq!(ctx.blocker, LaneBlocker::None);
|
||||
assert_eq!(ctx.green_level, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_reconciled_lane_does_not_trigger_reconcile_rule() {
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"reconcile-closeout",
|
||||
PolicyCondition::LaneReconciled,
|
||||
PolicyAction::Reconcile {
|
||||
reason: ReconcileReason::EmptyDiff,
|
||||
},
|
||||
5,
|
||||
)]);
|
||||
// Normal completed lane — not reconciled
|
||||
let context = LaneContext::new(
|
||||
"lane-7",
|
||||
0,
|
||||
Duration::from_secs(0),
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
true,
|
||||
);
|
||||
|
||||
let actions = engine.evaluate(&context);
|
||||
assert!(actions.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reconcile_reason_variants_are_distinct() {
|
||||
assert_ne!(ReconcileReason::AlreadyMerged, ReconcileReason::Superseded);
|
||||
assert_ne!(ReconcileReason::EmptyDiff, ReconcileReason::ManualClose);
|
||||
}
|
||||
}
|
||||
@@ -4,7 +4,9 @@ use std::path::{Path, PathBuf};
|
||||
use std::process::Command;
|
||||
|
||||
use crate::config::{ConfigError, ConfigLoader, RuntimeConfig};
|
||||
use crate::git_context::GitContext;
|
||||
|
||||
/// Errors raised while assembling the final system prompt.
|
||||
#[derive(Debug)]
|
||||
pub enum PromptBuildError {
|
||||
Io(std::io::Error),
|
||||
@@ -34,24 +36,29 @@ impl From<ConfigError> for PromptBuildError {
|
||||
}
|
||||
}
|
||||
|
||||
/// Marker separating static prompt scaffolding from dynamic runtime context.
|
||||
pub const SYSTEM_PROMPT_DYNAMIC_BOUNDARY: &str = "__SYSTEM_PROMPT_DYNAMIC_BOUNDARY__";
|
||||
/// Human-readable default frontier model name embedded into generated prompts.
|
||||
pub const FRONTIER_MODEL_NAME: &str = "Claude Opus 4.6";
|
||||
const MAX_INSTRUCTION_FILE_CHARS: usize = 4_000;
|
||||
const MAX_TOTAL_INSTRUCTION_CHARS: usize = 12_000;
|
||||
|
||||
/// Contents of an instruction file included in prompt construction.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ContextFile {
|
||||
pub path: PathBuf,
|
||||
pub content: String,
|
||||
}
|
||||
|
||||
/// Project-local context injected into the rendered system prompt.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct ProjectContext {
|
||||
pub cwd: PathBuf,
|
||||
pub current_date: String,
|
||||
pub git_status: Option<String>,
|
||||
pub git_diff: Option<String>,
|
||||
pub git_context: Option<GitContext>,
|
||||
pub instruction_files: Vec<ContextFile>,
|
||||
pub memory_files: Vec<ContextFile>,
|
||||
}
|
||||
|
||||
impl ProjectContext {
|
||||
@@ -61,13 +68,13 @@ impl ProjectContext {
|
||||
) -> std::io::Result<Self> {
|
||||
let cwd = cwd.into();
|
||||
let instruction_files = discover_instruction_files(&cwd)?;
|
||||
let memory_files = discover_memory_files(&cwd)?;
|
||||
Ok(Self {
|
||||
cwd,
|
||||
current_date: current_date.into(),
|
||||
git_status: None,
|
||||
git_diff: None,
|
||||
git_context: None,
|
||||
instruction_files,
|
||||
memory_files,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -77,10 +84,13 @@ impl ProjectContext {
|
||||
) -> std::io::Result<Self> {
|
||||
let mut context = Self::discover(cwd, current_date)?;
|
||||
context.git_status = read_git_status(&context.cwd);
|
||||
context.git_diff = read_git_diff(&context.cwd);
|
||||
context.git_context = GitContext::detect(&context.cwd);
|
||||
Ok(context)
|
||||
}
|
||||
}
|
||||
|
||||
/// Builder for the runtime system prompt and dynamic environment sections.
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
pub struct SystemPromptBuilder {
|
||||
output_style_name: Option<String>,
|
||||
@@ -147,9 +157,6 @@ impl SystemPromptBuilder {
|
||||
if !project_context.instruction_files.is_empty() {
|
||||
sections.push(render_instruction_files(&project_context.instruction_files));
|
||||
}
|
||||
if !project_context.memory_files.is_empty() {
|
||||
sections.push(render_memory_files(&project_context.memory_files));
|
||||
}
|
||||
}
|
||||
if let Some(config) = &self.config {
|
||||
sections.push(render_config_section(config));
|
||||
@@ -187,12 +194,13 @@ impl SystemPromptBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
/// Formats each item as an indented bullet for prompt sections.
|
||||
#[must_use]
|
||||
pub fn prepend_bullets(items: Vec<String>) -> Vec<String> {
|
||||
items.into_iter().map(|item| format!(" - {item}")).collect()
|
||||
}
|
||||
|
||||
fn discover_context_directories(cwd: &Path) -> Vec<PathBuf> {
|
||||
fn discover_instruction_files(cwd: &Path) -> std::io::Result<Vec<ContextFile>> {
|
||||
let mut directories = Vec::new();
|
||||
let mut cursor = Some(cwd);
|
||||
while let Some(dir) = cursor {
|
||||
@@ -200,19 +208,14 @@ fn discover_context_directories(cwd: &Path) -> Vec<PathBuf> {
|
||||
cursor = dir.parent();
|
||||
}
|
||||
directories.reverse();
|
||||
directories
|
||||
}
|
||||
|
||||
fn discover_instruction_files(cwd: &Path) -> std::io::Result<Vec<ContextFile>> {
|
||||
let directories = discover_context_directories(cwd);
|
||||
|
||||
let mut files = Vec::new();
|
||||
for dir in directories {
|
||||
for candidate in [
|
||||
dir.join("CLAUDE.md"),
|
||||
dir.join("CLAUDE.local.md"),
|
||||
dir.join(".claude").join("CLAUDE.md"),
|
||||
dir.join(".claude").join("instructions.md"),
|
||||
dir.join(".claw").join("CLAUDE.md"),
|
||||
dir.join(".claw").join("instructions.md"),
|
||||
] {
|
||||
push_context_file(&mut files, candidate)?;
|
||||
}
|
||||
@@ -220,26 +223,6 @@ fn discover_instruction_files(cwd: &Path) -> std::io::Result<Vec<ContextFile>> {
|
||||
Ok(dedupe_instruction_files(files))
|
||||
}
|
||||
|
||||
fn discover_memory_files(cwd: &Path) -> std::io::Result<Vec<ContextFile>> {
|
||||
let mut files = Vec::new();
|
||||
for dir in discover_context_directories(cwd) {
|
||||
let memory_dir = dir.join(".claude").join("memory");
|
||||
let Ok(entries) = fs::read_dir(&memory_dir) else {
|
||||
continue;
|
||||
};
|
||||
let mut paths = entries
|
||||
.flatten()
|
||||
.map(|entry| entry.path())
|
||||
.filter(|path| path.is_file())
|
||||
.collect::<Vec<_>>();
|
||||
paths.sort();
|
||||
for path in paths {
|
||||
push_context_file(&mut files, path)?;
|
||||
}
|
||||
}
|
||||
Ok(dedupe_instruction_files(files))
|
||||
}
|
||||
|
||||
fn push_context_file(files: &mut Vec<ContextFile>, path: PathBuf) -> std::io::Result<()> {
|
||||
match fs::read_to_string(&path) {
|
||||
Ok(content) if !content.trim().is_empty() => {
|
||||
@@ -270,6 +253,38 @@ fn read_git_status(cwd: &Path) -> Option<String> {
|
||||
}
|
||||
}
|
||||
|
||||
fn read_git_diff(cwd: &Path) -> Option<String> {
|
||||
let mut sections = Vec::new();
|
||||
|
||||
let staged = read_git_output(cwd, &["diff", "--cached"])?;
|
||||
if !staged.trim().is_empty() {
|
||||
sections.push(format!("Staged changes:\n{}", staged.trim_end()));
|
||||
}
|
||||
|
||||
let unstaged = read_git_output(cwd, &["diff"])?;
|
||||
if !unstaged.trim().is_empty() {
|
||||
sections.push(format!("Unstaged changes:\n{}", unstaged.trim_end()));
|
||||
}
|
||||
|
||||
if sections.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(sections.join("\n\n"))
|
||||
}
|
||||
}
|
||||
|
||||
fn read_git_output(cwd: &Path, args: &[&str]) -> Option<String> {
|
||||
let output = Command::new("git")
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
.ok()?;
|
||||
if !output.status.success() {
|
||||
return None;
|
||||
}
|
||||
String::from_utf8(output.stdout).ok()
|
||||
}
|
||||
|
||||
fn render_project_context(project_context: &ProjectContext) -> String {
|
||||
let mut lines = vec!["# Project context".to_string()];
|
||||
let mut bullets = vec![
|
||||
@@ -282,31 +297,38 @@ fn render_project_context(project_context: &ProjectContext) -> String {
|
||||
project_context.instruction_files.len()
|
||||
));
|
||||
}
|
||||
if !project_context.memory_files.is_empty() {
|
||||
bullets.push(format!(
|
||||
"Project memory files discovered: {}.",
|
||||
project_context.memory_files.len()
|
||||
));
|
||||
}
|
||||
lines.extend(prepend_bullets(bullets));
|
||||
if let Some(status) = &project_context.git_status {
|
||||
lines.push(String::new());
|
||||
lines.push("Git status snapshot:".to_string());
|
||||
lines.push(status.clone());
|
||||
}
|
||||
if let Some(ref gc) = project_context.git_context {
|
||||
if !gc.recent_commits.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push("Recent commits (last 5):".to_string());
|
||||
for c in &gc.recent_commits {
|
||||
lines.push(format!(" {} {}", c.hash, c.subject));
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(diff) = &project_context.git_diff {
|
||||
lines.push(String::new());
|
||||
lines.push("Git diff snapshot:".to_string());
|
||||
lines.push(diff.clone());
|
||||
}
|
||||
if let Some(git_context) = &project_context.git_context {
|
||||
let rendered = git_context.render();
|
||||
if !rendered.is_empty() {
|
||||
lines.push(String::new());
|
||||
lines.push(rendered);
|
||||
}
|
||||
}
|
||||
lines.join("\n")
|
||||
}
|
||||
|
||||
fn render_instruction_files(files: &[ContextFile]) -> String {
|
||||
render_context_file_section("# Claude instructions", files)
|
||||
}
|
||||
|
||||
fn render_memory_files(files: &[ContextFile]) -> String {
|
||||
render_context_file_section("# Project memory", files)
|
||||
}
|
||||
|
||||
fn render_context_file_section(title: &str, files: &[ContextFile]) -> String {
|
||||
let mut sections = vec![title.to_string()];
|
||||
let mut sections = vec!["# Claude instructions".to_string()];
|
||||
let mut remaining_chars = MAX_TOTAL_INSTRUCTION_CHARS;
|
||||
for file in files {
|
||||
if remaining_chars == 0 {
|
||||
@@ -406,6 +428,7 @@ fn collapse_blank_lines(content: &str) -> String {
|
||||
result
|
||||
}
|
||||
|
||||
/// Loads config and project context, then renders the system prompt text.
|
||||
pub fn load_system_prompt(
|
||||
cwd: impl Into<PathBuf>,
|
||||
current_date: impl Into<String>,
|
||||
@@ -426,7 +449,7 @@ fn render_config_section(config: &RuntimeConfig) -> String {
|
||||
let mut lines = vec!["# Runtime config".to_string()];
|
||||
if config.loaded_entries().is_empty() {
|
||||
lines.extend(prepend_bullets(vec![
|
||||
"No Claude Code settings files loaded.".to_string(),
|
||||
"No Claw Code settings files loaded.".to_string()
|
||||
]));
|
||||
return lines.join("\n");
|
||||
}
|
||||
@@ -498,9 +521,8 @@ fn get_actions_section() -> String {
|
||||
mod tests {
|
||||
use super::{
|
||||
collapse_blank_lines, display_context_path, normalize_instruction_content,
|
||||
render_instruction_content, render_instruction_files, render_memory_files,
|
||||
truncate_instruction_content, ContextFile, ProjectContext, SystemPromptBuilder,
|
||||
SYSTEM_PROMPT_DYNAMIC_BOUNDARY,
|
||||
render_instruction_content, render_instruction_files, truncate_instruction_content,
|
||||
ContextFile, ProjectContext, SystemPromptBuilder, SYSTEM_PROMPT_DYNAMIC_BOUNDARY,
|
||||
};
|
||||
use crate::config::ConfigLoader;
|
||||
use std::fs;
|
||||
@@ -519,27 +541,34 @@ mod tests {
|
||||
crate::test_env_lock()
|
||||
}
|
||||
|
||||
fn ensure_valid_cwd() {
|
||||
if std::env::current_dir().is_err() {
|
||||
std::env::set_current_dir(env!("CARGO_MANIFEST_DIR"))
|
||||
.expect("test cwd should be recoverable");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discovers_instruction_files_from_ancestor_chain() {
|
||||
let root = temp_dir();
|
||||
let nested = root.join("apps").join("api");
|
||||
fs::create_dir_all(nested.join(".claude")).expect("nested claude dir");
|
||||
fs::create_dir_all(nested.join(".claw")).expect("nested claw dir");
|
||||
fs::write(root.join("CLAUDE.md"), "root instructions").expect("write root instructions");
|
||||
fs::write(root.join("CLAUDE.local.md"), "local instructions")
|
||||
.expect("write local instructions");
|
||||
fs::create_dir_all(root.join("apps")).expect("apps dir");
|
||||
fs::create_dir_all(root.join("apps").join(".claude")).expect("apps claude dir");
|
||||
fs::create_dir_all(root.join("apps").join(".claw")).expect("apps claw dir");
|
||||
fs::write(root.join("apps").join("CLAUDE.md"), "apps instructions")
|
||||
.expect("write apps instructions");
|
||||
fs::write(
|
||||
root.join("apps").join(".claude").join("instructions.md"),
|
||||
root.join("apps").join(".claw").join("instructions.md"),
|
||||
"apps dot claude instructions",
|
||||
)
|
||||
.expect("write apps dot claude instructions");
|
||||
fs::write(nested.join(".claude").join("CLAUDE.md"), "nested rules")
|
||||
fs::write(nested.join(".claw").join("CLAUDE.md"), "nested rules")
|
||||
.expect("write nested rules");
|
||||
fs::write(
|
||||
nested.join(".claude").join("instructions.md"),
|
||||
nested.join(".claw").join("instructions.md"),
|
||||
"nested instructions",
|
||||
)
|
||||
.expect("write nested instructions");
|
||||
@@ -565,35 +594,6 @@ mod tests {
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discovers_project_memory_files_from_ancestor_chain() {
|
||||
let root = temp_dir();
|
||||
let nested = root.join("apps").join("api");
|
||||
fs::create_dir_all(root.join(".claude").join("memory")).expect("root memory dir");
|
||||
fs::create_dir_all(nested.join(".claude").join("memory")).expect("nested memory dir");
|
||||
fs::write(
|
||||
root.join(".claude").join("memory").join("2026-03-30.md"),
|
||||
"root memory",
|
||||
)
|
||||
.expect("write root memory");
|
||||
fs::write(
|
||||
nested.join(".claude").join("memory").join("2026-03-31.md"),
|
||||
"nested memory",
|
||||
)
|
||||
.expect("write nested memory");
|
||||
|
||||
let context = ProjectContext::discover(&nested, "2026-03-31").expect("context should load");
|
||||
let contents = context
|
||||
.memory_files
|
||||
.iter()
|
||||
.map(|file| file.content.as_str())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(contents, vec!["root memory", "nested memory"]);
|
||||
assert!(render_memory_files(&context.memory_files).contains("# Project memory"));
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dedupes_identical_instruction_content_across_scopes() {
|
||||
let root = temp_dir();
|
||||
@@ -628,13 +628,15 @@ mod tests {
|
||||
#[test]
|
||||
fn displays_context_paths_compactly() {
|
||||
assert_eq!(
|
||||
display_context_path(Path::new("/tmp/project/.claude/CLAUDE.md")),
|
||||
display_context_path(Path::new("/tmp/project/.claw/CLAUDE.md")),
|
||||
"CLAUDE.md"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_with_git_includes_status_snapshot() {
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir");
|
||||
std::process::Command::new("git")
|
||||
@@ -652,6 +654,133 @@ mod tests {
|
||||
assert!(status.contains("## No commits yet on") || status.contains("## "));
|
||||
assert!(status.contains("?? CLAUDE.md"));
|
||||
assert!(status.contains("?? tracked.txt"));
|
||||
assert!(context.git_diff.is_none());
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_with_git_includes_recent_commits_and_renders_them() {
|
||||
// given: a git repo with three commits and a current branch
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir");
|
||||
std::process::Command::new("git")
|
||||
.args(["init", "--quiet", "-b", "main"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git init should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["config", "user.email", "tests@example.com"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git config email should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["config", "user.name", "Runtime Prompt Tests"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git config name should run");
|
||||
for (file, message) in [
|
||||
("a.txt", "first commit"),
|
||||
("b.txt", "second commit"),
|
||||
("c.txt", "third commit"),
|
||||
] {
|
||||
fs::write(root.join(file), "x\n").expect("write commit file");
|
||||
std::process::Command::new("git")
|
||||
.args(["add", file])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git add should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["commit", "-m", message, "--quiet"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git commit should run");
|
||||
}
|
||||
fs::write(root.join("d.txt"), "staged\n").expect("write staged file");
|
||||
std::process::Command::new("git")
|
||||
.args(["add", "d.txt"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git add staged should run");
|
||||
|
||||
// when: discovering project context with git auto-include
|
||||
let context =
|
||||
ProjectContext::discover_with_git(&root, "2026-03-31").expect("context should load");
|
||||
let rendered = SystemPromptBuilder::new()
|
||||
.with_os("linux", "6.8")
|
||||
.with_project_context(context.clone())
|
||||
.render();
|
||||
|
||||
// then: branch, recent commits and staged files are present in context
|
||||
let gc = context
|
||||
.git_context
|
||||
.as_ref()
|
||||
.expect("git context should be present");
|
||||
let commits: String = gc
|
||||
.recent_commits
|
||||
.iter()
|
||||
.map(|c| c.subject.clone())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
assert!(commits.contains("first commit"));
|
||||
assert!(commits.contains("second commit"));
|
||||
assert!(commits.contains("third commit"));
|
||||
assert_eq!(gc.recent_commits.len(), 3);
|
||||
|
||||
let status = context.git_status.as_deref().expect("status snapshot");
|
||||
assert!(status.contains("## main"));
|
||||
assert!(status.contains("A d.txt"));
|
||||
|
||||
assert!(rendered.contains("Recent commits (last 5):"));
|
||||
assert!(rendered.contains("first commit"));
|
||||
assert!(rendered.contains("Git status snapshot:"));
|
||||
assert!(rendered.contains("## main"));
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn discover_with_git_includes_diff_snapshot_for_tracked_changes() {
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir");
|
||||
std::process::Command::new("git")
|
||||
.args(["init", "--quiet"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git init should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["config", "user.email", "tests@example.com"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git config email should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["config", "user.name", "Runtime Prompt Tests"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git config name should run");
|
||||
fs::write(root.join("tracked.txt"), "hello\n").expect("write tracked file");
|
||||
std::process::Command::new("git")
|
||||
.args(["add", "tracked.txt"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git add should run");
|
||||
std::process::Command::new("git")
|
||||
.args(["commit", "-m", "init", "--quiet"])
|
||||
.current_dir(&root)
|
||||
.status()
|
||||
.expect("git commit should run");
|
||||
fs::write(root.join("tracked.txt"), "hello\nworld\n").expect("rewrite tracked file");
|
||||
|
||||
let context =
|
||||
ProjectContext::discover_with_git(&root, "2026-03-31").expect("context should load");
|
||||
|
||||
let diff = context.git_diff.expect("git diff should be present");
|
||||
assert!(diff.contains("Unstaged changes:"));
|
||||
assert!(diff.contains("tracked.txt"));
|
||||
|
||||
fs::remove_dir_all(root).expect("cleanup temp dir");
|
||||
}
|
||||
@@ -659,20 +788,21 @@ mod tests {
|
||||
#[test]
|
||||
fn load_system_prompt_reads_claude_files_and_config() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(root.join(".claude")).expect("claude dir");
|
||||
fs::create_dir_all(root.join(".claw")).expect("claw dir");
|
||||
fs::write(root.join("CLAUDE.md"), "Project rules").expect("write instructions");
|
||||
fs::write(
|
||||
root.join(".claude").join("settings.json"),
|
||||
root.join(".claw").join("settings.json"),
|
||||
r#"{"permissionMode":"acceptEdits"}"#,
|
||||
)
|
||||
.expect("write settings");
|
||||
|
||||
let _guard = env_lock();
|
||||
ensure_valid_cwd();
|
||||
let previous = std::env::current_dir().expect("cwd");
|
||||
let original_home = std::env::var("HOME").ok();
|
||||
let original_claude_home = std::env::var("CLAUDE_CONFIG_HOME").ok();
|
||||
let original_claw_home = std::env::var("CLAW_CONFIG_HOME").ok();
|
||||
std::env::set_var("HOME", &root);
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", root.join("missing-home"));
|
||||
std::env::set_var("CLAW_CONFIG_HOME", root.join("missing-home"));
|
||||
std::env::set_current_dir(&root).expect("change cwd");
|
||||
let prompt = super::load_system_prompt(&root, "2026-03-31", "linux", "6.8")
|
||||
.expect("system prompt should load")
|
||||
@@ -687,10 +817,10 @@ mod tests {
|
||||
} else {
|
||||
std::env::remove_var("HOME");
|
||||
}
|
||||
if let Some(value) = original_claude_home {
|
||||
std::env::set_var("CLAUDE_CONFIG_HOME", value);
|
||||
if let Some(value) = original_claw_home {
|
||||
std::env::set_var("CLAW_CONFIG_HOME", value);
|
||||
} else {
|
||||
std::env::remove_var("CLAUDE_CONFIG_HOME");
|
||||
std::env::remove_var("CLAW_CONFIG_HOME");
|
||||
}
|
||||
|
||||
assert!(prompt.contains("Project rules"));
|
||||
@@ -701,10 +831,10 @@ mod tests {
|
||||
#[test]
|
||||
fn renders_claude_code_style_sections_with_project_context() {
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(root.join(".claude")).expect("claude dir");
|
||||
fs::create_dir_all(root.join(".claw")).expect("claw dir");
|
||||
fs::write(root.join("CLAUDE.md"), "Project rules").expect("write CLAUDE.md");
|
||||
fs::write(
|
||||
root.join(".claude").join("settings.json"),
|
||||
root.join(".claw").join("settings.json"),
|
||||
r#"{"permissionMode":"acceptEdits"}"#,
|
||||
)
|
||||
.expect("write settings");
|
||||
@@ -743,9 +873,9 @@ mod tests {
|
||||
fn discovers_dot_claude_instructions_markdown() {
|
||||
let root = temp_dir();
|
||||
let nested = root.join("apps").join("api");
|
||||
fs::create_dir_all(nested.join(".claude")).expect("nested claude dir");
|
||||
fs::create_dir_all(nested.join(".claw")).expect("nested claw dir");
|
||||
fs::write(
|
||||
nested.join(".claude").join("instructions.md"),
|
||||
nested.join(".claw").join("instructions.md"),
|
||||
"instruction markdown",
|
||||
)
|
||||
.expect("write instructions.md");
|
||||
@@ -754,7 +884,7 @@ mod tests {
|
||||
assert!(context
|
||||
.instruction_files
|
||||
.iter()
|
||||
.any(|file| file.path.ends_with(".claude/instructions.md")));
|
||||
.any(|file| file.path.ends_with(".claw/instructions.md")));
|
||||
assert!(
|
||||
render_instruction_files(&context.instruction_files).contains("instruction markdown")
|
||||
);
|
||||