Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ef48f58ef1 | ||
|
|
9a86aa6444 | ||
|
|
21b0887469 | ||
|
|
0d89231caa | ||
|
|
b445a3320f | ||
|
|
650a24b6e2 | ||
|
|
d018276fc1 | ||
|
|
387a8bb13f | ||
|
|
243a1ff74f | ||
|
|
583d191527 | ||
|
|
074bd5b7b7 | ||
|
|
bec07658b8 | ||
|
|
f403d3b107 | ||
|
|
bd494184fc | ||
|
|
a22700562d | ||
|
|
c14196c730 | ||
|
|
f544125c01 | ||
|
|
ccebabe605 | ||
|
|
cdf24b87b4 | ||
|
|
770fb8d0e7 | ||
|
|
e38e3ee4d7 | ||
|
|
331b8fc811 | ||
|
|
72b5f2fe80 | ||
|
|
b200198df7 | ||
|
|
2fd6241bd8 | ||
|
|
5b046836b9 | ||
|
|
549deb9a89 | ||
|
|
146260083c | ||
|
|
3ba60be514 | ||
|
|
d6341d54c1 | ||
|
|
cd01d0e387 | ||
|
|
863958b94c | ||
|
|
9455280f24 | ||
|
|
c92403994d | ||
|
|
8d4a739c05 | ||
|
|
e2f061fd08 | ||
|
|
c139fe9bee | ||
|
|
6a7cea810e | ||
|
|
842abcfe85 | ||
|
|
807e29c8a1 | ||
|
|
e84133527e | ||
|
|
32e89df631 | ||
|
|
1f8cfbce38 | ||
|
|
1e5002b521 | ||
|
|
d5d99af2d0 | ||
|
|
5180cc5658 | ||
|
|
964cc25821 | ||
|
|
8ab16276bf | ||
|
|
b8dadbfbf5 | ||
|
|
46581fe442 | ||
|
|
92f33c75c0 | ||
|
|
5f46fec5ad | ||
|
|
771f716625 | ||
|
|
d3e41be7f1 | ||
|
|
691ea57832 | ||
|
|
4d65f5c1a2 | ||
|
|
8b6bf4cee7 | ||
|
|
647b407444 | ||
|
|
5eeb7be4cc | ||
|
|
f8bc5cf264 | ||
|
|
346ea0b91b | ||
|
|
6076041f19 | ||
|
|
9f3be03463 | ||
|
|
c30bb8aa59 | ||
|
|
88cd2e31df | ||
|
|
1adf11d572 | ||
|
|
9b0c9b5739 | ||
|
|
cf8d5a8389 | ||
|
|
cba31c4f95 | ||
|
|
fa30059790 | ||
|
|
d9c5f60598 | ||
|
|
9b7fe16edb | ||
|
|
c8f95cd72b | ||
|
|
66dde1b74a | ||
|
|
99b78d6ea4 | ||
|
|
3db3dfa60d | ||
|
|
0ac188caad | ||
|
|
0794e76f07 | ||
|
|
b510387045 | ||
|
|
6e378185e9 | ||
|
|
019e9900ed | ||
|
|
67423d005a | ||
|
|
4db21e9595 | ||
|
|
daf98cc750 | ||
|
|
2ad2ec087f | ||
|
|
0346b7dd3a | ||
|
|
a8f5da6427 | ||
|
|
c996eb7b1b | ||
|
|
14757e0780 | ||
|
|
188c35f8a6 | ||
|
|
2de0b0e2af | ||
|
|
c024d8b21f | ||
|
|
a66c301fa3 | ||
|
|
321a1a681a | ||
|
|
2d1cade31b | ||
|
|
6fe404329d | ||
|
|
add5513ac5 | ||
|
|
8465b6923b | ||
|
|
32981ffa28 | ||
|
|
cb24430c56 | ||
|
|
071045f556 | ||
|
|
a96bb6c60f | ||
|
|
d6a814258c | ||
|
|
4bae5ee132 | ||
|
|
619ae71866 | ||
|
|
6037aaeff1 | ||
|
|
5b106b840d | ||
|
|
4586764a0e | ||
|
|
3faf8dd365 | ||
|
|
450556559a | ||
|
|
44e4758078 | ||
|
|
01bf54ad15 | ||
|
|
507c2460b9 |
@@ -0,0 +1 @@
|
||||
{"messages":[],"version":1}
|
||||
@@ -0,0 +1 @@
|
||||
{"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 @@
|
||||
{"messages":[],"version":1}
|
||||
@@ -0,0 +1 @@
|
||||
{"messages":[{"blocks":[{"text":"What is 2+2? Reply with just the number.","type":"text"}],"role":"user"},{"blocks":[{"text":"4","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":19,"output_tokens":5}}],"version":1}
|
||||
@@ -0,0 +1 @@
|
||||
{"messages":[{"blocks":[{"text":"Say hello in exactly 3 words","type":"text"}],"role":"user"},{"blocks":[{"text":"Hello there, friend!","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":14,"output_tokens":8}}],"version":1}
|
||||
@@ -0,0 +1 @@
|
||||
{"messages":[{"blocks":[{"text":"Say hi in one sentence","type":"text"}],"role":"user"},{"blocks":[{"text":"Hi! I'm Claude, ready to help you with any software engineering tasks or questions you have.","type":"text"}],"role":"assistant","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":11,"output_tokens":23}}],"version":1}
|
||||
@@ -1,5 +0,0 @@
|
||||
{
|
||||
"permissions": {
|
||||
"defaultMode": "dontAsk"
|
||||
}
|
||||
}
|
||||
@@ -1,3 +0,0 @@
|
||||
ANTHROPIC_API_KEY="9494feba6f7c45f48c3dfc35a85ffd89.2WUCscxcSp92ETNg"
|
||||
ANTHROPIC_BASE_URL="https://open.bigmodel.cn/api/anthropic"
|
||||
CLAW_MODEL="glm-5"
|
||||
@@ -0,0 +1 @@
|
||||
github: instructkr
|
||||
@@ -1,36 +0,0 @@
|
||||
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,6 +1,4 @@
|
||||
target/
|
||||
__pycache__/
|
||||
archive/
|
||||
.omx/
|
||||
.clawd-agents/
|
||||
# Claw Code local artifacts
|
||||
.claw/
|
||||
.claude/
|
||||
|
||||
@@ -1,113 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,15 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,43 +0,0 @@
|
||||
# 贡献指南
|
||||
|
||||
感谢你为 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 的范围仅限于一个明确的更改。
|
||||
- 说明更改动机、实现摘要以及你运行的验证。
|
||||
- 在请求审查之前,确保本地检查已通过。
|
||||
- 如果审查反馈导致行为更改,请重新运行相关的验证命令。
|
||||
@@ -1,55 +0,0 @@
|
||||
[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,121 +1,191 @@
|
||||
# Claw Code
|
||||
# Rewriting Project Claw Code
|
||||
|
||||
Claw Code 是一个使用安全 Rust 实现的本地编程代理(coding-agent)命令行工具。它的设计灵感来自 **Claude Code**,并作为一个**净室实现(clean-room implementation)**开发:旨在提供强大的本地代理体验,但它**不是** Claude 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>
|
||||
|
||||
Rust 工作区是当前主要的产品界面。`claw` 二进制文件在单个工作区内提供交互式会话、单次提示、工作区感知工具、本地代理工作流以及支持插件的操作。
|
||||
<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>
|
||||
|
||||
## 当前状态
|
||||
<p align="center">
|
||||
<img src="assets/clawd-hero.jpeg" alt="Claw" width="300" />
|
||||
</p>
|
||||
|
||||
- **版本:** `0.1.0`
|
||||
- **发布阶段:** 初始公开发布,源码编译分发
|
||||
- **主要实现:** 本仓库中的 Rust 工作区
|
||||
- **平台焦点:** macOS 和 Linux 开发工作站
|
||||
<p align="center">
|
||||
<strong>Better Harness Tools, not merely storing the archive of leaked Claude Code</strong>
|
||||
</p>
|
||||
|
||||
## 安装、构建与运行
|
||||
<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.
|
||||
|
||||
- Rust 稳定版工具链
|
||||
- Cargo
|
||||
- 你想使用的模型的提供商凭据
|
||||
> If you find this work useful, consider [sponsoring @instructkr on GitHub](https://github.com/sponsors/instructkr) to support continued open-source harness engineering research.
|
||||
|
||||
你可以通过环境变量或在项目根目录创建 **`.env`** 文件来配置 API 密钥:
|
||||
---
|
||||
|
||||
**配置 Claude (推荐):**
|
||||
```bash
|
||||
ANTHROPIC_API_KEY="..."
|
||||
# 使用兼容的端点时可选
|
||||
export ANTHROPIC_BASE_URL="https://api.anthropic.com"
|
||||
## 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
|
||||
```
|
||||
|
||||
Grok 模型:
|
||||
## 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:
|
||||
|
||||
```bash
|
||||
export XAI_API_KEY="..."
|
||||
# 使用兼容的端点时可选
|
||||
export XAI_BASE_URL="https://api.x.ai"
|
||||
python3 -m src.main summary
|
||||
```
|
||||
|
||||
也可以使用 OAuth 登录:
|
||||
Print the current Python workspace manifest:
|
||||
|
||||
```bash
|
||||
cargo run --bin claw -- login
|
||||
python3 -m src.main manifest
|
||||
```
|
||||
|
||||
### 本地安装
|
||||
List the current Python modules:
|
||||
|
||||
```bash
|
||||
cargo install --path crates/claw-cli --locked
|
||||
python3 -m src.main subsystems --limit 16
|
||||
```
|
||||
|
||||
### 从源码构建
|
||||
Run verification:
|
||||
|
||||
```bash
|
||||
cargo build --release -p claw-cli
|
||||
python3 -m unittest discover -s tests -v
|
||||
```
|
||||
|
||||
### 运行
|
||||
|
||||
在工作区内运行:
|
||||
Run the parity audit against the local ignored archive (when present):
|
||||
|
||||
```bash
|
||||
cargo run --bin claw -- --help
|
||||
cargo run --bin claw --
|
||||
cargo run --bin claw -- prompt "总结此工作区"
|
||||
cargo run --bin claw -- --model sonnet "审查最新更改"
|
||||
python3 -m src.main parity-audit
|
||||
```
|
||||
|
||||
运行发布版本:
|
||||
Inspect mirrored command/tool inventories:
|
||||
|
||||
```bash
|
||||
./target/release/claw
|
||||
./target/release/claw prompt "解释 crates/runtime"
|
||||
python3 -m src.main commands --limit 10
|
||||
python3 -m src.main tools --limit 10
|
||||
```
|
||||
|
||||
## 支持的功能
|
||||
## Current Parity Checkpoint
|
||||
|
||||
- 交互式 REPL 和单次提示执行
|
||||
- 已保存会话的检查和恢复流程
|
||||
- 内置工作区工具:shell、文件读/写/编辑、搜索、网页获取/搜索、待办事项和笔记本更新
|
||||
- 斜杠命令:状态、压缩、配置检查、差异(diff)、导出、会话管理和版本报告
|
||||
- 本地代理和技能发现:通过 `claw agents` 和 `claw skills`
|
||||
- 通过命令行和斜杠命令界面发现并管理插件
|
||||
- OAuth 登录/注销,以及从命令行选择模型/提供商
|
||||
- 工作区感知的指令/配置加载(`CLAW.md`、配置文件、权限、插件设置)
|
||||
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.
|
||||
|
||||
## 当前限制
|
||||
|
||||
- 目前公开发布**仅限源码构建**;此工作区尚未设置 crates.io 发布
|
||||
- GitHub CI 验证 `cargo check`、`cargo test` 和发布构建,但尚未提供自动化的发布打包
|
||||
- 当前 CI 目标为 Ubuntu 和 macOS;Windows 的发布就绪性仍待建立
|
||||
- 一些实时提供商集成覆盖是可选的,因为它们需要外部凭据 and 网络访问
|
||||
- 命令界面可能会在 `0.x` 系列期间继续演进
|
||||
## Built with `oh-my-codex`
|
||||
|
||||
## 实现现状
|
||||
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.
|
||||
|
||||
Rust 工作区是当前的产品实现。目前包含以下 crate:
|
||||
- **`$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
|
||||
|
||||
- `claw-cli` — 面向用户的二进制文件
|
||||
- `api` — 提供商客户端和流式处理
|
||||
- `runtime` — 会话、配置、权限、提示词和运行时循环
|
||||
- `tools` — 内置工具实现
|
||||
- `commands` — 斜杠命令注册和处理程序
|
||||
- `plugins` — 插件发现、注册和生命周期支持
|
||||
- `lsp` — 语言服务器协议支持类型和进程助手
|
||||
- `server` 和 `compat-harness` — 支持服务和兼容性工具
|
||||
### OmX workflow screenshots
|
||||
|
||||
## 路线图
|
||||

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

|
||||
|
||||
- 0.1.0 发行说明草案:[`docs/releases/0.1.0.md`](docs/releases/0.1.0.md)
|
||||
*Split-pane review and verification flow during the final README wording pass.*
|
||||
|
||||
## 许可
|
||||
## 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**.
|
||||
|
||||
|
After Width: | Height: | Size: 233 KiB |
|
After Width: | Height: | Size: 4.8 KiB |
|
After Width: | Height: | Size: 1.3 MiB |
|
After Width: | Height: | Size: 1.4 MiB |
|
After Width: | Height: | Size: 312 KiB |
|
After Width: | Height: | Size: 812 KiB |
|
After Width: | Height: | Size: 873 KiB |
@@ -1,53 +0,0 @@
|
||||
# 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?;
|
||||
```
|
||||
@@ -1,241 +0,0 @@
|
||||
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
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,572 +0,0 @@
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -1,344 +0,0 @@
|
||||
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:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,39 +0,0 @@
|
||||
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,
|
||||
};
|
||||
@@ -1,735 +0,0 @@
|
||||
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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,928 +0,0 @@
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -1,531 +0,0 @@
|
||||
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),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,88 +0,0 @@
|
||||
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),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,173 +0,0 @@
|
||||
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,58 +0,0 @@
|
||||
# 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
|
||||
```
|
||||
@@ -1,432 +0,0 @@
|
||||
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");
|
||||
}
|
||||
}
|
||||
@@ -1,55 +0,0 @@
|
||||
# 命令模块 (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"]
|
||||
```
|
||||
@@ -1,48 +0,0 @@
|
||||
# 兼容性测试套件模块 (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());
|
||||
}
|
||||
```
|
||||
@@ -1,56 +0,0 @@
|
||||
# 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());
|
||||
```
|
||||
@@ -1,465 +0,0 @@
|
||||
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()
|
||||
}
|
||||
@@ -1,62 +0,0 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -1,284 +0,0 @@
|
||||
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");
|
||||
}
|
||||
}
|
||||
@@ -1,191 +0,0 @@
|
||||
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,
|
||||
))
|
||||
});
|
||||
}
|
||||
@@ -1,186 +0,0 @@
|
||||
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",
|
||||
}
|
||||
}
|
||||
@@ -1,18 +0,0 @@
|
||||
[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
|
||||
@@ -1,34 +0,0 @@
|
||||
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(())
|
||||
}
|
||||
@@ -1,13 +0,0 @@
|
||||
[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
|
||||
@@ -1,59 +0,0 @@
|
||||
# 插件模块 (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"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
@@ -1,10 +0,0 @@
|
||||
{
|
||||
"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"]
|
||||
}
|
||||
}
|
||||
@@ -1,2 +0,0 @@
|
||||
#!/bin/sh
|
||||
printf '%s\n' 'example bundled post hook'
|
||||
@@ -1,2 +0,0 @@
|
||||
#!/bin/sh
|
||||
printf '%s\n' 'example bundled pre hook'
|
||||
@@ -1,10 +0,0 @@
|
||||
{
|
||||
"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"]
|
||||
}
|
||||
}
|
||||
@@ -1,2 +0,0 @@
|
||||
#!/bin/sh
|
||||
printf 'sample bundled post hook'
|
||||
@@ -1,2 +0,0 @@
|
||||
#!/bin/sh
|
||||
printf 'sample bundled pre hook'
|
||||
@@ -1,563 +0,0 @@
|
||||
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}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,21 +0,0 @@
|
||||
[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
|
||||
@@ -1,60 +0,0 @@
|
||||
# 运行时模块 (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());
|
||||
```
|
||||
@@ -1,111 +0,0 @@
|
||||
#[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);
|
||||
}
|
||||
}
|
||||
@@ -1,144 +0,0 @@
|
||||
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());
|
||||
}
|
||||
}
|
||||
@@ -1,827 +0,0 @@
|
||||
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"));
|
||||
}
|
||||
}
|
||||
@@ -1,901 +0,0 @@
|
||||
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"#
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,324 +0,0 @@
|
||||
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");
|
||||
}
|
||||
}
|
||||
@@ -1,152 +0,0 @@
|
||||
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,
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,988 +0,0 @@
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -1,383 +0,0 @@
|
||||
#![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"));
|
||||
}
|
||||
}
|
||||
@@ -1,179 +0,0 @@
|
||||
//! 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)
|
||||
}
|
||||
@@ -1,747 +0,0 @@
|
||||
#![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"));
|
||||
}
|
||||
}
|
||||
@@ -1,843 +0,0 @@
|
||||
#![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)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,440 +0,0 @@
|
||||
//! 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);
|
||||
}
|
||||
}
|
||||
@@ -1,920 +0,0 @@
|
||||
#![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());
|
||||
}
|
||||
}
|
||||
@@ -1,551 +0,0 @@
|
||||
#![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:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,683 +0,0 @@
|
||||
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")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,533 +0,0 @@
|
||||
#![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);
|
||||
}
|
||||
}
|
||||
@@ -1,581 +0,0 @@
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -1,631 +0,0 @@
|
||||
#![allow(clippy::cast_possible_truncation, clippy::uninlined_format_args)]
|
||||
//! Recovery recipes for common failure scenarios.
|
||||
//!
|
||||
//! Encodes known automatic recoveries for the six failure scenarios
|
||||
//! listed in ROADMAP item 8, and enforces one automatic recovery
|
||||
//! attempt before escalation. Each attempt is emitted as a structured
|
||||
//! recovery event.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::worker_boot::WorkerFailureKind;
|
||||
|
||||
/// The six failure scenarios that have known recovery recipes.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum FailureScenario {
|
||||
TrustPromptUnresolved,
|
||||
PromptMisdelivery,
|
||||
StaleBranch,
|
||||
CompileRedCrossCrate,
|
||||
McpHandshakeFailure,
|
||||
PartialPluginStartup,
|
||||
ProviderFailure,
|
||||
}
|
||||
|
||||
impl FailureScenario {
|
||||
/// Returns all known failure scenarios.
|
||||
#[must_use]
|
||||
pub fn all() -> &'static [FailureScenario] {
|
||||
&[
|
||||
Self::TrustPromptUnresolved,
|
||||
Self::PromptMisdelivery,
|
||||
Self::StaleBranch,
|
||||
Self::CompileRedCrossCrate,
|
||||
Self::McpHandshakeFailure,
|
||||
Self::PartialPluginStartup,
|
||||
Self::ProviderFailure,
|
||||
]
|
||||
}
|
||||
|
||||
/// Map a `WorkerFailureKind` to the corresponding `FailureScenario`.
|
||||
/// This is the bridge that lets recovery policy consume worker boot events.
|
||||
#[must_use]
|
||||
pub fn from_worker_failure_kind(kind: WorkerFailureKind) -> Self {
|
||||
match kind {
|
||||
WorkerFailureKind::TrustGate => Self::TrustPromptUnresolved,
|
||||
WorkerFailureKind::PromptDelivery => Self::PromptMisdelivery,
|
||||
WorkerFailureKind::Protocol => Self::McpHandshakeFailure,
|
||||
WorkerFailureKind::Provider => Self::ProviderFailure,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::fmt::Display for FailureScenario {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::TrustPromptUnresolved => write!(f, "trust_prompt_unresolved"),
|
||||
Self::PromptMisdelivery => write!(f, "prompt_misdelivery"),
|
||||
Self::StaleBranch => write!(f, "stale_branch"),
|
||||
Self::CompileRedCrossCrate => write!(f, "compile_red_cross_crate"),
|
||||
Self::McpHandshakeFailure => write!(f, "mcp_handshake_failure"),
|
||||
Self::PartialPluginStartup => write!(f, "partial_plugin_startup"),
|
||||
Self::ProviderFailure => write!(f, "provider_failure"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Individual step that can be executed as part of a recovery recipe.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RecoveryStep {
|
||||
AcceptTrustPrompt,
|
||||
RedirectPromptToAgent,
|
||||
RebaseBranch,
|
||||
CleanBuild,
|
||||
RetryMcpHandshake { timeout: u64 },
|
||||
RestartPlugin { name: String },
|
||||
RestartWorker,
|
||||
EscalateToHuman { reason: String },
|
||||
}
|
||||
|
||||
/// Policy governing what happens when automatic recovery is exhausted.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum EscalationPolicy {
|
||||
AlertHuman,
|
||||
LogAndContinue,
|
||||
Abort,
|
||||
}
|
||||
|
||||
/// A recovery recipe encodes the sequence of steps to attempt for a
|
||||
/// given failure scenario, along with the maximum number of automatic
|
||||
/// attempts and the escalation policy.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RecoveryRecipe {
|
||||
pub scenario: FailureScenario,
|
||||
pub steps: Vec<RecoveryStep>,
|
||||
pub max_attempts: u32,
|
||||
pub escalation_policy: EscalationPolicy,
|
||||
}
|
||||
|
||||
/// Outcome of a recovery attempt.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RecoveryResult {
|
||||
Recovered {
|
||||
steps_taken: u32,
|
||||
},
|
||||
PartialRecovery {
|
||||
recovered: Vec<RecoveryStep>,
|
||||
remaining: Vec<RecoveryStep>,
|
||||
},
|
||||
EscalationRequired {
|
||||
reason: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// Structured event emitted during recovery.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum RecoveryEvent {
|
||||
RecoveryAttempted {
|
||||
scenario: FailureScenario,
|
||||
recipe: RecoveryRecipe,
|
||||
result: RecoveryResult,
|
||||
},
|
||||
RecoverySucceeded,
|
||||
RecoveryFailed,
|
||||
Escalated,
|
||||
}
|
||||
|
||||
/// Minimal context for tracking recovery state and emitting events.
|
||||
///
|
||||
/// Holds per-scenario attempt counts, a structured event log, and an
|
||||
/// optional simulation knob for controlling step outcomes during tests.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct RecoveryContext {
|
||||
attempts: HashMap<FailureScenario, u32>,
|
||||
events: Vec<RecoveryEvent>,
|
||||
/// Optional step index at which simulated execution fails.
|
||||
/// `None` means all steps succeed.
|
||||
fail_at_step: Option<usize>,
|
||||
}
|
||||
|
||||
impl RecoveryContext {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Configure a step index at which simulated execution will fail.
|
||||
#[must_use]
|
||||
pub fn with_fail_at_step(mut self, index: usize) -> Self {
|
||||
self.fail_at_step = Some(index);
|
||||
self
|
||||
}
|
||||
|
||||
/// Returns the structured event log populated during recovery.
|
||||
#[must_use]
|
||||
pub fn events(&self) -> &[RecoveryEvent] {
|
||||
&self.events
|
||||
}
|
||||
|
||||
/// Returns the number of recovery attempts made for a scenario.
|
||||
#[must_use]
|
||||
pub fn attempt_count(&self, scenario: &FailureScenario) -> u32 {
|
||||
self.attempts.get(scenario).copied().unwrap_or(0)
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the known recovery recipe for the given failure scenario.
|
||||
#[must_use]
|
||||
pub fn recipe_for(scenario: &FailureScenario) -> RecoveryRecipe {
|
||||
match scenario {
|
||||
FailureScenario::TrustPromptUnresolved => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::AcceptTrustPrompt],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
FailureScenario::PromptMisdelivery => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::RedirectPromptToAgent],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
FailureScenario::StaleBranch => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::RebaseBranch, RecoveryStep::CleanBuild],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
FailureScenario::CompileRedCrossCrate => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::CleanBuild],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
FailureScenario::McpHandshakeFailure => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::RetryMcpHandshake { timeout: 5000 }],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::Abort,
|
||||
},
|
||||
FailureScenario::PartialPluginStartup => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![
|
||||
RecoveryStep::RestartPlugin {
|
||||
name: "stalled".to_string(),
|
||||
},
|
||||
RecoveryStep::RetryMcpHandshake { timeout: 3000 },
|
||||
],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::LogAndContinue,
|
||||
},
|
||||
FailureScenario::ProviderFailure => RecoveryRecipe {
|
||||
scenario: *scenario,
|
||||
steps: vec![RecoveryStep::RestartWorker],
|
||||
max_attempts: 1,
|
||||
escalation_policy: EscalationPolicy::AlertHuman,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Attempts automatic recovery for the given failure scenario.
|
||||
///
|
||||
/// Looks up the recipe, enforces the one-attempt-before-escalation
|
||||
/// policy, simulates step execution (controlled by the context), and
|
||||
/// emits structured [`RecoveryEvent`]s for every attempt.
|
||||
pub fn attempt_recovery(scenario: &FailureScenario, ctx: &mut RecoveryContext) -> RecoveryResult {
|
||||
let recipe = recipe_for(scenario);
|
||||
let attempt_count = ctx.attempts.entry(*scenario).or_insert(0);
|
||||
|
||||
// Enforce one automatic recovery attempt before escalation.
|
||||
if *attempt_count >= recipe.max_attempts {
|
||||
let result = RecoveryResult::EscalationRequired {
|
||||
reason: format!(
|
||||
"max recovery attempts ({}) exceeded for {}",
|
||||
recipe.max_attempts, scenario
|
||||
),
|
||||
};
|
||||
ctx.events.push(RecoveryEvent::RecoveryAttempted {
|
||||
scenario: *scenario,
|
||||
recipe,
|
||||
result: result.clone(),
|
||||
});
|
||||
ctx.events.push(RecoveryEvent::Escalated);
|
||||
return result;
|
||||
}
|
||||
|
||||
*attempt_count += 1;
|
||||
|
||||
// Execute steps, honoring the optional fail_at_step simulation.
|
||||
let fail_index = ctx.fail_at_step;
|
||||
let mut executed = Vec::new();
|
||||
let mut failed = false;
|
||||
|
||||
for (i, step) in recipe.steps.iter().enumerate() {
|
||||
if fail_index == Some(i) {
|
||||
failed = true;
|
||||
break;
|
||||
}
|
||||
executed.push(step.clone());
|
||||
}
|
||||
|
||||
let result = if failed {
|
||||
let remaining: Vec<RecoveryStep> = recipe.steps[executed.len()..].to_vec();
|
||||
if executed.is_empty() {
|
||||
RecoveryResult::EscalationRequired {
|
||||
reason: format!("recovery failed at first step for {}", scenario),
|
||||
}
|
||||
} else {
|
||||
RecoveryResult::PartialRecovery {
|
||||
recovered: executed,
|
||||
remaining,
|
||||
}
|
||||
}
|
||||
} else {
|
||||
RecoveryResult::Recovered {
|
||||
steps_taken: recipe.steps.len() as u32,
|
||||
}
|
||||
};
|
||||
|
||||
// Emit the attempt as structured event data.
|
||||
ctx.events.push(RecoveryEvent::RecoveryAttempted {
|
||||
scenario: *scenario,
|
||||
recipe,
|
||||
result: result.clone(),
|
||||
});
|
||||
|
||||
match &result {
|
||||
RecoveryResult::Recovered { .. } => {
|
||||
ctx.events.push(RecoveryEvent::RecoverySucceeded);
|
||||
}
|
||||
RecoveryResult::PartialRecovery { .. } => {
|
||||
ctx.events.push(RecoveryEvent::RecoveryFailed);
|
||||
}
|
||||
RecoveryResult::EscalationRequired { .. } => {
|
||||
ctx.events.push(RecoveryEvent::Escalated);
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn each_scenario_has_a_matching_recipe() {
|
||||
// given
|
||||
let scenarios = FailureScenario::all();
|
||||
|
||||
// when / then
|
||||
for scenario in scenarios {
|
||||
let recipe = recipe_for(scenario);
|
||||
assert_eq!(
|
||||
recipe.scenario, *scenario,
|
||||
"recipe scenario should match requested scenario"
|
||||
);
|
||||
assert!(
|
||||
!recipe.steps.is_empty(),
|
||||
"recipe for {} should have at least one step",
|
||||
scenario
|
||||
);
|
||||
assert!(
|
||||
recipe.max_attempts >= 1,
|
||||
"recipe for {} should allow at least one attempt",
|
||||
scenario
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn successful_recovery_returns_recovered_and_emits_events() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::TrustPromptUnresolved;
|
||||
|
||||
// when
|
||||
let result = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
assert_eq!(result, RecoveryResult::Recovered { steps_taken: 1 });
|
||||
assert_eq!(ctx.events().len(), 2);
|
||||
assert!(matches!(
|
||||
&ctx.events()[0],
|
||||
RecoveryEvent::RecoveryAttempted {
|
||||
scenario: s,
|
||||
result: r,
|
||||
..
|
||||
} if *s == FailureScenario::TrustPromptUnresolved
|
||||
&& matches!(r, RecoveryResult::Recovered { steps_taken: 1 })
|
||||
));
|
||||
assert_eq!(ctx.events()[1], RecoveryEvent::RecoverySucceeded);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escalation_after_max_attempts_exceeded() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::PromptMisdelivery;
|
||||
|
||||
// when — first attempt succeeds
|
||||
let first = attempt_recovery(&scenario, &mut ctx);
|
||||
assert!(matches!(first, RecoveryResult::Recovered { .. }));
|
||||
|
||||
// when — second attempt should escalate
|
||||
let second = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
assert!(
|
||||
matches!(
|
||||
&second,
|
||||
RecoveryResult::EscalationRequired { reason }
|
||||
if reason.contains("max recovery attempts")
|
||||
),
|
||||
"second attempt should require escalation, got: {second:?}"
|
||||
);
|
||||
assert_eq!(ctx.attempt_count(&scenario), 1);
|
||||
assert!(ctx
|
||||
.events()
|
||||
.iter()
|
||||
.any(|e| matches!(e, RecoveryEvent::Escalated)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_recovery_when_step_fails_midway() {
|
||||
// given — PartialPluginStartup has two steps; fail at step index 1
|
||||
let mut ctx = RecoveryContext::new().with_fail_at_step(1);
|
||||
let scenario = FailureScenario::PartialPluginStartup;
|
||||
|
||||
// when
|
||||
let result = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
match &result {
|
||||
RecoveryResult::PartialRecovery {
|
||||
recovered,
|
||||
remaining,
|
||||
} => {
|
||||
assert_eq!(recovered.len(), 1, "one step should have succeeded");
|
||||
assert_eq!(remaining.len(), 1, "one step should remain");
|
||||
assert!(matches!(recovered[0], RecoveryStep::RestartPlugin { .. }));
|
||||
assert!(matches!(
|
||||
remaining[0],
|
||||
RecoveryStep::RetryMcpHandshake { .. }
|
||||
));
|
||||
}
|
||||
other => panic!("expected PartialRecovery, got {other:?}"),
|
||||
}
|
||||
assert!(ctx
|
||||
.events()
|
||||
.iter()
|
||||
.any(|e| matches!(e, RecoveryEvent::RecoveryFailed)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_step_failure_escalates_immediately() {
|
||||
// given — fail at step index 0
|
||||
let mut ctx = RecoveryContext::new().with_fail_at_step(0);
|
||||
let scenario = FailureScenario::CompileRedCrossCrate;
|
||||
|
||||
// when
|
||||
let result = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
assert!(
|
||||
matches!(
|
||||
&result,
|
||||
RecoveryResult::EscalationRequired { reason }
|
||||
if reason.contains("failed at first step")
|
||||
),
|
||||
"zero-step failure should escalate, got: {result:?}"
|
||||
);
|
||||
assert!(ctx
|
||||
.events()
|
||||
.iter()
|
||||
.any(|e| matches!(e, RecoveryEvent::Escalated)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn emitted_events_include_structured_attempt_data() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::McpHandshakeFailure;
|
||||
|
||||
// when
|
||||
let _ = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then — verify the RecoveryAttempted event carries full context
|
||||
let attempted = ctx
|
||||
.events()
|
||||
.iter()
|
||||
.find(|e| matches!(e, RecoveryEvent::RecoveryAttempted { .. }))
|
||||
.expect("should have emitted RecoveryAttempted event");
|
||||
|
||||
match attempted {
|
||||
RecoveryEvent::RecoveryAttempted {
|
||||
scenario: s,
|
||||
recipe,
|
||||
result,
|
||||
} => {
|
||||
assert_eq!(*s, scenario);
|
||||
assert_eq!(recipe.scenario, scenario);
|
||||
assert!(!recipe.steps.is_empty());
|
||||
assert!(matches!(result, RecoveryResult::Recovered { .. }));
|
||||
}
|
||||
_ => unreachable!(),
|
||||
}
|
||||
|
||||
// Verify the event is serializable as structured JSON
|
||||
let json = serde_json::to_string(&ctx.events()[0])
|
||||
.expect("recovery event should be serializable to JSON");
|
||||
assert!(
|
||||
json.contains("mcp_handshake_failure"),
|
||||
"serialized event should contain scenario name"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recovery_context_tracks_attempts_per_scenario() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
|
||||
// when
|
||||
assert_eq!(ctx.attempt_count(&FailureScenario::StaleBranch), 0);
|
||||
attempt_recovery(&FailureScenario::StaleBranch, &mut ctx);
|
||||
|
||||
// then
|
||||
assert_eq!(ctx.attempt_count(&FailureScenario::StaleBranch), 1);
|
||||
assert_eq!(ctx.attempt_count(&FailureScenario::PromptMisdelivery), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_branch_recipe_has_rebase_then_clean_build() {
|
||||
// given
|
||||
let recipe = recipe_for(&FailureScenario::StaleBranch);
|
||||
|
||||
// then
|
||||
assert_eq!(recipe.steps.len(), 2);
|
||||
assert_eq!(recipe.steps[0], RecoveryStep::RebaseBranch);
|
||||
assert_eq!(recipe.steps[1], RecoveryStep::CleanBuild);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_plugin_startup_recipe_has_restart_then_handshake() {
|
||||
// given
|
||||
let recipe = recipe_for(&FailureScenario::PartialPluginStartup);
|
||||
|
||||
// then
|
||||
assert_eq!(recipe.steps.len(), 2);
|
||||
assert!(matches!(
|
||||
recipe.steps[0],
|
||||
RecoveryStep::RestartPlugin { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
recipe.steps[1],
|
||||
RecoveryStep::RetryMcpHandshake { timeout: 3000 }
|
||||
));
|
||||
assert_eq!(recipe.escalation_policy, EscalationPolicy::LogAndContinue);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failure_scenario_display_all_variants() {
|
||||
// given
|
||||
let cases = [
|
||||
(
|
||||
FailureScenario::TrustPromptUnresolved,
|
||||
"trust_prompt_unresolved",
|
||||
),
|
||||
(FailureScenario::PromptMisdelivery, "prompt_misdelivery"),
|
||||
(FailureScenario::StaleBranch, "stale_branch"),
|
||||
(
|
||||
FailureScenario::CompileRedCrossCrate,
|
||||
"compile_red_cross_crate",
|
||||
),
|
||||
(
|
||||
FailureScenario::McpHandshakeFailure,
|
||||
"mcp_handshake_failure",
|
||||
),
|
||||
(
|
||||
FailureScenario::PartialPluginStartup,
|
||||
"partial_plugin_startup",
|
||||
),
|
||||
];
|
||||
|
||||
// when / then
|
||||
for (scenario, expected) in &cases {
|
||||
assert_eq!(scenario.to_string(), *expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multi_step_success_reports_correct_steps_taken() {
|
||||
// given — StaleBranch has 2 steps, no simulated failure
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::StaleBranch;
|
||||
|
||||
// when
|
||||
let result = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then
|
||||
assert_eq!(result, RecoveryResult::Recovered { steps_taken: 2 });
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mcp_handshake_recipe_uses_abort_escalation_policy() {
|
||||
// given
|
||||
let recipe = recipe_for(&FailureScenario::McpHandshakeFailure);
|
||||
|
||||
// then
|
||||
assert_eq!(recipe.escalation_policy, EscalationPolicy::Abort);
|
||||
assert_eq!(recipe.max_attempts, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_failure_kind_maps_to_failure_scenario() {
|
||||
// given / when / then — verify the bridge is correct
|
||||
assert_eq!(
|
||||
FailureScenario::from_worker_failure_kind(WorkerFailureKind::TrustGate),
|
||||
FailureScenario::TrustPromptUnresolved,
|
||||
);
|
||||
assert_eq!(
|
||||
FailureScenario::from_worker_failure_kind(WorkerFailureKind::PromptDelivery),
|
||||
FailureScenario::PromptMisdelivery,
|
||||
);
|
||||
assert_eq!(
|
||||
FailureScenario::from_worker_failure_kind(WorkerFailureKind::Protocol),
|
||||
FailureScenario::McpHandshakeFailure,
|
||||
);
|
||||
assert_eq!(
|
||||
FailureScenario::from_worker_failure_kind(WorkerFailureKind::Provider),
|
||||
FailureScenario::ProviderFailure,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_failure_recipe_uses_restart_worker_step() {
|
||||
// given
|
||||
let recipe = recipe_for(&FailureScenario::ProviderFailure);
|
||||
|
||||
// then
|
||||
assert_eq!(recipe.scenario, FailureScenario::ProviderFailure);
|
||||
assert!(recipe.steps.contains(&RecoveryStep::RestartWorker));
|
||||
assert_eq!(recipe.escalation_policy, EscalationPolicy::AlertHuman);
|
||||
assert_eq!(recipe.max_attempts, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_failure_recovery_attempt_succeeds_then_escalates() {
|
||||
// given
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let scenario = FailureScenario::ProviderFailure;
|
||||
|
||||
// when — first attempt
|
||||
let first = attempt_recovery(&scenario, &mut ctx);
|
||||
assert!(matches!(first, RecoveryResult::Recovered { .. }));
|
||||
|
||||
// when — second attempt should escalate (max_attempts=1)
|
||||
let second = attempt_recovery(&scenario, &mut ctx);
|
||||
assert!(matches!(second, RecoveryResult::EscalationRequired { .. }));
|
||||
assert!(ctx
|
||||
.events()
|
||||
.iter()
|
||||
.any(|e| matches!(e, RecoveryEvent::Escalated)));
|
||||
}
|
||||
}
|
||||
@@ -1,873 +0,0 @@
|
||||
#![allow(dead_code)]
|
||||
use std::env;
|
||||
use std::fmt::{Display, Formatter};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use crate::session::{Session, SessionError};
|
||||
|
||||
/// Per-worktree session store that namespaces on-disk session files by
|
||||
/// workspace fingerprint so that parallel `opencode serve` instances never
|
||||
/// collide.
|
||||
///
|
||||
/// Create via [`SessionStore::from_cwd`] (derives the store path from the
|
||||
/// server's working directory) or [`SessionStore::from_data_dir`] (honours an
|
||||
/// explicit `--data-dir` flag). Both constructors produce a directory layout
|
||||
/// of `<data_dir>/sessions/<workspace_hash>/` where `<workspace_hash>` is a
|
||||
/// stable hex digest of the canonical workspace root.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SessionStore {
|
||||
/// Resolved root of the session namespace, e.g.
|
||||
/// `/home/user/project/.claw/sessions/a1b2c3d4e5f60718/`.
|
||||
sessions_root: PathBuf,
|
||||
/// The canonical workspace path that was fingerprinted.
|
||||
workspace_root: PathBuf,
|
||||
}
|
||||
|
||||
impl SessionStore {
|
||||
/// Build a store from the server's current working directory.
|
||||
///
|
||||
/// The on-disk layout becomes `<cwd>/.claw/sessions/<workspace_hash>/`.
|
||||
pub fn from_cwd(cwd: impl AsRef<Path>) -> Result<Self, SessionControlError> {
|
||||
let cwd = cwd.as_ref();
|
||||
let sessions_root = cwd
|
||||
.join(".claw")
|
||||
.join("sessions")
|
||||
.join(workspace_fingerprint(cwd));
|
||||
fs::create_dir_all(&sessions_root)?;
|
||||
Ok(Self {
|
||||
sessions_root,
|
||||
workspace_root: cwd.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build a store from an explicit `--data-dir` flag.
|
||||
///
|
||||
/// The on-disk layout becomes `<data_dir>/sessions/<workspace_hash>/`
|
||||
/// where `<workspace_hash>` is derived from `workspace_root`.
|
||||
pub fn from_data_dir(
|
||||
data_dir: impl AsRef<Path>,
|
||||
workspace_root: impl AsRef<Path>,
|
||||
) -> Result<Self, SessionControlError> {
|
||||
let workspace_root = workspace_root.as_ref();
|
||||
let sessions_root = data_dir
|
||||
.as_ref()
|
||||
.join("sessions")
|
||||
.join(workspace_fingerprint(workspace_root));
|
||||
fs::create_dir_all(&sessions_root)?;
|
||||
Ok(Self {
|
||||
sessions_root,
|
||||
workspace_root: workspace_root.to_path_buf(),
|
||||
})
|
||||
}
|
||||
|
||||
/// The fully resolved sessions directory for this namespace.
|
||||
#[must_use]
|
||||
pub fn sessions_dir(&self) -> &Path {
|
||||
&self.sessions_root
|
||||
}
|
||||
|
||||
/// The workspace root this store is bound to.
|
||||
#[must_use]
|
||||
pub fn workspace_root(&self) -> &Path {
|
||||
&self.workspace_root
|
||||
}
|
||||
|
||||
pub fn create_handle(&self, session_id: &str) -> SessionHandle {
|
||||
let id = session_id.to_string();
|
||||
let path = self
|
||||
.sessions_root
|
||||
.join(format!("{id}.{PRIMARY_SESSION_EXTENSION}"));
|
||||
SessionHandle { id, path }
|
||||
}
|
||||
|
||||
pub fn resolve_reference(&self, reference: &str) -> Result<SessionHandle, SessionControlError> {
|
||||
if is_session_reference_alias(reference) {
|
||||
let latest = self.latest_session()?;
|
||||
return Ok(SessionHandle {
|
||||
id: latest.id,
|
||||
path: latest.path,
|
||||
});
|
||||
}
|
||||
|
||||
let direct = PathBuf::from(reference);
|
||||
let candidate = if direct.is_absolute() {
|
||||
direct.clone()
|
||||
} else {
|
||||
self.workspace_root.join(&direct)
|
||||
};
|
||||
let looks_like_path = direct.extension().is_some() || direct.components().count() > 1;
|
||||
let path = if candidate.exists() {
|
||||
candidate
|
||||
} else if looks_like_path {
|
||||
return Err(SessionControlError::Format(
|
||||
format_missing_session_reference(reference),
|
||||
));
|
||||
} else {
|
||||
self.resolve_managed_path(reference)?
|
||||
};
|
||||
|
||||
Ok(SessionHandle {
|
||||
id: session_id_from_path(&path).unwrap_or_else(|| reference.to_string()),
|
||||
path,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve_managed_path(&self, session_id: &str) -> Result<PathBuf, SessionControlError> {
|
||||
for extension in [PRIMARY_SESSION_EXTENSION, LEGACY_SESSION_EXTENSION] {
|
||||
let path = self.sessions_root.join(format!("{session_id}.{extension}"));
|
||||
if path.exists() {
|
||||
return Ok(path);
|
||||
}
|
||||
}
|
||||
Err(SessionControlError::Format(
|
||||
format_missing_session_reference(session_id),
|
||||
))
|
||||
}
|
||||
|
||||
pub fn list_sessions(&self) -> Result<Vec<ManagedSessionSummary>, SessionControlError> {
|
||||
let mut sessions = Vec::new();
|
||||
let read_result = fs::read_dir(&self.sessions_root);
|
||||
let entries = match read_result {
|
||||
Ok(entries) => entries,
|
||||
Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(sessions),
|
||||
Err(err) => return Err(err.into()),
|
||||
};
|
||||
for entry in entries {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if !is_managed_session_file(&path) {
|
||||
continue;
|
||||
}
|
||||
let metadata = entry.metadata()?;
|
||||
let modified_epoch_millis = metadata
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|time| time.duration_since(UNIX_EPOCH).ok())
|
||||
.map(|duration| duration.as_millis())
|
||||
.unwrap_or_default();
|
||||
let (id, message_count, parent_session_id, branch_name) =
|
||||
match Session::load_from_path(&path) {
|
||||
Ok(session) => {
|
||||
let parent_session_id = session
|
||||
.fork
|
||||
.as_ref()
|
||||
.map(|fork| fork.parent_session_id.clone());
|
||||
let branch_name = session
|
||||
.fork
|
||||
.as_ref()
|
||||
.and_then(|fork| fork.branch_name.clone());
|
||||
(
|
||||
session.session_id,
|
||||
session.messages.len(),
|
||||
parent_session_id,
|
||||
branch_name,
|
||||
)
|
||||
}
|
||||
Err(_) => (
|
||||
path.file_stem()
|
||||
.and_then(|value| value.to_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string(),
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
),
|
||||
};
|
||||
sessions.push(ManagedSessionSummary {
|
||||
id,
|
||||
path,
|
||||
modified_epoch_millis,
|
||||
message_count,
|
||||
parent_session_id,
|
||||
branch_name,
|
||||
});
|
||||
}
|
||||
sessions.sort_by(|left, right| {
|
||||
right
|
||||
.modified_epoch_millis
|
||||
.cmp(&left.modified_epoch_millis)
|
||||
.then_with(|| right.id.cmp(&left.id))
|
||||
});
|
||||
Ok(sessions)
|
||||
}
|
||||
|
||||
pub fn latest_session(&self) -> Result<ManagedSessionSummary, SessionControlError> {
|
||||
self.list_sessions()?
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| SessionControlError::Format(format_no_managed_sessions()))
|
||||
}
|
||||
|
||||
pub fn load_session(
|
||||
&self,
|
||||
reference: &str,
|
||||
) -> Result<LoadedManagedSession, SessionControlError> {
|
||||
let handle = self.resolve_reference(reference)?;
|
||||
let session = Session::load_from_path(&handle.path)?;
|
||||
Ok(LoadedManagedSession {
|
||||
handle: SessionHandle {
|
||||
id: session.session_id.clone(),
|
||||
path: handle.path,
|
||||
},
|
||||
session,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn fork_session(
|
||||
&self,
|
||||
session: &Session,
|
||||
branch_name: Option<String>,
|
||||
) -> Result<ForkedManagedSession, SessionControlError> {
|
||||
let parent_session_id = session.session_id.clone();
|
||||
let forked = session.fork(branch_name);
|
||||
let handle = self.create_handle(&forked.session_id);
|
||||
let branch_name = forked
|
||||
.fork
|
||||
.as_ref()
|
||||
.and_then(|fork| fork.branch_name.clone());
|
||||
let forked = forked.with_persistence_path(handle.path.clone());
|
||||
forked.save_to_path(&handle.path)?;
|
||||
Ok(ForkedManagedSession {
|
||||
parent_session_id,
|
||||
handle,
|
||||
session: forked,
|
||||
branch_name,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Stable hex fingerprint of a workspace path.
|
||||
///
|
||||
/// Uses FNV-1a (64-bit) to produce a 16-char hex string that partitions the
|
||||
/// on-disk session directory per workspace root.
|
||||
#[must_use]
|
||||
pub fn workspace_fingerprint(workspace_root: &Path) -> String {
|
||||
let input = workspace_root.to_string_lossy();
|
||||
let mut hash = 0xcbf2_9ce4_8422_2325_u64;
|
||||
for byte in input.as_bytes() {
|
||||
hash ^= u64::from(*byte);
|
||||
hash = hash.wrapping_mul(0x0100_0000_01b3);
|
||||
}
|
||||
format!("{hash:016x}")
|
||||
}
|
||||
|
||||
pub const PRIMARY_SESSION_EXTENSION: &str = "jsonl";
|
||||
pub const LEGACY_SESSION_EXTENSION: &str = "json";
|
||||
pub const LATEST_SESSION_REFERENCE: &str = "latest";
|
||||
|
||||
const SESSION_REFERENCE_ALIASES: &[&str] = &[LATEST_SESSION_REFERENCE, "last", "recent"];
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SessionHandle {
|
||||
pub id: String,
|
||||
pub path: PathBuf,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ManagedSessionSummary {
|
||||
pub id: String,
|
||||
pub path: PathBuf,
|
||||
pub modified_epoch_millis: u128,
|
||||
pub message_count: usize,
|
||||
pub parent_session_id: Option<String>,
|
||||
pub branch_name: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct LoadedManagedSession {
|
||||
pub handle: SessionHandle,
|
||||
pub session: Session,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ForkedManagedSession {
|
||||
pub parent_session_id: String,
|
||||
pub handle: SessionHandle,
|
||||
pub session: Session,
|
||||
pub branch_name: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum SessionControlError {
|
||||
Io(std::io::Error),
|
||||
Session(SessionError),
|
||||
Format(String),
|
||||
}
|
||||
|
||||
impl Display for SessionControlError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Io(error) => write!(f, "{error}"),
|
||||
Self::Session(error) => write!(f, "{error}"),
|
||||
Self::Format(error) => write!(f, "{error}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for SessionControlError {}
|
||||
|
||||
impl From<std::io::Error> for SessionControlError {
|
||||
fn from(value: std::io::Error) -> Self {
|
||||
Self::Io(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<SessionError> for SessionControlError {
|
||||
fn from(value: SessionError) -> Self {
|
||||
Self::Session(value)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sessions_dir() -> Result<PathBuf, SessionControlError> {
|
||||
managed_sessions_dir_for(env::current_dir()?)
|
||||
}
|
||||
|
||||
pub fn managed_sessions_dir_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
) -> Result<PathBuf, SessionControlError> {
|
||||
let path = base_dir.as_ref().join(".claw").join("sessions");
|
||||
fs::create_dir_all(&path)?;
|
||||
Ok(path)
|
||||
}
|
||||
|
||||
pub fn create_managed_session_handle(
|
||||
session_id: &str,
|
||||
) -> Result<SessionHandle, SessionControlError> {
|
||||
create_managed_session_handle_for(env::current_dir()?, session_id)
|
||||
}
|
||||
|
||||
pub fn create_managed_session_handle_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
session_id: &str,
|
||||
) -> Result<SessionHandle, SessionControlError> {
|
||||
let id = session_id.to_string();
|
||||
let path =
|
||||
managed_sessions_dir_for(base_dir)?.join(format!("{id}.{PRIMARY_SESSION_EXTENSION}"));
|
||||
Ok(SessionHandle { id, path })
|
||||
}
|
||||
|
||||
pub fn resolve_session_reference(reference: &str) -> Result<SessionHandle, SessionControlError> {
|
||||
resolve_session_reference_for(env::current_dir()?, reference)
|
||||
}
|
||||
|
||||
pub fn resolve_session_reference_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
reference: &str,
|
||||
) -> Result<SessionHandle, SessionControlError> {
|
||||
let base_dir = base_dir.as_ref();
|
||||
if is_session_reference_alias(reference) {
|
||||
let latest = latest_managed_session_for(base_dir)?;
|
||||
return Ok(SessionHandle {
|
||||
id: latest.id,
|
||||
path: latest.path,
|
||||
});
|
||||
}
|
||||
|
||||
let direct = PathBuf::from(reference);
|
||||
let candidate = if direct.is_absolute() {
|
||||
direct.clone()
|
||||
} else {
|
||||
base_dir.join(&direct)
|
||||
};
|
||||
let looks_like_path = direct.extension().is_some() || direct.components().count() > 1;
|
||||
let path = if candidate.exists() {
|
||||
candidate
|
||||
} else if looks_like_path {
|
||||
return Err(SessionControlError::Format(
|
||||
format_missing_session_reference(reference),
|
||||
));
|
||||
} else {
|
||||
resolve_managed_session_path_for(base_dir, reference)?
|
||||
};
|
||||
|
||||
Ok(SessionHandle {
|
||||
id: session_id_from_path(&path).unwrap_or_else(|| reference.to_string()),
|
||||
path,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn resolve_managed_session_path(session_id: &str) -> Result<PathBuf, SessionControlError> {
|
||||
resolve_managed_session_path_for(env::current_dir()?, session_id)
|
||||
}
|
||||
|
||||
pub fn resolve_managed_session_path_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
session_id: &str,
|
||||
) -> Result<PathBuf, SessionControlError> {
|
||||
let directory = managed_sessions_dir_for(base_dir)?;
|
||||
for extension in [PRIMARY_SESSION_EXTENSION, LEGACY_SESSION_EXTENSION] {
|
||||
let path = directory.join(format!("{session_id}.{extension}"));
|
||||
if path.exists() {
|
||||
return Ok(path);
|
||||
}
|
||||
}
|
||||
Err(SessionControlError::Format(
|
||||
format_missing_session_reference(session_id),
|
||||
))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_managed_session_file(path: &Path) -> bool {
|
||||
path.extension()
|
||||
.and_then(|ext| ext.to_str())
|
||||
.is_some_and(|extension| {
|
||||
extension == PRIMARY_SESSION_EXTENSION || extension == LEGACY_SESSION_EXTENSION
|
||||
})
|
||||
}
|
||||
|
||||
pub fn list_managed_sessions() -> Result<Vec<ManagedSessionSummary>, SessionControlError> {
|
||||
list_managed_sessions_for(env::current_dir()?)
|
||||
}
|
||||
|
||||
pub fn list_managed_sessions_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
) -> Result<Vec<ManagedSessionSummary>, SessionControlError> {
|
||||
let mut sessions = Vec::new();
|
||||
for entry in fs::read_dir(managed_sessions_dir_for(base_dir)?)? {
|
||||
let entry = entry?;
|
||||
let path = entry.path();
|
||||
if !is_managed_session_file(&path) {
|
||||
continue;
|
||||
}
|
||||
let metadata = entry.metadata()?;
|
||||
let modified_epoch_millis = metadata
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|time| time.duration_since(UNIX_EPOCH).ok())
|
||||
.map(|duration| duration.as_millis())
|
||||
.unwrap_or_default();
|
||||
let (id, message_count, parent_session_id, branch_name) =
|
||||
match Session::load_from_path(&path) {
|
||||
Ok(session) => {
|
||||
let parent_session_id = session
|
||||
.fork
|
||||
.as_ref()
|
||||
.map(|fork| fork.parent_session_id.clone());
|
||||
let branch_name = session
|
||||
.fork
|
||||
.as_ref()
|
||||
.and_then(|fork| fork.branch_name.clone());
|
||||
(
|
||||
session.session_id,
|
||||
session.messages.len(),
|
||||
parent_session_id,
|
||||
branch_name,
|
||||
)
|
||||
}
|
||||
Err(_) => (
|
||||
path.file_stem()
|
||||
.and_then(|value| value.to_str())
|
||||
.unwrap_or("unknown")
|
||||
.to_string(),
|
||||
0,
|
||||
None,
|
||||
None,
|
||||
),
|
||||
};
|
||||
sessions.push(ManagedSessionSummary {
|
||||
id,
|
||||
path,
|
||||
modified_epoch_millis,
|
||||
message_count,
|
||||
parent_session_id,
|
||||
branch_name,
|
||||
});
|
||||
}
|
||||
sessions.sort_by(|left, right| {
|
||||
right
|
||||
.modified_epoch_millis
|
||||
.cmp(&left.modified_epoch_millis)
|
||||
.then_with(|| right.id.cmp(&left.id))
|
||||
});
|
||||
Ok(sessions)
|
||||
}
|
||||
|
||||
pub fn latest_managed_session() -> Result<ManagedSessionSummary, SessionControlError> {
|
||||
latest_managed_session_for(env::current_dir()?)
|
||||
}
|
||||
|
||||
pub fn latest_managed_session_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
) -> Result<ManagedSessionSummary, SessionControlError> {
|
||||
list_managed_sessions_for(base_dir)?
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| SessionControlError::Format(format_no_managed_sessions()))
|
||||
}
|
||||
|
||||
pub fn load_managed_session(reference: &str) -> Result<LoadedManagedSession, SessionControlError> {
|
||||
load_managed_session_for(env::current_dir()?, reference)
|
||||
}
|
||||
|
||||
pub fn load_managed_session_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
reference: &str,
|
||||
) -> Result<LoadedManagedSession, SessionControlError> {
|
||||
let handle = resolve_session_reference_for(base_dir, reference)?;
|
||||
let session = Session::load_from_path(&handle.path)?;
|
||||
Ok(LoadedManagedSession {
|
||||
handle: SessionHandle {
|
||||
id: session.session_id.clone(),
|
||||
path: handle.path,
|
||||
},
|
||||
session,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn fork_managed_session(
|
||||
session: &Session,
|
||||
branch_name: Option<String>,
|
||||
) -> Result<ForkedManagedSession, SessionControlError> {
|
||||
fork_managed_session_for(env::current_dir()?, session, branch_name)
|
||||
}
|
||||
|
||||
pub fn fork_managed_session_for(
|
||||
base_dir: impl AsRef<Path>,
|
||||
session: &Session,
|
||||
branch_name: Option<String>,
|
||||
) -> Result<ForkedManagedSession, SessionControlError> {
|
||||
let parent_session_id = session.session_id.clone();
|
||||
let forked = session.fork(branch_name);
|
||||
let handle = create_managed_session_handle_for(base_dir, &forked.session_id)?;
|
||||
let branch_name = forked
|
||||
.fork
|
||||
.as_ref()
|
||||
.and_then(|fork| fork.branch_name.clone());
|
||||
let forked = forked.with_persistence_path(handle.path.clone());
|
||||
forked.save_to_path(&handle.path)?;
|
||||
Ok(ForkedManagedSession {
|
||||
parent_session_id,
|
||||
handle,
|
||||
session: forked,
|
||||
branch_name,
|
||||
})
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_session_reference_alias(reference: &str) -> bool {
|
||||
SESSION_REFERENCE_ALIASES
|
||||
.iter()
|
||||
.any(|alias| reference.eq_ignore_ascii_case(alias))
|
||||
}
|
||||
|
||||
fn session_id_from_path(path: &Path) -> Option<String> {
|
||||
path.file_name()
|
||||
.and_then(|value| value.to_str())
|
||||
.and_then(|name| {
|
||||
name.strip_suffix(&format!(".{PRIMARY_SESSION_EXTENSION}"))
|
||||
.or_else(|| name.strip_suffix(&format!(".{LEGACY_SESSION_EXTENSION}")))
|
||||
})
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn format_missing_session_reference(reference: &str) -> String {
|
||||
format!(
|
||||
"session not found: {reference}\nHint: managed sessions live in .claw/sessions/. Try `{LATEST_SESSION_REFERENCE}` for the most recent session or `/session list` in the REPL."
|
||||
)
|
||||
}
|
||||
|
||||
fn format_no_managed_sessions() -> String {
|
||||
format!(
|
||||
"no managed sessions found in .claw/sessions/\nStart `claw` to create a session, then rerun with `--resume {LATEST_SESSION_REFERENCE}`."
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
create_managed_session_handle_for, fork_managed_session_for, is_session_reference_alias,
|
||||
list_managed_sessions_for, load_managed_session_for, resolve_session_reference_for,
|
||||
workspace_fingerprint, ManagedSessionSummary, SessionStore, LATEST_SESSION_REFERENCE,
|
||||
};
|
||||
use crate::session::Session;
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
fn temp_dir() -> PathBuf {
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_nanos();
|
||||
std::env::temp_dir().join(format!("runtime-session-control-{nanos}"))
|
||||
}
|
||||
|
||||
fn persist_session(root: &Path, text: &str) -> Session {
|
||||
let mut session = Session::new();
|
||||
session
|
||||
.push_user_text(text)
|
||||
.expect("session message should save");
|
||||
let handle = create_managed_session_handle_for(root, &session.session_id)
|
||||
.expect("managed session handle should build");
|
||||
let session = session.with_persistence_path(handle.path.clone());
|
||||
session
|
||||
.save_to_path(&handle.path)
|
||||
.expect("session should persist");
|
||||
session
|
||||
}
|
||||
|
||||
fn wait_for_next_millisecond() {
|
||||
let start = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_millis();
|
||||
while SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.expect("time should be after epoch")
|
||||
.as_millis()
|
||||
<= start
|
||||
{}
|
||||
}
|
||||
|
||||
fn summary_by_id<'a>(
|
||||
summaries: &'a [ManagedSessionSummary],
|
||||
id: &str,
|
||||
) -> &'a ManagedSessionSummary {
|
||||
summaries
|
||||
.iter()
|
||||
.find(|summary| summary.id == id)
|
||||
.expect("session summary should exist")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn creates_and_lists_managed_sessions() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir should exist");
|
||||
let older = persist_session(&root, "older session");
|
||||
wait_for_next_millisecond();
|
||||
let newer = persist_session(&root, "newer session");
|
||||
|
||||
// when
|
||||
let sessions = list_managed_sessions_for(&root).expect("managed sessions should list");
|
||||
|
||||
// then
|
||||
assert_eq!(sessions.len(), 2);
|
||||
assert_eq!(sessions[0].id, newer.session_id);
|
||||
assert_eq!(summary_by_id(&sessions, &older.session_id).message_count, 1);
|
||||
assert_eq!(summary_by_id(&sessions, &newer.session_id).message_count, 1);
|
||||
fs::remove_dir_all(root).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolves_latest_alias_and_loads_session_from_workspace_root() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir should exist");
|
||||
let older = persist_session(&root, "older session");
|
||||
wait_for_next_millisecond();
|
||||
let newer = persist_session(&root, "newer session");
|
||||
|
||||
// when
|
||||
let handle = resolve_session_reference_for(&root, LATEST_SESSION_REFERENCE)
|
||||
.expect("latest alias should resolve");
|
||||
let loaded = load_managed_session_for(&root, "recent")
|
||||
.expect("recent alias should load the latest session");
|
||||
|
||||
// then
|
||||
assert_eq!(handle.id, newer.session_id);
|
||||
assert_eq!(loaded.handle.id, newer.session_id);
|
||||
assert_eq!(loaded.session.messages.len(), 1);
|
||||
assert_ne!(loaded.handle.id, older.session_id);
|
||||
assert!(is_session_reference_alias("last"));
|
||||
fs::remove_dir_all(root).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forks_session_into_managed_storage_with_lineage() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("root dir should exist");
|
||||
let source = persist_session(&root, "parent session");
|
||||
|
||||
// when
|
||||
let forked = fork_managed_session_for(&root, &source, Some("incident-review".to_string()))
|
||||
.expect("session should fork");
|
||||
let sessions = list_managed_sessions_for(&root).expect("managed sessions should list");
|
||||
let summary = summary_by_id(&sessions, &forked.handle.id);
|
||||
|
||||
// then
|
||||
assert_eq!(forked.parent_session_id, source.session_id);
|
||||
assert_eq!(forked.branch_name.as_deref(), Some("incident-review"));
|
||||
assert_eq!(
|
||||
summary.parent_session_id.as_deref(),
|
||||
Some(source.session_id.as_str())
|
||||
);
|
||||
assert_eq!(summary.branch_name.as_deref(), Some("incident-review"));
|
||||
assert_eq!(
|
||||
forked.session.persistence_path(),
|
||||
Some(forked.handle.path.as_path())
|
||||
);
|
||||
fs::remove_dir_all(root).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// Per-worktree session isolation (SessionStore) tests
|
||||
// ------------------------------------------------------------------
|
||||
|
||||
fn persist_session_via_store(store: &SessionStore, text: &str) -> Session {
|
||||
let mut session = Session::new();
|
||||
session
|
||||
.push_user_text(text)
|
||||
.expect("session message should save");
|
||||
let handle = store.create_handle(&session.session_id);
|
||||
let session = session.with_persistence_path(handle.path.clone());
|
||||
session
|
||||
.save_to_path(&handle.path)
|
||||
.expect("session should persist");
|
||||
session
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workspace_fingerprint_is_deterministic_and_differs_per_path() {
|
||||
// given
|
||||
let path_a = Path::new("/tmp/worktree-alpha");
|
||||
let path_b = Path::new("/tmp/worktree-beta");
|
||||
|
||||
// when
|
||||
let fp_a1 = workspace_fingerprint(path_a);
|
||||
let fp_a2 = workspace_fingerprint(path_a);
|
||||
let fp_b = workspace_fingerprint(path_b);
|
||||
|
||||
// then
|
||||
assert_eq!(fp_a1, fp_a2, "same path must produce the same fingerprint");
|
||||
assert_ne!(
|
||||
fp_a1, fp_b,
|
||||
"different paths must produce different fingerprints"
|
||||
);
|
||||
assert_eq!(fp_a1.len(), 16, "fingerprint must be a 16-char hex string");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_from_cwd_isolates_sessions_by_workspace() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
let workspace_a = base.join("repo-alpha");
|
||||
let workspace_b = base.join("repo-beta");
|
||||
fs::create_dir_all(&workspace_a).expect("workspace a should exist");
|
||||
fs::create_dir_all(&workspace_b).expect("workspace b should exist");
|
||||
|
||||
let store_a = SessionStore::from_cwd(&workspace_a).expect("store a should build");
|
||||
let store_b = SessionStore::from_cwd(&workspace_b).expect("store b should build");
|
||||
|
||||
// when
|
||||
let session_a = persist_session_via_store(&store_a, "alpha work");
|
||||
let _session_b = persist_session_via_store(&store_b, "beta work");
|
||||
|
||||
// then — each store only sees its own sessions
|
||||
let list_a = store_a.list_sessions().expect("list a");
|
||||
let list_b = store_b.list_sessions().expect("list b");
|
||||
assert_eq!(list_a.len(), 1, "store a should see exactly one session");
|
||||
assert_eq!(list_b.len(), 1, "store b should see exactly one session");
|
||||
assert_eq!(list_a[0].id, session_a.session_id);
|
||||
assert_ne!(
|
||||
store_a.sessions_dir(),
|
||||
store_b.sessions_dir(),
|
||||
"session directories must differ across workspaces"
|
||||
);
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_from_data_dir_namespaces_by_workspace() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
let data_dir = base.join("global-data");
|
||||
let workspace_a = PathBuf::from("/tmp/project-one");
|
||||
let workspace_b = PathBuf::from("/tmp/project-two");
|
||||
fs::create_dir_all(&data_dir).expect("data dir should exist");
|
||||
|
||||
let store_a =
|
||||
SessionStore::from_data_dir(&data_dir, &workspace_a).expect("store a should build");
|
||||
let store_b =
|
||||
SessionStore::from_data_dir(&data_dir, &workspace_b).expect("store b should build");
|
||||
|
||||
// when
|
||||
persist_session_via_store(&store_a, "work in project-one");
|
||||
persist_session_via_store(&store_b, "work in project-two");
|
||||
|
||||
// then
|
||||
assert_ne!(
|
||||
store_a.sessions_dir(),
|
||||
store_b.sessions_dir(),
|
||||
"data-dir stores must namespace by workspace"
|
||||
);
|
||||
assert_eq!(store_a.list_sessions().expect("list a").len(), 1);
|
||||
assert_eq!(store_b.list_sessions().expect("list b").len(), 1);
|
||||
assert_eq!(store_a.workspace_root(), workspace_a.as_path());
|
||||
assert_eq!(store_b.workspace_root(), workspace_b.as_path());
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_create_and_load_round_trip() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
fs::create_dir_all(&base).expect("base dir should exist");
|
||||
let store = SessionStore::from_cwd(&base).expect("store should build");
|
||||
let session = persist_session_via_store(&store, "round-trip message");
|
||||
|
||||
// when
|
||||
let loaded = store
|
||||
.load_session(&session.session_id)
|
||||
.expect("session should load via store");
|
||||
|
||||
// then
|
||||
assert_eq!(loaded.handle.id, session.session_id);
|
||||
assert_eq!(loaded.session.messages.len(), 1);
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_latest_and_resolve_reference() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
fs::create_dir_all(&base).expect("base dir should exist");
|
||||
let store = SessionStore::from_cwd(&base).expect("store should build");
|
||||
let _older = persist_session_via_store(&store, "older");
|
||||
wait_for_next_millisecond();
|
||||
let newer = persist_session_via_store(&store, "newer");
|
||||
|
||||
// when
|
||||
let latest = store.latest_session().expect("latest should resolve");
|
||||
let handle = store
|
||||
.resolve_reference("latest")
|
||||
.expect("latest alias should resolve");
|
||||
|
||||
// then
|
||||
assert_eq!(latest.id, newer.session_id);
|
||||
assert_eq!(handle.id, newer.session_id);
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_store_fork_stays_in_same_namespace() {
|
||||
// given
|
||||
let base = temp_dir();
|
||||
fs::create_dir_all(&base).expect("base dir should exist");
|
||||
let store = SessionStore::from_cwd(&base).expect("store should build");
|
||||
let source = persist_session_via_store(&store, "parent work");
|
||||
|
||||
// when
|
||||
let forked = store
|
||||
.fork_session(&source, Some("bugfix".to_string()))
|
||||
.expect("fork should succeed");
|
||||
let sessions = store.list_sessions().expect("list sessions");
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
sessions.len(),
|
||||
2,
|
||||
"forked session must land in the same namespace"
|
||||
);
|
||||
assert_eq!(forked.parent_session_id, source.session_id);
|
||||
assert_eq!(forked.branch_name.as_deref(), Some("bugfix"));
|
||||
assert!(
|
||||
forked.handle.path.starts_with(store.sessions_dir()),
|
||||
"forked session path must be inside the store namespace"
|
||||
);
|
||||
fs::remove_dir_all(base).expect("temp dir should clean up");
|
||||
}
|
||||
}
|
||||
@@ -1,429 +0,0 @@
|
||||
#![allow(clippy::must_use_candidate)]
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
/// Outcome of comparing the worktree HEAD against the expected base commit.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum BaseCommitState {
|
||||
/// HEAD matches the expected base commit.
|
||||
Matches,
|
||||
/// HEAD has diverged from the expected base.
|
||||
Diverged { expected: String, actual: String },
|
||||
/// No expected base was supplied (neither flag nor file).
|
||||
NoExpectedBase,
|
||||
/// The working directory is not inside a git repository.
|
||||
NotAGitRepo,
|
||||
}
|
||||
|
||||
/// Where the expected base commit originated from.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum BaseCommitSource {
|
||||
Flag(String),
|
||||
File(String),
|
||||
}
|
||||
|
||||
/// Read the `.claw-base` file from the given directory and return the trimmed
|
||||
/// commit hash, or `None` when the file is absent or empty.
|
||||
pub fn read_claw_base_file(cwd: &Path) -> Option<String> {
|
||||
let path = cwd.join(".claw-base");
|
||||
let content = std::fs::read_to_string(path).ok()?;
|
||||
let trimmed = content.trim();
|
||||
if trimmed.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(trimmed.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the expected base commit: prefer the `--base-commit` flag value,
|
||||
/// fall back to reading `.claw-base` from `cwd`.
|
||||
pub fn resolve_expected_base(flag_value: Option<&str>, cwd: &Path) -> Option<BaseCommitSource> {
|
||||
if let Some(value) = flag_value {
|
||||
let trimmed = value.trim();
|
||||
if !trimmed.is_empty() {
|
||||
return Some(BaseCommitSource::Flag(trimmed.to_string()));
|
||||
}
|
||||
}
|
||||
read_claw_base_file(cwd).map(BaseCommitSource::File)
|
||||
}
|
||||
|
||||
/// Verify that the worktree HEAD matches `expected_base`.
|
||||
///
|
||||
/// Returns [`BaseCommitState::NoExpectedBase`] when no expected commit is
|
||||
/// provided (the check is effectively a no-op in that case).
|
||||
pub fn check_base_commit(cwd: &Path, expected_base: Option<&BaseCommitSource>) -> BaseCommitState {
|
||||
let Some(source) = expected_base else {
|
||||
return BaseCommitState::NoExpectedBase;
|
||||
};
|
||||
let expected_raw = match source {
|
||||
BaseCommitSource::Flag(value) | BaseCommitSource::File(value) => value.as_str(),
|
||||
};
|
||||
|
||||
let Some(head_sha) = resolve_head_sha(cwd) else {
|
||||
return BaseCommitState::NotAGitRepo;
|
||||
};
|
||||
|
||||
let Some(expected_sha) = resolve_rev(cwd, expected_raw) else {
|
||||
// If the expected ref cannot be resolved, compare raw strings as a
|
||||
// best-effort fallback (e.g. partial SHA provided by the caller).
|
||||
return if head_sha.starts_with(expected_raw) || expected_raw.starts_with(&head_sha) {
|
||||
BaseCommitState::Matches
|
||||
} else {
|
||||
BaseCommitState::Diverged {
|
||||
expected: expected_raw.to_string(),
|
||||
actual: head_sha,
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
if head_sha == expected_sha {
|
||||
BaseCommitState::Matches
|
||||
} else {
|
||||
BaseCommitState::Diverged {
|
||||
expected: expected_sha,
|
||||
actual: head_sha,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Format a human-readable warning when the base commit has diverged.
|
||||
///
|
||||
/// Returns `None` for non-warning states (`Matches`, `NoExpectedBase`).
|
||||
pub fn format_stale_base_warning(state: &BaseCommitState) -> Option<String> {
|
||||
match state {
|
||||
BaseCommitState::Diverged { expected, actual } => Some(format!(
|
||||
"warning: worktree HEAD ({actual}) does not match expected base commit ({expected}). \
|
||||
Session may run against a stale codebase."
|
||||
)),
|
||||
BaseCommitState::NotAGitRepo => {
|
||||
Some("warning: stale-base check skipped — not inside a git repository.".to_string())
|
||||
}
|
||||
BaseCommitState::Matches | BaseCommitState::NoExpectedBase => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_head_sha(cwd: &Path) -> Option<String> {
|
||||
resolve_rev(cwd, "HEAD")
|
||||
}
|
||||
|
||||
fn resolve_rev(cwd: &Path, rev: &str) -> Option<String> {
|
||||
let output = Command::new("git")
|
||||
.args(["rev-parse", rev])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
.ok()?;
|
||||
if !output.status.success() {
|
||||
return None;
|
||||
}
|
||||
let sha = String::from_utf8(output.stdout).ok()?;
|
||||
let trimmed = sha.trim();
|
||||
if trimmed.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(trimmed.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
use std::process::Command;
|
||||
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!("runtime-stale-base-{nanos}"))
|
||||
}
|
||||
|
||||
fn init_repo(path: &std::path::Path) {
|
||||
fs::create_dir_all(path).expect("create repo dir");
|
||||
run(path, &["init", "--quiet", "-b", "main"]);
|
||||
run(path, &["config", "user.email", "tests@example.com"]);
|
||||
run(path, &["config", "user.name", "Stale Base Tests"]);
|
||||
fs::write(path.join("init.txt"), "initial\n").expect("write init file");
|
||||
run(path, &["add", "."]);
|
||||
run(path, &["commit", "-m", "initial commit", "--quiet"]);
|
||||
}
|
||||
|
||||
fn run(cwd: &std::path::Path, args: &[&str]) {
|
||||
let status = Command::new("git")
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.status()
|
||||
.unwrap_or_else(|e| panic!("git {} failed to execute: {e}", args.join(" ")));
|
||||
assert!(
|
||||
status.success(),
|
||||
"git {} exited with {status}",
|
||||
args.join(" ")
|
||||
);
|
||||
}
|
||||
|
||||
fn commit_file(repo: &std::path::Path, name: &str, msg: &str) {
|
||||
fs::write(repo.join(name), format!("{msg}\n")).expect("write file");
|
||||
run(repo, &["add", name]);
|
||||
run(repo, &["commit", "-m", msg, "--quiet"]);
|
||||
}
|
||||
|
||||
fn head_sha(repo: &std::path::Path) -> String {
|
||||
let output = Command::new("git")
|
||||
.args(["rev-parse", "HEAD"])
|
||||
.current_dir(repo)
|
||||
.output()
|
||||
.expect("git rev-parse HEAD");
|
||||
String::from_utf8(output.stdout)
|
||||
.expect("valid utf8")
|
||||
.trim()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_when_head_equals_expected_base() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
let sha = head_sha(&root);
|
||||
let source = BaseCommitSource::Flag(sha);
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, Some(&source));
|
||||
|
||||
// then
|
||||
assert_eq!(state, BaseCommitState::Matches);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diverged_when_head_moved_past_expected_base() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
let old_sha = head_sha(&root);
|
||||
commit_file(&root, "extra.txt", "move head forward");
|
||||
let new_sha = head_sha(&root);
|
||||
let source = BaseCommitSource::Flag(old_sha.clone());
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, Some(&source));
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
state,
|
||||
BaseCommitState::Diverged {
|
||||
expected: old_sha,
|
||||
actual: new_sha,
|
||||
}
|
||||
);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_expected_base_when_source_is_none() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, None);
|
||||
|
||||
// then
|
||||
assert_eq!(state, BaseCommitState::NoExpectedBase);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn not_a_git_repo_when_outside_repo() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
let source = BaseCommitSource::Flag("abc1234".to_string());
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, Some(&source));
|
||||
|
||||
// then
|
||||
assert_eq!(state, BaseCommitState::NotAGitRepo);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_claw_base_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
fs::write(root.join(".claw-base"), "abc1234def5678\n").expect("write .claw-base");
|
||||
|
||||
// when
|
||||
let value = read_claw_base_file(&root);
|
||||
|
||||
// then
|
||||
assert_eq!(value, Some("abc1234def5678".to_string()));
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_for_missing_claw_base_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
|
||||
// when
|
||||
let value = read_claw_base_file(&root);
|
||||
|
||||
// then
|
||||
assert!(value.is_none());
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_for_empty_claw_base_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
fs::write(root.join(".claw-base"), " \n").expect("write empty .claw-base");
|
||||
|
||||
// when
|
||||
let value = read_claw_base_file(&root);
|
||||
|
||||
// then
|
||||
assert!(value.is_none());
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_expected_base_prefers_flag_over_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
fs::write(root.join(".claw-base"), "from_file\n").expect("write .claw-base");
|
||||
|
||||
// when
|
||||
let source = resolve_expected_base(Some("from_flag"), &root);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
source,
|
||||
Some(BaseCommitSource::Flag("from_flag".to_string()))
|
||||
);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_expected_base_falls_back_to_file() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
fs::write(root.join(".claw-base"), "from_file\n").expect("write .claw-base");
|
||||
|
||||
// when
|
||||
let source = resolve_expected_base(None, &root);
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
source,
|
||||
Some(BaseCommitSource::File("from_file".to_string()))
|
||||
);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_expected_base_returns_none_when_nothing_available() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
fs::create_dir_all(&root).expect("create dir");
|
||||
|
||||
// when
|
||||
let source = resolve_expected_base(None, &root);
|
||||
|
||||
// then
|
||||
assert!(source.is_none());
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_warning_returns_message_for_diverged() {
|
||||
// given
|
||||
let state = BaseCommitState::Diverged {
|
||||
expected: "abc1234".to_string(),
|
||||
actual: "def5678".to_string(),
|
||||
};
|
||||
|
||||
// when
|
||||
let warning = format_stale_base_warning(&state);
|
||||
|
||||
// then
|
||||
let message = warning.expect("should produce warning");
|
||||
assert!(message.contains("abc1234"));
|
||||
assert!(message.contains("def5678"));
|
||||
assert!(message.contains("stale codebase"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_warning_returns_none_for_matches() {
|
||||
// given
|
||||
let state = BaseCommitState::Matches;
|
||||
|
||||
// when
|
||||
let warning = format_stale_base_warning(&state);
|
||||
|
||||
// then
|
||||
assert!(warning.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn format_warning_returns_none_for_no_expected_base() {
|
||||
// given
|
||||
let state = BaseCommitState::NoExpectedBase;
|
||||
|
||||
// when
|
||||
let warning = format_stale_base_warning(&state);
|
||||
|
||||
// then
|
||||
assert!(warning.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matches_with_claw_base_file_in_real_repo() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
let sha = head_sha(&root);
|
||||
fs::write(root.join(".claw-base"), format!("{sha}\n")).expect("write .claw-base");
|
||||
let source = resolve_expected_base(None, &root);
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, source.as_ref());
|
||||
|
||||
// then
|
||||
assert_eq!(state, BaseCommitState::Matches);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diverged_with_claw_base_file_after_new_commit() {
|
||||
// given
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
let old_sha = head_sha(&root);
|
||||
fs::write(root.join(".claw-base"), format!("{old_sha}\n")).expect("write .claw-base");
|
||||
commit_file(&root, "new.txt", "advance head");
|
||||
let new_sha = head_sha(&root);
|
||||
let source = resolve_expected_base(None, &root);
|
||||
|
||||
// when
|
||||
let state = check_base_commit(&root, source.as_ref());
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
state,
|
||||
BaseCommitState::Diverged {
|
||||
expected: old_sha,
|
||||
actual: new_sha,
|
||||
}
|
||||
);
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
}
|
||||
@@ -1,417 +0,0 @@
|
||||
#![allow(clippy::must_use_candidate)]
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum BranchFreshness {
|
||||
Fresh,
|
||||
Stale {
|
||||
commits_behind: usize,
|
||||
missing_fixes: Vec<String>,
|
||||
},
|
||||
Diverged {
|
||||
ahead: usize,
|
||||
behind: usize,
|
||||
missing_fixes: Vec<String>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum StaleBranchPolicy {
|
||||
AutoRebase,
|
||||
AutoMergeForward,
|
||||
WarnOnly,
|
||||
Block,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum StaleBranchEvent {
|
||||
BranchStaleAgainstMain {
|
||||
branch: String,
|
||||
commits_behind: usize,
|
||||
missing_fixes: Vec<String>,
|
||||
},
|
||||
RebaseAttempted {
|
||||
branch: String,
|
||||
result: String,
|
||||
},
|
||||
MergeForwardAttempted {
|
||||
branch: String,
|
||||
result: String,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum StaleBranchAction {
|
||||
Noop,
|
||||
Warn { message: String },
|
||||
Block { message: String },
|
||||
Rebase,
|
||||
MergeForward,
|
||||
}
|
||||
|
||||
pub fn check_freshness(branch: &str, main_ref: &str) -> BranchFreshness {
|
||||
check_freshness_in(branch, main_ref, Path::new("."))
|
||||
}
|
||||
|
||||
pub fn apply_policy(freshness: &BranchFreshness, policy: StaleBranchPolicy) -> StaleBranchAction {
|
||||
match freshness {
|
||||
BranchFreshness::Fresh => StaleBranchAction::Noop,
|
||||
BranchFreshness::Stale {
|
||||
commits_behind,
|
||||
missing_fixes,
|
||||
} => match policy {
|
||||
StaleBranchPolicy::WarnOnly => StaleBranchAction::Warn {
|
||||
message: format!(
|
||||
"Branch is {commits_behind} commit(s) behind main. Missing fixes: {}",
|
||||
if missing_fixes.is_empty() {
|
||||
"(none)".to_string()
|
||||
} else {
|
||||
missing_fixes.join("; ")
|
||||
}
|
||||
),
|
||||
},
|
||||
StaleBranchPolicy::Block => StaleBranchAction::Block {
|
||||
message: format!(
|
||||
"Branch is {commits_behind} commit(s) behind main and must be updated before proceeding."
|
||||
),
|
||||
},
|
||||
StaleBranchPolicy::AutoRebase => StaleBranchAction::Rebase,
|
||||
StaleBranchPolicy::AutoMergeForward => StaleBranchAction::MergeForward,
|
||||
},
|
||||
BranchFreshness::Diverged {
|
||||
ahead,
|
||||
behind,
|
||||
missing_fixes,
|
||||
} => match policy {
|
||||
StaleBranchPolicy::WarnOnly => StaleBranchAction::Warn {
|
||||
message: format!(
|
||||
"Branch has diverged: {ahead} commit(s) ahead, {behind} commit(s) behind main. Missing fixes: {}",
|
||||
format_missing_fixes(missing_fixes)
|
||||
),
|
||||
},
|
||||
StaleBranchPolicy::Block => StaleBranchAction::Block {
|
||||
message: format!(
|
||||
"Branch has diverged ({ahead} ahead, {behind} behind) and must be reconciled before proceeding. Missing fixes: {}",
|
||||
format_missing_fixes(missing_fixes)
|
||||
),
|
||||
},
|
||||
StaleBranchPolicy::AutoRebase => StaleBranchAction::Rebase,
|
||||
StaleBranchPolicy::AutoMergeForward => StaleBranchAction::MergeForward,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn check_freshness_in(
|
||||
branch: &str,
|
||||
main_ref: &str,
|
||||
repo_path: &Path,
|
||||
) -> BranchFreshness {
|
||||
let behind = rev_list_count(main_ref, branch, repo_path);
|
||||
let ahead = rev_list_count(branch, main_ref, repo_path);
|
||||
|
||||
if behind == 0 {
|
||||
return BranchFreshness::Fresh;
|
||||
}
|
||||
|
||||
if ahead > 0 {
|
||||
return BranchFreshness::Diverged {
|
||||
ahead,
|
||||
behind,
|
||||
missing_fixes: missing_fix_subjects(main_ref, branch, repo_path),
|
||||
};
|
||||
}
|
||||
|
||||
let missing_fixes = missing_fix_subjects(main_ref, branch, repo_path);
|
||||
BranchFreshness::Stale {
|
||||
commits_behind: behind,
|
||||
missing_fixes,
|
||||
}
|
||||
}
|
||||
|
||||
fn format_missing_fixes(missing_fixes: &[String]) -> String {
|
||||
if missing_fixes.is_empty() {
|
||||
"(none)".to_string()
|
||||
} else {
|
||||
missing_fixes.join("; ")
|
||||
}
|
||||
}
|
||||
|
||||
fn rev_list_count(a: &str, b: &str, repo_path: &Path) -> usize {
|
||||
let output = Command::new("git")
|
||||
.args(["rev-list", "--count", &format!("{b}..{a}")])
|
||||
.current_dir(repo_path)
|
||||
.output();
|
||||
match output {
|
||||
Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout)
|
||||
.trim()
|
||||
.parse::<usize>()
|
||||
.unwrap_or(0),
|
||||
_ => 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn missing_fix_subjects(a: &str, b: &str, repo_path: &Path) -> Vec<String> {
|
||||
let output = Command::new("git")
|
||||
.args(["log", "--format=%s", &format!("{b}..{a}")])
|
||||
.current_dir(repo_path)
|
||||
.output();
|
||||
match output {
|
||||
Ok(o) if o.status.success() => String::from_utf8_lossy(&o.stdout)
|
||||
.lines()
|
||||
.filter(|l| !l.is_empty())
|
||||
.map(String::from)
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs;
|
||||
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!("runtime-stale-branch-{nanos}"))
|
||||
}
|
||||
|
||||
fn init_repo(path: &Path) {
|
||||
fs::create_dir_all(path).expect("create repo dir");
|
||||
run(path, &["init", "--quiet", "-b", "main"]);
|
||||
run(path, &["config", "user.email", "tests@example.com"]);
|
||||
run(path, &["config", "user.name", "Stale Branch Tests"]);
|
||||
fs::write(path.join("init.txt"), "initial\n").expect("write init file");
|
||||
run(path, &["add", "."]);
|
||||
run(path, &["commit", "-m", "initial commit", "--quiet"]);
|
||||
}
|
||||
|
||||
fn run(cwd: &Path, args: &[&str]) {
|
||||
let status = Command::new("git")
|
||||
.args(args)
|
||||
.current_dir(cwd)
|
||||
.status()
|
||||
.unwrap_or_else(|e| panic!("git {} failed to execute: {e}", args.join(" ")));
|
||||
assert!(
|
||||
status.success(),
|
||||
"git {} exited with {status}",
|
||||
args.join(" ")
|
||||
);
|
||||
}
|
||||
|
||||
fn commit_file(repo: &Path, name: &str, msg: &str) {
|
||||
fs::write(repo.join(name), format!("{msg}\n")).expect("write file");
|
||||
run(repo, &["add", name]);
|
||||
run(repo, &["commit", "-m", msg, "--quiet"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fresh_branch_passes() {
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// given
|
||||
run(&root, &["checkout", "-b", "topic"]);
|
||||
|
||||
// when
|
||||
let freshness = check_freshness_in("topic", "main", &root);
|
||||
|
||||
// then
|
||||
assert_eq!(freshness, BranchFreshness::Fresh);
|
||||
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fresh_branch_ahead_of_main_still_fresh() {
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// given
|
||||
run(&root, &["checkout", "-b", "topic"]);
|
||||
commit_file(&root, "feature.txt", "add feature");
|
||||
|
||||
// when
|
||||
let freshness = check_freshness_in("topic", "main", &root);
|
||||
|
||||
// then
|
||||
assert_eq!(freshness, BranchFreshness::Fresh);
|
||||
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_branch_detected_with_correct_behind_count_and_missing_fixes() {
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// given
|
||||
run(&root, &["checkout", "-b", "topic"]);
|
||||
run(&root, &["checkout", "main"]);
|
||||
commit_file(&root, "fix1.txt", "fix: resolve timeout");
|
||||
commit_file(&root, "fix2.txt", "fix: handle null pointer");
|
||||
|
||||
// when
|
||||
let freshness = check_freshness_in("topic", "main", &root);
|
||||
|
||||
// then
|
||||
match freshness {
|
||||
BranchFreshness::Stale {
|
||||
commits_behind,
|
||||
missing_fixes,
|
||||
} => {
|
||||
assert_eq!(commits_behind, 2);
|
||||
assert_eq!(missing_fixes.len(), 2);
|
||||
assert_eq!(missing_fixes[0], "fix: handle null pointer");
|
||||
assert_eq!(missing_fixes[1], "fix: resolve timeout");
|
||||
}
|
||||
other => panic!("expected Stale, got {other:?}"),
|
||||
}
|
||||
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn diverged_branch_detection() {
|
||||
let root = temp_dir();
|
||||
init_repo(&root);
|
||||
|
||||
// given
|
||||
run(&root, &["checkout", "-b", "topic"]);
|
||||
commit_file(&root, "topic_work.txt", "topic work");
|
||||
run(&root, &["checkout", "main"]);
|
||||
commit_file(&root, "main_fix.txt", "main fix");
|
||||
|
||||
// when
|
||||
let freshness = check_freshness_in("topic", "main", &root);
|
||||
|
||||
// then
|
||||
match freshness {
|
||||
BranchFreshness::Diverged {
|
||||
ahead,
|
||||
behind,
|
||||
missing_fixes,
|
||||
} => {
|
||||
assert_eq!(ahead, 1);
|
||||
assert_eq!(behind, 1);
|
||||
assert_eq!(missing_fixes, vec!["main fix".to_string()]);
|
||||
}
|
||||
other => panic!("expected Diverged, got {other:?}"),
|
||||
}
|
||||
|
||||
fs::remove_dir_all(&root).expect("cleanup");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_noop_for_fresh_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Fresh;
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::WarnOnly);
|
||||
|
||||
// then
|
||||
assert_eq!(action, StaleBranchAction::Noop);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_warn_for_stale_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Stale {
|
||||
commits_behind: 3,
|
||||
missing_fixes: vec!["fix: timeout".into(), "fix: null ptr".into()],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::WarnOnly);
|
||||
|
||||
// then
|
||||
match action {
|
||||
StaleBranchAction::Warn { message } => {
|
||||
assert!(message.contains("3 commit(s) behind"));
|
||||
assert!(message.contains("fix: timeout"));
|
||||
assert!(message.contains("fix: null ptr"));
|
||||
}
|
||||
other => panic!("expected Warn, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_block_for_stale_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Stale {
|
||||
commits_behind: 1,
|
||||
missing_fixes: vec!["hotfix".into()],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::Block);
|
||||
|
||||
// then
|
||||
match action {
|
||||
StaleBranchAction::Block { message } => {
|
||||
assert!(message.contains("1 commit(s) behind"));
|
||||
}
|
||||
other => panic!("expected Block, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_auto_rebase_for_stale_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Stale {
|
||||
commits_behind: 2,
|
||||
missing_fixes: vec![],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::AutoRebase);
|
||||
|
||||
// then
|
||||
assert_eq!(action, StaleBranchAction::Rebase);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_auto_merge_forward_for_diverged_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Diverged {
|
||||
ahead: 5,
|
||||
behind: 2,
|
||||
missing_fixes: vec!["fix: merge main".into()],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::AutoMergeForward);
|
||||
|
||||
// then
|
||||
assert_eq!(action, StaleBranchAction::MergeForward);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn policy_warn_for_diverged_branch() {
|
||||
// given
|
||||
let freshness = BranchFreshness::Diverged {
|
||||
ahead: 3,
|
||||
behind: 1,
|
||||
missing_fixes: vec!["main hotfix".into()],
|
||||
};
|
||||
|
||||
// when
|
||||
let action = apply_policy(&freshness, StaleBranchPolicy::WarnOnly);
|
||||
|
||||
// then
|
||||
match action {
|
||||
StaleBranchAction::Warn { message } => {
|
||||
assert!(message.contains("diverged"));
|
||||
assert!(message.contains("3 commit(s) ahead"));
|
||||
assert!(message.contains("1 commit(s) behind"));
|
||||
assert!(message.contains("main hotfix"));
|
||||
}
|
||||
other => panic!("expected Warn, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,300 +0,0 @@
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
const DEFAULT_MAX_CHARS: usize = 1_200;
|
||||
const DEFAULT_MAX_LINES: usize = 24;
|
||||
const DEFAULT_MAX_LINE_CHARS: usize = 160;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct SummaryCompressionBudget {
|
||||
pub max_chars: usize,
|
||||
pub max_lines: usize,
|
||||
pub max_line_chars: usize,
|
||||
}
|
||||
|
||||
impl Default for SummaryCompressionBudget {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_chars: DEFAULT_MAX_CHARS,
|
||||
max_lines: DEFAULT_MAX_LINES,
|
||||
max_line_chars: DEFAULT_MAX_LINE_CHARS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SummaryCompressionResult {
|
||||
pub summary: String,
|
||||
pub original_chars: usize,
|
||||
pub compressed_chars: usize,
|
||||
pub original_lines: usize,
|
||||
pub compressed_lines: usize,
|
||||
pub removed_duplicate_lines: usize,
|
||||
pub omitted_lines: usize,
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn compress_summary(
|
||||
summary: &str,
|
||||
budget: SummaryCompressionBudget,
|
||||
) -> SummaryCompressionResult {
|
||||
let original_chars = summary.chars().count();
|
||||
let original_lines = summary.lines().count();
|
||||
|
||||
let normalized = normalize_lines(summary, budget.max_line_chars);
|
||||
if normalized.lines.is_empty() || budget.max_chars == 0 || budget.max_lines == 0 {
|
||||
return SummaryCompressionResult {
|
||||
summary: String::new(),
|
||||
original_chars,
|
||||
compressed_chars: 0,
|
||||
original_lines,
|
||||
compressed_lines: 0,
|
||||
removed_duplicate_lines: normalized.removed_duplicate_lines,
|
||||
omitted_lines: normalized.lines.len(),
|
||||
truncated: original_chars > 0,
|
||||
};
|
||||
}
|
||||
|
||||
let selected = select_line_indexes(&normalized.lines, budget);
|
||||
let mut compressed_lines = selected
|
||||
.iter()
|
||||
.map(|index| normalized.lines[*index].clone())
|
||||
.collect::<Vec<_>>();
|
||||
if compressed_lines.is_empty() {
|
||||
compressed_lines.push(truncate_line(&normalized.lines[0], budget.max_chars));
|
||||
}
|
||||
let omitted_lines = normalized
|
||||
.lines
|
||||
.len()
|
||||
.saturating_sub(compressed_lines.len());
|
||||
|
||||
if omitted_lines > 0 {
|
||||
let omission_notice = omission_notice(omitted_lines);
|
||||
push_line_with_budget(&mut compressed_lines, omission_notice, budget);
|
||||
}
|
||||
|
||||
let compressed_summary = compressed_lines.join("\n");
|
||||
|
||||
SummaryCompressionResult {
|
||||
compressed_chars: compressed_summary.chars().count(),
|
||||
compressed_lines: compressed_lines.len(),
|
||||
removed_duplicate_lines: normalized.removed_duplicate_lines,
|
||||
omitted_lines,
|
||||
truncated: compressed_summary != summary.trim(),
|
||||
summary: compressed_summary,
|
||||
original_chars,
|
||||
original_lines,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn compress_summary_text(summary: &str) -> String {
|
||||
compress_summary(summary, SummaryCompressionBudget::default()).summary
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct NormalizedSummary {
|
||||
lines: Vec<String>,
|
||||
removed_duplicate_lines: usize,
|
||||
}
|
||||
|
||||
fn normalize_lines(summary: &str, max_line_chars: usize) -> NormalizedSummary {
|
||||
let mut seen = BTreeSet::new();
|
||||
let mut lines = Vec::new();
|
||||
let mut removed_duplicate_lines = 0;
|
||||
|
||||
for raw_line in summary.lines() {
|
||||
let normalized = collapse_inline_whitespace(raw_line);
|
||||
if normalized.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let truncated = truncate_line(&normalized, max_line_chars);
|
||||
let dedupe_key = dedupe_key(&truncated);
|
||||
if !seen.insert(dedupe_key) {
|
||||
removed_duplicate_lines += 1;
|
||||
continue;
|
||||
}
|
||||
|
||||
lines.push(truncated);
|
||||
}
|
||||
|
||||
NormalizedSummary {
|
||||
lines,
|
||||
removed_duplicate_lines,
|
||||
}
|
||||
}
|
||||
|
||||
fn select_line_indexes(lines: &[String], budget: SummaryCompressionBudget) -> Vec<usize> {
|
||||
let mut selected = BTreeSet::<usize>::new();
|
||||
|
||||
for priority in 0..=3 {
|
||||
for (index, line) in lines.iter().enumerate() {
|
||||
if selected.contains(&index) || line_priority(line) != priority {
|
||||
continue;
|
||||
}
|
||||
|
||||
let candidate = selected
|
||||
.iter()
|
||||
.map(|selected_index| lines[*selected_index].as_str())
|
||||
.chain(std::iter::once(line.as_str()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if candidate.len() > budget.max_lines {
|
||||
continue;
|
||||
}
|
||||
|
||||
if joined_char_count(&candidate) > budget.max_chars {
|
||||
continue;
|
||||
}
|
||||
|
||||
selected.insert(index);
|
||||
}
|
||||
}
|
||||
|
||||
selected.into_iter().collect()
|
||||
}
|
||||
|
||||
fn push_line_with_budget(lines: &mut Vec<String>, line: String, budget: SummaryCompressionBudget) {
|
||||
let candidate = lines
|
||||
.iter()
|
||||
.map(String::as_str)
|
||||
.chain(std::iter::once(line.as_str()))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if candidate.len() <= budget.max_lines && joined_char_count(&candidate) <= budget.max_chars {
|
||||
lines.push(line);
|
||||
}
|
||||
}
|
||||
|
||||
fn joined_char_count(lines: &[&str]) -> usize {
|
||||
lines.iter().map(|line| line.chars().count()).sum::<usize>() + lines.len().saturating_sub(1)
|
||||
}
|
||||
|
||||
fn line_priority(line: &str) -> usize {
|
||||
if line == "Summary:" || line == "Conversation summary:" || is_core_detail(line) {
|
||||
0
|
||||
} else if is_section_header(line) {
|
||||
1
|
||||
} else if line.starts_with("- ") || line.starts_with(" - ") {
|
||||
2
|
||||
} else {
|
||||
3
|
||||
}
|
||||
}
|
||||
|
||||
fn is_core_detail(line: &str) -> bool {
|
||||
[
|
||||
"- Scope:",
|
||||
"- Current work:",
|
||||
"- Pending work:",
|
||||
"- Key files referenced:",
|
||||
"- Tools mentioned:",
|
||||
"- Recent user requests:",
|
||||
"- Previously compacted context:",
|
||||
"- Newly compacted context:",
|
||||
]
|
||||
.iter()
|
||||
.any(|prefix| line.starts_with(prefix))
|
||||
}
|
||||
|
||||
fn is_section_header(line: &str) -> bool {
|
||||
line.ends_with(':')
|
||||
}
|
||||
|
||||
fn omission_notice(omitted_lines: usize) -> String {
|
||||
format!("- … {omitted_lines} additional line(s) omitted.")
|
||||
}
|
||||
|
||||
fn collapse_inline_whitespace(line: &str) -> String {
|
||||
line.split_whitespace().collect::<Vec<_>>().join(" ")
|
||||
}
|
||||
|
||||
fn truncate_line(line: &str, max_chars: usize) -> String {
|
||||
if max_chars == 0 || line.chars().count() <= max_chars {
|
||||
return line.to_string();
|
||||
}
|
||||
|
||||
if max_chars == 1 {
|
||||
return "…".to_string();
|
||||
}
|
||||
|
||||
let mut truncated = line
|
||||
.chars()
|
||||
.take(max_chars.saturating_sub(1))
|
||||
.collect::<String>();
|
||||
truncated.push('…');
|
||||
truncated
|
||||
}
|
||||
|
||||
fn dedupe_key(line: &str) -> String {
|
||||
line.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{compress_summary, compress_summary_text, SummaryCompressionBudget};
|
||||
|
||||
#[test]
|
||||
fn collapses_whitespace_and_duplicate_lines() {
|
||||
// given
|
||||
let summary = "Conversation summary:\n\n- Scope: compact earlier messages.\n- Scope: compact earlier messages.\n- Current work: update runtime module.\n";
|
||||
|
||||
// when
|
||||
let result = compress_summary(summary, SummaryCompressionBudget::default());
|
||||
|
||||
// then
|
||||
assert_eq!(result.removed_duplicate_lines, 1);
|
||||
assert!(result
|
||||
.summary
|
||||
.contains("- Scope: compact earlier messages."));
|
||||
assert!(!result.summary.contains(" compact earlier"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_core_lines_when_budget_is_tight() {
|
||||
// given
|
||||
let summary = [
|
||||
"Conversation summary:",
|
||||
"- Scope: 18 earlier messages compacted.",
|
||||
"- Current work: finish summary compression.",
|
||||
"- Key timeline:",
|
||||
" - user: asked for a working implementation.",
|
||||
" - assistant: inspected runtime compaction flow.",
|
||||
" - tool: cargo check succeeded.",
|
||||
]
|
||||
.join("\n");
|
||||
|
||||
// when
|
||||
let result = compress_summary(
|
||||
&summary,
|
||||
SummaryCompressionBudget {
|
||||
max_chars: 120,
|
||||
max_lines: 3,
|
||||
max_line_chars: 80,
|
||||
},
|
||||
);
|
||||
|
||||
// then
|
||||
assert!(result.summary.contains("Conversation summary:"));
|
||||
assert!(result
|
||||
.summary
|
||||
.contains("- Scope: 18 earlier messages compacted."));
|
||||
assert!(result
|
||||
.summary
|
||||
.contains("- Current work: finish summary compression."));
|
||||
assert!(result.omitted_lines > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provides_a_default_text_only_helper() {
|
||||
// given
|
||||
let summary = "Summary:\n\nA short line.";
|
||||
|
||||
// when
|
||||
let compressed = compress_summary_text(summary);
|
||||
|
||||
// then
|
||||
assert_eq!(compressed, "Summary:\nA short line.");
|
||||
}
|
||||
}
|
||||
@@ -1,158 +0,0 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct TaskPacket {
|
||||
pub objective: String,
|
||||
pub scope: String,
|
||||
pub repo: String,
|
||||
pub branch_policy: String,
|
||||
pub acceptance_tests: Vec<String>,
|
||||
pub commit_policy: String,
|
||||
pub reporting_contract: String,
|
||||
pub escalation_policy: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct TaskPacketValidationError {
|
||||
errors: Vec<String>,
|
||||
}
|
||||
|
||||
impl TaskPacketValidationError {
|
||||
#[must_use]
|
||||
pub fn new(errors: Vec<String>) -> Self {
|
||||
Self { errors }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn errors(&self) -> &[String] {
|
||||
&self.errors
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for TaskPacketValidationError {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "{}", self.errors.join("; "))
|
||||
}
|
||||
}
|
||||
|
||||
impl std::error::Error for TaskPacketValidationError {}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct ValidatedPacket(TaskPacket);
|
||||
|
||||
impl ValidatedPacket {
|
||||
#[must_use]
|
||||
pub fn packet(&self) -> &TaskPacket {
|
||||
&self.0
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn into_inner(self) -> TaskPacket {
|
||||
self.0
|
||||
}
|
||||
}
|
||||
|
||||
pub fn validate_packet(packet: TaskPacket) -> Result<ValidatedPacket, TaskPacketValidationError> {
|
||||
let mut errors = Vec::new();
|
||||
|
||||
validate_required("objective", &packet.objective, &mut errors);
|
||||
validate_required("scope", &packet.scope, &mut errors);
|
||||
validate_required("repo", &packet.repo, &mut errors);
|
||||
validate_required("branch_policy", &packet.branch_policy, &mut errors);
|
||||
validate_required("commit_policy", &packet.commit_policy, &mut errors);
|
||||
validate_required(
|
||||
"reporting_contract",
|
||||
&packet.reporting_contract,
|
||||
&mut errors,
|
||||
);
|
||||
validate_required("escalation_policy", &packet.escalation_policy, &mut errors);
|
||||
|
||||
for (index, test) in packet.acceptance_tests.iter().enumerate() {
|
||||
if test.trim().is_empty() {
|
||||
errors.push(format!(
|
||||
"acceptance_tests contains an empty value at index {index}"
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if errors.is_empty() {
|
||||
Ok(ValidatedPacket(packet))
|
||||
} else {
|
||||
Err(TaskPacketValidationError::new(errors))
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_required(field: &str, value: &str, errors: &mut Vec<String>) {
|
||||
if value.trim().is_empty() {
|
||||
errors.push(format!("{field} must not be empty"));
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn sample_packet() -> TaskPacket {
|
||||
TaskPacket {
|
||||
objective: "Implement typed task packet format".to_string(),
|
||||
scope: "runtime/task system".to_string(),
|
||||
repo: "claw-code-parity".to_string(),
|
||||
branch_policy: "origin/main only".to_string(),
|
||||
acceptance_tests: vec![
|
||||
"cargo build --workspace".to_string(),
|
||||
"cargo test --workspace".to_string(),
|
||||
],
|
||||
commit_policy: "single verified commit".to_string(),
|
||||
reporting_contract: "print build result, test result, commit sha".to_string(),
|
||||
escalation_policy: "stop only on destructive ambiguity".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn valid_packet_passes_validation() {
|
||||
let packet = sample_packet();
|
||||
let validated = validate_packet(packet.clone()).expect("packet should validate");
|
||||
assert_eq!(validated.packet(), &packet);
|
||||
assert_eq!(validated.into_inner(), packet);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_packet_accumulates_errors() {
|
||||
let packet = TaskPacket {
|
||||
objective: " ".to_string(),
|
||||
scope: String::new(),
|
||||
repo: String::new(),
|
||||
branch_policy: "\t".to_string(),
|
||||
acceptance_tests: vec!["ok".to_string(), " ".to_string()],
|
||||
commit_policy: String::new(),
|
||||
reporting_contract: String::new(),
|
||||
escalation_policy: String::new(),
|
||||
};
|
||||
|
||||
let error = validate_packet(packet).expect_err("packet should be rejected");
|
||||
|
||||
assert!(error.errors().len() >= 7);
|
||||
assert!(error
|
||||
.errors()
|
||||
.contains(&"objective must not be empty".to_string()));
|
||||
assert!(error
|
||||
.errors()
|
||||
.contains(&"scope must not be empty".to_string()));
|
||||
assert!(error
|
||||
.errors()
|
||||
.contains(&"repo must not be empty".to_string()));
|
||||
assert!(error
|
||||
.errors()
|
||||
.contains(&"acceptance_tests contains an empty value at index 1".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn serialization_roundtrip_preserves_packet() {
|
||||
let packet = sample_packet();
|
||||
let serialized = serde_json::to_string(&packet).expect("packet should serialize");
|
||||
let deserialized: TaskPacket =
|
||||
serde_json::from_str(&serialized).expect("packet should deserialize");
|
||||
assert_eq!(deserialized, packet);
|
||||
}
|
||||
}
|
||||
@@ -1,503 +0,0 @@
|
||||
#![allow(clippy::must_use_candidate, clippy::unnecessary_map_or)]
|
||||
//! In-memory task registry for sub-agent task lifecycle management.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::{validate_packet, TaskPacket, TaskPacketValidationError};
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TaskStatus {
|
||||
Created,
|
||||
Running,
|
||||
Completed,
|
||||
Failed,
|
||||
Stopped,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TaskStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Created => write!(f, "created"),
|
||||
Self::Running => write!(f, "running"),
|
||||
Self::Completed => write!(f, "completed"),
|
||||
Self::Failed => write!(f, "failed"),
|
||||
Self::Stopped => write!(f, "stopped"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Task {
|
||||
pub task_id: String,
|
||||
pub prompt: String,
|
||||
pub description: Option<String>,
|
||||
pub task_packet: Option<TaskPacket>,
|
||||
pub status: TaskStatus,
|
||||
pub created_at: u64,
|
||||
pub updated_at: u64,
|
||||
pub messages: Vec<TaskMessage>,
|
||||
pub output: String,
|
||||
pub team_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TaskMessage {
|
||||
pub role: String,
|
||||
pub content: String,
|
||||
pub timestamp: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TaskRegistry {
|
||||
inner: Arc<Mutex<RegistryInner>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct RegistryInner {
|
||||
tasks: HashMap<String, Task>,
|
||||
counter: u64,
|
||||
}
|
||||
|
||||
fn now_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs()
|
||||
}
|
||||
|
||||
impl TaskRegistry {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn create(&self, prompt: &str, description: Option<&str>) -> Task {
|
||||
self.create_task(prompt.to_owned(), description.map(str::to_owned), None)
|
||||
}
|
||||
|
||||
pub fn create_from_packet(
|
||||
&self,
|
||||
packet: TaskPacket,
|
||||
) -> Result<Task, TaskPacketValidationError> {
|
||||
let packet = validate_packet(packet)?.into_inner();
|
||||
Ok(self.create_task(
|
||||
packet.objective.clone(),
|
||||
Some(packet.scope.clone()),
|
||||
Some(packet),
|
||||
))
|
||||
}
|
||||
|
||||
fn create_task(
|
||||
&self,
|
||||
prompt: String,
|
||||
description: Option<String>,
|
||||
task_packet: Option<TaskPacket>,
|
||||
) -> Task {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner.counter += 1;
|
||||
let ts = now_secs();
|
||||
let task_id = format!("task_{:08x}_{}", ts, inner.counter);
|
||||
let task = Task {
|
||||
task_id: task_id.clone(),
|
||||
prompt,
|
||||
description,
|
||||
task_packet,
|
||||
status: TaskStatus::Created,
|
||||
created_at: ts,
|
||||
updated_at: ts,
|
||||
messages: Vec::new(),
|
||||
output: String::new(),
|
||||
team_id: None,
|
||||
};
|
||||
inner.tasks.insert(task_id, task.clone());
|
||||
task
|
||||
}
|
||||
|
||||
pub fn get(&self, task_id: &str) -> Option<Task> {
|
||||
let inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner.tasks.get(task_id).cloned()
|
||||
}
|
||||
|
||||
pub fn list(&self, status_filter: Option<TaskStatus>) -> Vec<Task> {
|
||||
let inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner
|
||||
.tasks
|
||||
.values()
|
||||
.filter(|t| status_filter.map_or(true, |s| t.status == s))
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn stop(&self, task_id: &str) -> Result<Task, String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
|
||||
match task.status {
|
||||
TaskStatus::Completed | TaskStatus::Failed | TaskStatus::Stopped => {
|
||||
return Err(format!(
|
||||
"task {task_id} is already in terminal state: {}",
|
||||
task.status
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
task.status = TaskStatus::Stopped;
|
||||
task.updated_at = now_secs();
|
||||
Ok(task.clone())
|
||||
}
|
||||
|
||||
pub fn update(&self, task_id: &str, message: &str) -> Result<Task, String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
|
||||
task.messages.push(TaskMessage {
|
||||
role: String::from("user"),
|
||||
content: message.to_owned(),
|
||||
timestamp: now_secs(),
|
||||
});
|
||||
task.updated_at = now_secs();
|
||||
Ok(task.clone())
|
||||
}
|
||||
|
||||
pub fn output(&self, task_id: &str) -> Result<String, String> {
|
||||
let inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
Ok(task.output.clone())
|
||||
}
|
||||
|
||||
pub fn append_output(&self, task_id: &str, output: &str) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
task.output.push_str(output);
|
||||
task.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn set_status(&self, task_id: &str, status: TaskStatus) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
task.status = status;
|
||||
task.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn assign_team(&self, task_id: &str, team_id: &str) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
let task = inner
|
||||
.tasks
|
||||
.get_mut(task_id)
|
||||
.ok_or_else(|| format!("task not found: {task_id}"))?;
|
||||
task.team_id = Some(team_id.to_owned());
|
||||
task.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn remove(&self, task_id: &str) -> Option<Task> {
|
||||
let mut inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner.tasks.remove(task_id)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn len(&self) -> usize {
|
||||
let inner = self.inner.lock().expect("registry lock poisoned");
|
||||
inner.tasks.len()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn creates_and_retrieves_tasks() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Do something", Some("A test task"));
|
||||
assert_eq!(task.status, TaskStatus::Created);
|
||||
assert_eq!(task.prompt, "Do something");
|
||||
assert_eq!(task.description.as_deref(), Some("A test task"));
|
||||
assert_eq!(task.task_packet, None);
|
||||
|
||||
let fetched = registry.get(&task.task_id).expect("task should exist");
|
||||
assert_eq!(fetched.task_id, task.task_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn creates_task_from_packet() {
|
||||
let registry = TaskRegistry::new();
|
||||
let packet = TaskPacket {
|
||||
objective: "Ship task packet support".to_string(),
|
||||
scope: "runtime/task system".to_string(),
|
||||
repo: "claw-code-parity".to_string(),
|
||||
branch_policy: "origin/main only".to_string(),
|
||||
acceptance_tests: vec!["cargo test --workspace".to_string()],
|
||||
commit_policy: "single commit".to_string(),
|
||||
reporting_contract: "print commit sha".to_string(),
|
||||
escalation_policy: "manual escalation".to_string(),
|
||||
};
|
||||
|
||||
let task = registry
|
||||
.create_from_packet(packet.clone())
|
||||
.expect("packet-backed task should be created");
|
||||
|
||||
assert_eq!(task.prompt, packet.objective);
|
||||
assert_eq!(task.description.as_deref(), Some("runtime/task system"));
|
||||
assert_eq!(task.task_packet, Some(packet.clone()));
|
||||
|
||||
let fetched = registry.get(&task.task_id).expect("task should exist");
|
||||
assert_eq!(fetched.task_packet, Some(packet));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lists_tasks_with_optional_filter() {
|
||||
let registry = TaskRegistry::new();
|
||||
registry.create("Task A", None);
|
||||
let task_b = registry.create("Task B", None);
|
||||
registry
|
||||
.set_status(&task_b.task_id, TaskStatus::Running)
|
||||
.expect("set status should succeed");
|
||||
|
||||
let all = registry.list(None);
|
||||
assert_eq!(all.len(), 2);
|
||||
|
||||
let running = registry.list(Some(TaskStatus::Running));
|
||||
assert_eq!(running.len(), 1);
|
||||
assert_eq!(running[0].task_id, task_b.task_id);
|
||||
|
||||
let created = registry.list(Some(TaskStatus::Created));
|
||||
assert_eq!(created.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stops_running_task() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Stoppable", None);
|
||||
registry
|
||||
.set_status(&task.task_id, TaskStatus::Running)
|
||||
.unwrap();
|
||||
|
||||
let stopped = registry.stop(&task.task_id).expect("stop should succeed");
|
||||
assert_eq!(stopped.status, TaskStatus::Stopped);
|
||||
|
||||
// Stopping again should fail
|
||||
let result = registry.stop(&task.task_id);
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn updates_task_with_messages() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Messageable", None);
|
||||
let updated = registry
|
||||
.update(&task.task_id, "Here's more context")
|
||||
.expect("update should succeed");
|
||||
assert_eq!(updated.messages.len(), 1);
|
||||
assert_eq!(updated.messages[0].content, "Here's more context");
|
||||
assert_eq!(updated.messages[0].role, "user");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn appends_and_retrieves_output() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Output task", None);
|
||||
registry
|
||||
.append_output(&task.task_id, "line 1\n")
|
||||
.expect("append should succeed");
|
||||
registry
|
||||
.append_output(&task.task_id, "line 2\n")
|
||||
.expect("append should succeed");
|
||||
|
||||
let output = registry.output(&task.task_id).expect("output should exist");
|
||||
assert_eq!(output, "line 1\nline 2\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assigns_team_and_removes_task() {
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("Team task", None);
|
||||
registry
|
||||
.assign_team(&task.task_id, "team_abc")
|
||||
.expect("assign should succeed");
|
||||
|
||||
let fetched = registry.get(&task.task_id).unwrap();
|
||||
assert_eq!(fetched.team_id.as_deref(), Some("team_abc"));
|
||||
|
||||
let removed = registry.remove(&task.task_id);
|
||||
assert!(removed.is_some());
|
||||
assert!(registry.get(&task.task_id).is_none());
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_operations_on_missing_task() {
|
||||
let registry = TaskRegistry::new();
|
||||
assert!(registry.stop("nonexistent").is_err());
|
||||
assert!(registry.update("nonexistent", "msg").is_err());
|
||||
assert!(registry.output("nonexistent").is_err());
|
||||
assert!(registry.append_output("nonexistent", "data").is_err());
|
||||
assert!(registry
|
||||
.set_status("nonexistent", TaskStatus::Running)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn task_status_display_all_variants() {
|
||||
// given
|
||||
let cases = [
|
||||
(TaskStatus::Created, "created"),
|
||||
(TaskStatus::Running, "running"),
|
||||
(TaskStatus::Completed, "completed"),
|
||||
(TaskStatus::Failed, "failed"),
|
||||
(TaskStatus::Stopped, "stopped"),
|
||||
];
|
||||
|
||||
// when
|
||||
let rendered: Vec<_> = cases
|
||||
.into_iter()
|
||||
.map(|(status, expected)| (status.to_string(), expected))
|
||||
.collect();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
rendered,
|
||||
vec![
|
||||
("created".to_string(), "created"),
|
||||
("running".to_string(), "running"),
|
||||
("completed".to_string(), "completed"),
|
||||
("failed".to_string(), "failed"),
|
||||
("stopped".to_string(), "stopped"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_rejects_completed_task() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("done", None);
|
||||
registry
|
||||
.set_status(&task.task_id, TaskStatus::Completed)
|
||||
.expect("set status should succeed");
|
||||
|
||||
// when
|
||||
let result = registry.stop(&task.task_id);
|
||||
|
||||
// then
|
||||
let error = result.expect_err("completed task should be rejected");
|
||||
assert!(error.contains("already in terminal state"));
|
||||
assert!(error.contains("completed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_rejects_failed_task() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("failed", None);
|
||||
registry
|
||||
.set_status(&task.task_id, TaskStatus::Failed)
|
||||
.expect("set status should succeed");
|
||||
|
||||
// when
|
||||
let result = registry.stop(&task.task_id);
|
||||
|
||||
// then
|
||||
let error = result.expect_err("failed task should be rejected");
|
||||
assert!(error.contains("already in terminal state"));
|
||||
assert!(error.contains("failed"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stop_succeeds_from_created_state() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
let task = registry.create("created task", None);
|
||||
|
||||
// when
|
||||
let stopped = registry.stop(&task.task_id).expect("stop should succeed");
|
||||
|
||||
// then
|
||||
assert_eq!(stopped.status, TaskStatus::Stopped);
|
||||
assert!(stopped.updated_at >= task.updated_at);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_registry_is_empty() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
|
||||
// when
|
||||
let all_tasks = registry.list(None);
|
||||
|
||||
// then
|
||||
assert!(registry.is_empty());
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(all_tasks.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn create_without_description() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
|
||||
// when
|
||||
let task = registry.create("Do the thing", None);
|
||||
|
||||
// then
|
||||
assert!(task.task_id.starts_with("task_"));
|
||||
assert_eq!(task.description, None);
|
||||
assert_eq!(task.task_packet, None);
|
||||
assert!(task.messages.is_empty());
|
||||
assert!(task.output.is_empty());
|
||||
assert_eq!(task.team_id, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_nonexistent_returns_none() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
|
||||
// when
|
||||
let removed = registry.remove("missing");
|
||||
|
||||
// then
|
||||
assert!(removed.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assign_team_rejects_missing_task() {
|
||||
// given
|
||||
let registry = TaskRegistry::new();
|
||||
|
||||
// when
|
||||
let result = registry.assign_team("missing", "team_123");
|
||||
|
||||
// then
|
||||
let error = result.expect_err("missing task should be rejected");
|
||||
assert_eq!(error, "task not found: missing");
|
||||
}
|
||||
}
|
||||
@@ -1,509 +0,0 @@
|
||||
#![allow(clippy::must_use_candidate)]
|
||||
//! In-memory registries for Team and Cron lifecycle management.
|
||||
//!
|
||||
//! Provides TeamCreate/Delete and CronCreate/Delete/List runtime backing
|
||||
//! to replace the stub implementations in the tools crate.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{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, Serialize, Deserialize)]
|
||||
pub struct Team {
|
||||
pub team_id: String,
|
||||
pub name: String,
|
||||
pub task_ids: Vec<String>,
|
||||
pub status: TeamStatus,
|
||||
pub created_at: u64,
|
||||
pub updated_at: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TeamStatus {
|
||||
Created,
|
||||
Running,
|
||||
Completed,
|
||||
Deleted,
|
||||
}
|
||||
|
||||
impl std::fmt::Display for TeamStatus {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
Self::Created => write!(f, "created"),
|
||||
Self::Running => write!(f, "running"),
|
||||
Self::Completed => write!(f, "completed"),
|
||||
Self::Deleted => write!(f, "deleted"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TeamRegistry {
|
||||
inner: Arc<Mutex<TeamInner>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct TeamInner {
|
||||
teams: HashMap<String, Team>,
|
||||
counter: u64,
|
||||
}
|
||||
|
||||
impl TeamRegistry {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn create(&self, name: &str, task_ids: Vec<String>) -> Team {
|
||||
let mut inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.counter += 1;
|
||||
let ts = now_secs();
|
||||
let team_id = format!("team_{:08x}_{}", ts, inner.counter);
|
||||
let team = Team {
|
||||
team_id: team_id.clone(),
|
||||
name: name.to_owned(),
|
||||
task_ids,
|
||||
status: TeamStatus::Created,
|
||||
created_at: ts,
|
||||
updated_at: ts,
|
||||
};
|
||||
inner.teams.insert(team_id, team.clone());
|
||||
team
|
||||
}
|
||||
|
||||
pub fn get(&self, team_id: &str) -> Option<Team> {
|
||||
let inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.teams.get(team_id).cloned()
|
||||
}
|
||||
|
||||
pub fn list(&self) -> Vec<Team> {
|
||||
let inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.teams.values().cloned().collect()
|
||||
}
|
||||
|
||||
pub fn delete(&self, team_id: &str) -> Result<Team, String> {
|
||||
let mut inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
let team = inner
|
||||
.teams
|
||||
.get_mut(team_id)
|
||||
.ok_or_else(|| format!("team not found: {team_id}"))?;
|
||||
team.status = TeamStatus::Deleted;
|
||||
team.updated_at = now_secs();
|
||||
Ok(team.clone())
|
||||
}
|
||||
|
||||
pub fn remove(&self, team_id: &str) -> Option<Team> {
|
||||
let mut inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.teams.remove(team_id)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn len(&self) -> usize {
|
||||
let inner = self.inner.lock().expect("team registry lock poisoned");
|
||||
inner.teams.len()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CronEntry {
|
||||
pub cron_id: String,
|
||||
pub schedule: String,
|
||||
pub prompt: String,
|
||||
pub description: Option<String>,
|
||||
pub enabled: bool,
|
||||
pub created_at: u64,
|
||||
pub updated_at: u64,
|
||||
pub last_run_at: Option<u64>,
|
||||
pub run_count: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct CronRegistry {
|
||||
inner: Arc<Mutex<CronInner>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct CronInner {
|
||||
entries: HashMap<String, CronEntry>,
|
||||
counter: u64,
|
||||
}
|
||||
|
||||
impl CronRegistry {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn create(&self, schedule: &str, prompt: &str, description: Option<&str>) -> CronEntry {
|
||||
let mut inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner.counter += 1;
|
||||
let ts = now_secs();
|
||||
let cron_id = format!("cron_{:08x}_{}", ts, inner.counter);
|
||||
let entry = CronEntry {
|
||||
cron_id: cron_id.clone(),
|
||||
schedule: schedule.to_owned(),
|
||||
prompt: prompt.to_owned(),
|
||||
description: description.map(str::to_owned),
|
||||
enabled: true,
|
||||
created_at: ts,
|
||||
updated_at: ts,
|
||||
last_run_at: None,
|
||||
run_count: 0,
|
||||
};
|
||||
inner.entries.insert(cron_id, entry.clone());
|
||||
entry
|
||||
}
|
||||
|
||||
pub fn get(&self, cron_id: &str) -> Option<CronEntry> {
|
||||
let inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner.entries.get(cron_id).cloned()
|
||||
}
|
||||
|
||||
pub fn list(&self, enabled_only: bool) -> Vec<CronEntry> {
|
||||
let inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner
|
||||
.entries
|
||||
.values()
|
||||
.filter(|e| !enabled_only || e.enabled)
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn delete(&self, cron_id: &str) -> Result<CronEntry, String> {
|
||||
let mut inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner
|
||||
.entries
|
||||
.remove(cron_id)
|
||||
.ok_or_else(|| format!("cron not found: {cron_id}"))
|
||||
}
|
||||
|
||||
/// Disable a cron entry without removing it.
|
||||
pub fn disable(&self, cron_id: &str) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
let entry = inner
|
||||
.entries
|
||||
.get_mut(cron_id)
|
||||
.ok_or_else(|| format!("cron not found: {cron_id}"))?;
|
||||
entry.enabled = false;
|
||||
entry.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Record a cron run.
|
||||
pub fn record_run(&self, cron_id: &str) -> Result<(), String> {
|
||||
let mut inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
let entry = inner
|
||||
.entries
|
||||
.get_mut(cron_id)
|
||||
.ok_or_else(|| format!("cron not found: {cron_id}"))?;
|
||||
entry.last_run_at = Some(now_secs());
|
||||
entry.run_count += 1;
|
||||
entry.updated_at = now_secs();
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn len(&self) -> usize {
|
||||
let inner = self.inner.lock().expect("cron registry lock poisoned");
|
||||
inner.entries.len()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.len() == 0
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// ── Team tests ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn creates_and_retrieves_team() {
|
||||
let registry = TeamRegistry::new();
|
||||
let team = registry.create("Alpha Squad", vec!["task_001".into(), "task_002".into()]);
|
||||
assert_eq!(team.name, "Alpha Squad");
|
||||
assert_eq!(team.task_ids.len(), 2);
|
||||
assert_eq!(team.status, TeamStatus::Created);
|
||||
|
||||
let fetched = registry.get(&team.team_id).expect("team should exist");
|
||||
assert_eq!(fetched.team_id, team.team_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lists_and_deletes_teams() {
|
||||
let registry = TeamRegistry::new();
|
||||
let t1 = registry.create("Team A", vec![]);
|
||||
let t2 = registry.create("Team B", vec![]);
|
||||
|
||||
let all = registry.list();
|
||||
assert_eq!(all.len(), 2);
|
||||
|
||||
let deleted = registry.delete(&t1.team_id).expect("delete should succeed");
|
||||
assert_eq!(deleted.status, TeamStatus::Deleted);
|
||||
|
||||
// Team is still listable (soft delete)
|
||||
let still_there = registry.get(&t1.team_id).unwrap();
|
||||
assert_eq!(still_there.status, TeamStatus::Deleted);
|
||||
|
||||
// Hard remove
|
||||
registry.remove(&t2.team_id);
|
||||
assert_eq!(registry.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_missing_team_operations() {
|
||||
let registry = TeamRegistry::new();
|
||||
assert!(registry.delete("nonexistent").is_err());
|
||||
assert!(registry.get("nonexistent").is_none());
|
||||
}
|
||||
|
||||
// ── Cron tests ──────────────────────────────────────
|
||||
|
||||
#[test]
|
||||
fn creates_and_retrieves_cron() {
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("0 * * * *", "Check status", Some("hourly check"));
|
||||
assert_eq!(entry.schedule, "0 * * * *");
|
||||
assert_eq!(entry.prompt, "Check status");
|
||||
assert!(entry.enabled);
|
||||
assert_eq!(entry.run_count, 0);
|
||||
assert!(entry.last_run_at.is_none());
|
||||
|
||||
let fetched = registry.get(&entry.cron_id).expect("cron should exist");
|
||||
assert_eq!(fetched.cron_id, entry.cron_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lists_with_enabled_filter() {
|
||||
let registry = CronRegistry::new();
|
||||
let c1 = registry.create("* * * * *", "Task 1", None);
|
||||
let c2 = registry.create("0 * * * *", "Task 2", None);
|
||||
registry
|
||||
.disable(&c1.cron_id)
|
||||
.expect("disable should succeed");
|
||||
|
||||
let all = registry.list(false);
|
||||
assert_eq!(all.len(), 2);
|
||||
|
||||
let enabled_only = registry.list(true);
|
||||
assert_eq!(enabled_only.len(), 1);
|
||||
assert_eq!(enabled_only[0].cron_id, c2.cron_id);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deletes_cron_entry() {
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("* * * * *", "To delete", None);
|
||||
let deleted = registry
|
||||
.delete(&entry.cron_id)
|
||||
.expect("delete should succeed");
|
||||
assert_eq!(deleted.cron_id, entry.cron_id);
|
||||
assert!(registry.get(&entry.cron_id).is_none());
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn records_cron_runs() {
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("*/5 * * * *", "Recurring", None);
|
||||
registry.record_run(&entry.cron_id).unwrap();
|
||||
registry.record_run(&entry.cron_id).unwrap();
|
||||
|
||||
let fetched = registry.get(&entry.cron_id).unwrap();
|
||||
assert_eq!(fetched.run_count, 2);
|
||||
assert!(fetched.last_run_at.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_missing_cron_operations() {
|
||||
let registry = CronRegistry::new();
|
||||
assert!(registry.delete("nonexistent").is_err());
|
||||
assert!(registry.disable("nonexistent").is_err());
|
||||
assert!(registry.record_run("nonexistent").is_err());
|
||||
assert!(registry.get("nonexistent").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn team_status_display_all_variants() {
|
||||
// given
|
||||
let cases = [
|
||||
(TeamStatus::Created, "created"),
|
||||
(TeamStatus::Running, "running"),
|
||||
(TeamStatus::Completed, "completed"),
|
||||
(TeamStatus::Deleted, "deleted"),
|
||||
];
|
||||
|
||||
// when
|
||||
let rendered: Vec<_> = cases
|
||||
.into_iter()
|
||||
.map(|(status, expected)| (status.to_string(), expected))
|
||||
.collect();
|
||||
|
||||
// then
|
||||
assert_eq!(
|
||||
rendered,
|
||||
vec![
|
||||
("created".to_string(), "created"),
|
||||
("running".to_string(), "running"),
|
||||
("completed".to_string(), "completed"),
|
||||
("deleted".to_string(), "deleted"),
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_team_registry_is_empty() {
|
||||
// given
|
||||
let registry = TeamRegistry::new();
|
||||
|
||||
// when
|
||||
let teams = registry.list();
|
||||
|
||||
// then
|
||||
assert!(registry.is_empty());
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(teams.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn team_remove_nonexistent_returns_none() {
|
||||
// given
|
||||
let registry = TeamRegistry::new();
|
||||
|
||||
// when
|
||||
let removed = registry.remove("missing");
|
||||
|
||||
// then
|
||||
assert!(removed.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn team_len_transitions() {
|
||||
// given
|
||||
let registry = TeamRegistry::new();
|
||||
|
||||
// when
|
||||
let alpha = registry.create("Alpha", vec![]);
|
||||
let beta = registry.create("Beta", vec![]);
|
||||
let after_create = registry.len();
|
||||
registry.remove(&alpha.team_id);
|
||||
let after_first_remove = registry.len();
|
||||
registry.remove(&beta.team_id);
|
||||
|
||||
// then
|
||||
assert_eq!(after_create, 2);
|
||||
assert_eq!(after_first_remove, 1);
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(registry.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_list_all_disabled_returns_empty_for_enabled_only() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
let first = registry.create("* * * * *", "Task 1", None);
|
||||
let second = registry.create("0 * * * *", "Task 2", None);
|
||||
registry
|
||||
.disable(&first.cron_id)
|
||||
.expect("disable should succeed");
|
||||
registry
|
||||
.disable(&second.cron_id)
|
||||
.expect("disable should succeed");
|
||||
|
||||
// when
|
||||
let enabled_only = registry.list(true);
|
||||
let all_entries = registry.list(false);
|
||||
|
||||
// then
|
||||
assert!(enabled_only.is_empty());
|
||||
assert_eq!(all_entries.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_create_without_description() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
|
||||
// when
|
||||
let entry = registry.create("*/15 * * * *", "Check health", None);
|
||||
|
||||
// then
|
||||
assert!(entry.cron_id.starts_with("cron_"));
|
||||
assert_eq!(entry.description, None);
|
||||
assert!(entry.enabled);
|
||||
assert_eq!(entry.run_count, 0);
|
||||
assert_eq!(entry.last_run_at, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_cron_registry_is_empty() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
|
||||
// when
|
||||
let enabled_only = registry.list(true);
|
||||
let all_entries = registry.list(false);
|
||||
|
||||
// then
|
||||
assert!(registry.is_empty());
|
||||
assert_eq!(registry.len(), 0);
|
||||
assert!(enabled_only.is_empty());
|
||||
assert!(all_entries.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_record_run_updates_timestamp_and_counter() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("*/5 * * * *", "Recurring", None);
|
||||
|
||||
// when
|
||||
registry
|
||||
.record_run(&entry.cron_id)
|
||||
.expect("first run should succeed");
|
||||
registry
|
||||
.record_run(&entry.cron_id)
|
||||
.expect("second run should succeed");
|
||||
let fetched = registry.get(&entry.cron_id).expect("entry should exist");
|
||||
|
||||
// then
|
||||
assert_eq!(fetched.run_count, 2);
|
||||
assert!(fetched.last_run_at.is_some());
|
||||
assert!(fetched.updated_at >= entry.updated_at);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cron_disable_updates_timestamp() {
|
||||
// given
|
||||
let registry = CronRegistry::new();
|
||||
let entry = registry.create("0 0 * * *", "Nightly", None);
|
||||
|
||||
// when
|
||||
registry
|
||||
.disable(&entry.cron_id)
|
||||
.expect("disable should succeed");
|
||||
let fetched = registry.get(&entry.cron_id).expect("entry should exist");
|
||||
|
||||
// then
|
||||
assert!(!fetched.enabled);
|
||||
assert!(fetched.updated_at >= entry.updated_at);
|
||||
}
|
||||
}
|
||||
@@ -1,299 +0,0 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
const TRUST_PROMPT_CUES: &[&str] = &[
|
||||
"do you trust the files in this folder",
|
||||
"trust the files in this folder",
|
||||
"trust this folder",
|
||||
"allow and continue",
|
||||
"yes, proceed",
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum TrustPolicy {
|
||||
AutoTrust,
|
||||
RequireApproval,
|
||||
Deny,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum TrustEvent {
|
||||
TrustRequired { cwd: String },
|
||||
TrustResolved { cwd: String, policy: TrustPolicy },
|
||||
TrustDenied { cwd: String, reason: String },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct TrustConfig {
|
||||
allowlisted: Vec<PathBuf>,
|
||||
denied: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
impl TrustConfig {
|
||||
#[must_use]
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_allowlisted(mut self, path: impl Into<PathBuf>) -> Self {
|
||||
self.allowlisted.push(path.into());
|
||||
self
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn with_denied(mut self, path: impl Into<PathBuf>) -> Self {
|
||||
self.denied.push(path.into());
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum TrustDecision {
|
||||
NotRequired,
|
||||
Required {
|
||||
policy: TrustPolicy,
|
||||
events: Vec<TrustEvent>,
|
||||
},
|
||||
}
|
||||
|
||||
impl TrustDecision {
|
||||
#[must_use]
|
||||
pub fn policy(&self) -> Option<TrustPolicy> {
|
||||
match self {
|
||||
Self::NotRequired => None,
|
||||
Self::Required { policy, .. } => Some(*policy),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn events(&self) -> &[TrustEvent] {
|
||||
match self {
|
||||
Self::NotRequired => &[],
|
||||
Self::Required { events, .. } => events,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TrustResolver {
|
||||
config: TrustConfig,
|
||||
}
|
||||
|
||||
impl TrustResolver {
|
||||
#[must_use]
|
||||
pub fn new(config: TrustConfig) -> Self {
|
||||
Self { config }
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn resolve(&self, cwd: &str, screen_text: &str) -> TrustDecision {
|
||||
if !detect_trust_prompt(screen_text) {
|
||||
return TrustDecision::NotRequired;
|
||||
}
|
||||
|
||||
let mut events = vec![TrustEvent::TrustRequired {
|
||||
cwd: cwd.to_owned(),
|
||||
}];
|
||||
|
||||
if let Some(matched_root) = self
|
||||
.config
|
||||
.denied
|
||||
.iter()
|
||||
.find(|root| path_matches(cwd, root))
|
||||
{
|
||||
let reason = format!("cwd matches denied trust root: {}", matched_root.display());
|
||||
events.push(TrustEvent::TrustDenied {
|
||||
cwd: cwd.to_owned(),
|
||||
reason,
|
||||
});
|
||||
return TrustDecision::Required {
|
||||
policy: TrustPolicy::Deny,
|
||||
events,
|
||||
};
|
||||
}
|
||||
|
||||
if self
|
||||
.config
|
||||
.allowlisted
|
||||
.iter()
|
||||
.any(|root| path_matches(cwd, root))
|
||||
{
|
||||
events.push(TrustEvent::TrustResolved {
|
||||
cwd: cwd.to_owned(),
|
||||
policy: TrustPolicy::AutoTrust,
|
||||
});
|
||||
return TrustDecision::Required {
|
||||
policy: TrustPolicy::AutoTrust,
|
||||
events,
|
||||
};
|
||||
}
|
||||
|
||||
TrustDecision::Required {
|
||||
policy: TrustPolicy::RequireApproval,
|
||||
events,
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn trusts(&self, cwd: &str) -> bool {
|
||||
!self
|
||||
.config
|
||||
.denied
|
||||
.iter()
|
||||
.any(|root| path_matches(cwd, root))
|
||||
&& self
|
||||
.config
|
||||
.allowlisted
|
||||
.iter()
|
||||
.any(|root| path_matches(cwd, root))
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn detect_trust_prompt(screen_text: &str) -> bool {
|
||||
let lowered = screen_text.to_ascii_lowercase();
|
||||
TRUST_PROMPT_CUES
|
||||
.iter()
|
||||
.any(|needle| lowered.contains(needle))
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn path_matches_trusted_root(cwd: &str, trusted_root: &str) -> bool {
|
||||
path_matches(cwd, &normalize_path(Path::new(trusted_root)))
|
||||
}
|
||||
|
||||
fn path_matches(candidate: &str, root: &Path) -> bool {
|
||||
let candidate = normalize_path(Path::new(candidate));
|
||||
let root = normalize_path(root);
|
||||
candidate == root || candidate.starts_with(&root)
|
||||
}
|
||||
|
||||
fn normalize_path(path: &Path) -> PathBuf {
|
||||
std::fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
detect_trust_prompt, path_matches_trusted_root, TrustConfig, TrustDecision, TrustEvent,
|
||||
TrustPolicy, TrustResolver,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn detects_known_trust_prompt_copy() {
|
||||
// given
|
||||
let screen_text = "Do you trust the files in this folder?\n1. Yes, proceed\n2. No";
|
||||
|
||||
// when
|
||||
let detected = detect_trust_prompt(screen_text);
|
||||
|
||||
// then
|
||||
assert!(detected);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn does_not_emit_events_when_prompt_is_absent() {
|
||||
// given
|
||||
let resolver = TrustResolver::new(TrustConfig::new().with_allowlisted("/tmp/worktrees"));
|
||||
|
||||
// when
|
||||
let decision = resolver.resolve("/tmp/worktrees/repo-a", "Ready for your input\n>");
|
||||
|
||||
// then
|
||||
assert_eq!(decision, TrustDecision::NotRequired);
|
||||
assert_eq!(decision.events(), &[]);
|
||||
assert_eq!(decision.policy(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_trusts_allowlisted_cwd_after_prompt_detection() {
|
||||
// given
|
||||
let resolver = TrustResolver::new(TrustConfig::new().with_allowlisted("/tmp/worktrees"));
|
||||
|
||||
// when
|
||||
let decision = resolver.resolve(
|
||||
"/tmp/worktrees/repo-a",
|
||||
"Do you trust the files in this folder?\n1. Yes, proceed\n2. No",
|
||||
);
|
||||
|
||||
// then
|
||||
assert_eq!(decision.policy(), Some(TrustPolicy::AutoTrust));
|
||||
assert_eq!(
|
||||
decision.events(),
|
||||
&[
|
||||
TrustEvent::TrustRequired {
|
||||
cwd: "/tmp/worktrees/repo-a".to_string(),
|
||||
},
|
||||
TrustEvent::TrustResolved {
|
||||
cwd: "/tmp/worktrees/repo-a".to_string(),
|
||||
policy: TrustPolicy::AutoTrust,
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn requires_approval_for_unknown_cwd_after_prompt_detection() {
|
||||
// given
|
||||
let resolver = TrustResolver::new(TrustConfig::new().with_allowlisted("/tmp/worktrees"));
|
||||
|
||||
// when
|
||||
let decision = resolver.resolve(
|
||||
"/tmp/other/repo-b",
|
||||
"Do you trust the files in this folder?\n1. Yes, proceed\n2. No",
|
||||
);
|
||||
|
||||
// then
|
||||
assert_eq!(decision.policy(), Some(TrustPolicy::RequireApproval));
|
||||
assert_eq!(
|
||||
decision.events(),
|
||||
&[TrustEvent::TrustRequired {
|
||||
cwd: "/tmp/other/repo-b".to_string(),
|
||||
}]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn denied_root_takes_precedence_over_allowlist() {
|
||||
// given
|
||||
let resolver = TrustResolver::new(
|
||||
TrustConfig::new()
|
||||
.with_allowlisted("/tmp/worktrees")
|
||||
.with_denied("/tmp/worktrees/repo-c"),
|
||||
);
|
||||
|
||||
// when
|
||||
let decision = resolver.resolve(
|
||||
"/tmp/worktrees/repo-c",
|
||||
"Do you trust the files in this folder?\n1. Yes, proceed\n2. No",
|
||||
);
|
||||
|
||||
// then
|
||||
assert_eq!(decision.policy(), Some(TrustPolicy::Deny));
|
||||
assert_eq!(
|
||||
decision.events(),
|
||||
&[
|
||||
TrustEvent::TrustRequired {
|
||||
cwd: "/tmp/worktrees/repo-c".to_string(),
|
||||
},
|
||||
TrustEvent::TrustDenied {
|
||||
cwd: "/tmp/worktrees/repo-c".to_string(),
|
||||
reason: "cwd matches denied trust root: /tmp/worktrees/repo-c".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sibling_prefix_does_not_match_trusted_root() {
|
||||
// given
|
||||
let trusted_root = "/tmp/worktrees";
|
||||
let sibling_path = "/tmp/worktrees-other/repo-d";
|
||||
|
||||
// when
|
||||
let matched = path_matches_trusted_root(sibling_path, trusted_root);
|
||||
|
||||
// then
|
||||
assert!(!matched);
|
||||
}
|
||||
}
|
||||
@@ -1,386 +0,0 @@
|
||||
#![allow(clippy::doc_markdown, clippy::uninlined_format_args, unused_imports)]
|
||||
//! Integration tests for cross-module wiring.
|
||||
//!
|
||||
//! These tests verify that adjacent modules in the runtime crate actually
|
||||
//! connect correctly — catching wiring gaps that unit tests miss.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use runtime::green_contract::{GreenContract, GreenContractOutcome, GreenLevel};
|
||||
use runtime::{
|
||||
apply_policy, BranchFreshness, DiffScope, LaneBlocker, LaneContext, PolicyAction,
|
||||
PolicyCondition, PolicyEngine, PolicyRule, ReconcileReason, ReviewStatus, StaleBranchAction,
|
||||
StaleBranchPolicy,
|
||||
};
|
||||
|
||||
/// stale_branch + policy_engine integration:
|
||||
/// When a branch is detected stale, does it correctly flow through
|
||||
/// PolicyCondition::StaleBranch to generate the expected action?
|
||||
#[test]
|
||||
fn stale_branch_detection_flows_into_policy_engine() {
|
||||
// given — a stale branch context (2 hours behind main, threshold is 1 hour)
|
||||
let stale_context = LaneContext::new(
|
||||
"stale-lane",
|
||||
0,
|
||||
Duration::from_secs(2 * 60 * 60), // 2 hours stale
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
);
|
||||
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"stale-merge-forward",
|
||||
PolicyCondition::StaleBranch,
|
||||
PolicyAction::MergeForward,
|
||||
10,
|
||||
)]);
|
||||
|
||||
// when
|
||||
let actions = engine.evaluate(&stale_context);
|
||||
|
||||
// then
|
||||
assert_eq!(actions, vec![PolicyAction::MergeForward]);
|
||||
}
|
||||
|
||||
/// stale_branch + policy_engine: Fresh branch does NOT trigger stale rules
|
||||
#[test]
|
||||
fn fresh_branch_does_not_trigger_stale_policy() {
|
||||
let fresh_context = LaneContext::new(
|
||||
"fresh-lane",
|
||||
0,
|
||||
Duration::from_secs(30 * 60), // 30 min stale — under 1 hour threshold
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
);
|
||||
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"stale-merge-forward",
|
||||
PolicyCondition::StaleBranch,
|
||||
PolicyAction::MergeForward,
|
||||
10,
|
||||
)]);
|
||||
|
||||
let actions = engine.evaluate(&fresh_context);
|
||||
assert!(actions.is_empty());
|
||||
}
|
||||
|
||||
/// green_contract + policy_engine integration:
|
||||
/// A lane that meets its green contract should be mergeable
|
||||
#[test]
|
||||
fn green_contract_satisfied_allows_merge() {
|
||||
let contract = GreenContract::new(GreenLevel::Workspace);
|
||||
let satisfied = contract.is_satisfied_by(GreenLevel::Workspace);
|
||||
assert!(satisfied);
|
||||
|
||||
let exceeded = contract.is_satisfied_by(GreenLevel::MergeReady);
|
||||
assert!(exceeded);
|
||||
|
||||
let insufficient = contract.is_satisfied_by(GreenLevel::Package);
|
||||
assert!(!insufficient);
|
||||
}
|
||||
|
||||
/// green_contract + policy_engine:
|
||||
/// Lane with green level below contract requirement gets blocked
|
||||
#[test]
|
||||
fn green_contract_unsatisfied_blocks_merge() {
|
||||
let context = LaneContext::new(
|
||||
"partial-green-lane",
|
||||
1, // GreenLevel::Package as u8
|
||||
Duration::from_secs(0),
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Pending,
|
||||
DiffScope::Full,
|
||||
false,
|
||||
);
|
||||
|
||||
// This is a conceptual test — we need a way to express "requires workspace green"
|
||||
// Currently LaneContext has raw green_level: u8, not a contract
|
||||
// For now we just verify the policy condition works
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"workspace-green-required",
|
||||
PolicyCondition::GreenAt { level: 3 }, // GreenLevel::Workspace
|
||||
PolicyAction::MergeToDev,
|
||||
10,
|
||||
)]);
|
||||
|
||||
let actions = engine.evaluate(&context);
|
||||
assert!(actions.is_empty()); // level 1 < 3, so no merge
|
||||
}
|
||||
|
||||
/// reconciliation + policy_engine integration:
|
||||
/// A reconciled lane should be handled by reconcile rules, not generic closeout
|
||||
#[test]
|
||||
fn reconciled_lane_matches_reconcile_condition() {
|
||||
let context = LaneContext::reconciled("reconciled-lane");
|
||||
|
||||
let engine = PolicyEngine::new(vec![
|
||||
PolicyRule::new(
|
||||
"reconcile-first",
|
||||
PolicyCondition::LaneReconciled,
|
||||
PolicyAction::Reconcile {
|
||||
reason: ReconcileReason::AlreadyMerged,
|
||||
},
|
||||
5,
|
||||
),
|
||||
PolicyRule::new(
|
||||
"generic-closeout",
|
||||
PolicyCondition::LaneCompleted,
|
||||
PolicyAction::CloseoutLane,
|
||||
30,
|
||||
),
|
||||
]);
|
||||
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// Both rules fire — reconcile (priority 5) first, then closeout (priority 30)
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::Reconcile {
|
||||
reason: ReconcileReason::AlreadyMerged,
|
||||
},
|
||||
PolicyAction::CloseoutLane,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
/// stale_branch module: apply_policy generates correct actions
|
||||
#[test]
|
||||
fn stale_branch_apply_policy_produces_rebase_action() {
|
||||
let stale = BranchFreshness::Stale {
|
||||
commits_behind: 5,
|
||||
missing_fixes: vec!["fix-123".to_string()],
|
||||
};
|
||||
|
||||
let action = apply_policy(&stale, StaleBranchPolicy::AutoRebase);
|
||||
assert_eq!(action, StaleBranchAction::Rebase);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_branch_apply_policy_produces_merge_forward_action() {
|
||||
let stale = BranchFreshness::Stale {
|
||||
commits_behind: 3,
|
||||
missing_fixes: vec![],
|
||||
};
|
||||
|
||||
let action = apply_policy(&stale, StaleBranchPolicy::AutoMergeForward);
|
||||
assert_eq!(action, StaleBranchAction::MergeForward);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_branch_apply_policy_warn_only() {
|
||||
let stale = BranchFreshness::Stale {
|
||||
commits_behind: 2,
|
||||
missing_fixes: vec!["fix-456".to_string()],
|
||||
};
|
||||
|
||||
let action = apply_policy(&stale, StaleBranchPolicy::WarnOnly);
|
||||
match action {
|
||||
StaleBranchAction::Warn { message } => {
|
||||
assert!(message.contains("2 commit(s) behind main"));
|
||||
assert!(message.contains("fix-456"));
|
||||
}
|
||||
_ => panic!("expected Warn action, got {:?}", action),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stale_branch_fresh_produces_noop() {
|
||||
let fresh = BranchFreshness::Fresh;
|
||||
let action = apply_policy(&fresh, StaleBranchPolicy::AutoRebase);
|
||||
assert_eq!(action, StaleBranchAction::Noop);
|
||||
}
|
||||
|
||||
/// Combined flow: stale detection + policy + action
|
||||
#[test]
|
||||
fn end_to_end_stale_lane_gets_merge_forward_action() {
|
||||
// Simulating what a harness would do:
|
||||
// 1. Detect branch freshness
|
||||
// 2. Build lane context from freshness + other signals
|
||||
// 3. Run policy engine
|
||||
// 4. Return actions
|
||||
|
||||
// given: detected stale state
|
||||
let _freshness = BranchFreshness::Stale {
|
||||
commits_behind: 5,
|
||||
missing_fixes: vec!["fix-123".to_string()],
|
||||
};
|
||||
|
||||
// when: build context and evaluate policy
|
||||
let context = LaneContext::new(
|
||||
"lane-9411",
|
||||
3, // Workspace green
|
||||
Duration::from_secs(5 * 60 * 60), // 5 hours stale, definitely over threshold
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Approved,
|
||||
DiffScope::Scoped,
|
||||
false,
|
||||
);
|
||||
|
||||
let engine = PolicyEngine::new(vec![
|
||||
// Priority 5: Check if stale first
|
||||
PolicyRule::new(
|
||||
"auto-merge-forward-if-stale-and-approved",
|
||||
PolicyCondition::And(vec![
|
||||
PolicyCondition::StaleBranch,
|
||||
PolicyCondition::ReviewPassed,
|
||||
]),
|
||||
PolicyAction::MergeForward,
|
||||
5,
|
||||
),
|
||||
// Priority 10: Normal stale handling
|
||||
PolicyRule::new(
|
||||
"stale-warning",
|
||||
PolicyCondition::StaleBranch,
|
||||
PolicyAction::Notify {
|
||||
channel: "#build-status".to_string(),
|
||||
},
|
||||
10,
|
||||
),
|
||||
]);
|
||||
|
||||
let actions = engine.evaluate(&context);
|
||||
|
||||
// then: both rules should fire (stale + approved matches both)
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![
|
||||
PolicyAction::MergeForward,
|
||||
PolicyAction::Notify {
|
||||
channel: "#build-status".to_string(),
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
/// Fresh branch with approved review should merge (not stale-blocked)
|
||||
#[test]
|
||||
fn fresh_approved_lane_gets_merge_action() {
|
||||
let context = LaneContext::new(
|
||||
"fresh-approved-lane",
|
||||
3, // Workspace green
|
||||
Duration::from_secs(30 * 60), // 30 min — under 1 hour threshold = fresh
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Approved,
|
||||
DiffScope::Scoped,
|
||||
false,
|
||||
);
|
||||
|
||||
let engine = PolicyEngine::new(vec![PolicyRule::new(
|
||||
"merge-if-green-approved-not-stale",
|
||||
PolicyCondition::And(vec![
|
||||
PolicyCondition::GreenAt { level: 3 },
|
||||
PolicyCondition::ReviewPassed,
|
||||
// NOT PolicyCondition::StaleBranch — fresh lanes bypass this
|
||||
]),
|
||||
PolicyAction::MergeToDev,
|
||||
5,
|
||||
)]);
|
||||
|
||||
let actions = engine.evaluate(&context);
|
||||
assert_eq!(actions, vec![PolicyAction::MergeToDev]);
|
||||
}
|
||||
|
||||
/// worker_boot + recovery_recipes + policy_engine integration:
|
||||
/// When a session completes with a provider failure, does the worker
|
||||
/// status transition trigger the correct recovery recipe, and does
|
||||
/// the resulting recovery state feed into policy decisions?
|
||||
#[test]
|
||||
fn worker_provider_failure_flows_through_recovery_to_policy() {
|
||||
use runtime::recovery_recipes::{
|
||||
attempt_recovery, FailureScenario, RecoveryContext, RecoveryResult, RecoveryStep,
|
||||
};
|
||||
use runtime::worker_boot::{WorkerFailureKind, WorkerRegistry, WorkerStatus};
|
||||
|
||||
// given — a worker that encounters a provider failure during session completion
|
||||
let registry = WorkerRegistry::new();
|
||||
let worker = registry.create("/tmp/repo-recovery-test", &[], true);
|
||||
|
||||
// Worker reaches ready state
|
||||
registry
|
||||
.observe(&worker.worker_id, "Ready for your input\n>")
|
||||
.expect("ready observe should succeed");
|
||||
registry
|
||||
.send_prompt(&worker.worker_id, Some("Run analysis"))
|
||||
.expect("prompt send should succeed");
|
||||
|
||||
// Session completes with provider failure (finish="unknown", tokens=0)
|
||||
let failed_worker = registry
|
||||
.observe_completion(&worker.worker_id, "unknown", 0)
|
||||
.expect("completion observe should succeed");
|
||||
assert_eq!(failed_worker.status, WorkerStatus::Failed);
|
||||
let failure = failed_worker
|
||||
.last_error
|
||||
.expect("worker should have recorded error");
|
||||
assert_eq!(failure.kind, WorkerFailureKind::Provider);
|
||||
|
||||
// Bridge: WorkerFailureKind -> FailureScenario
|
||||
let scenario = FailureScenario::from_worker_failure_kind(failure.kind);
|
||||
assert_eq!(scenario, FailureScenario::ProviderFailure);
|
||||
|
||||
// Recovery recipe lookup and execution
|
||||
let mut ctx = RecoveryContext::new();
|
||||
let result = attempt_recovery(&scenario, &mut ctx);
|
||||
|
||||
// then — recovery should recommend RestartWorker step
|
||||
assert!(
|
||||
matches!(result, RecoveryResult::Recovered { steps_taken: 1 }),
|
||||
"provider failure should recover via single RestartWorker step, got: {result:?}"
|
||||
);
|
||||
assert!(
|
||||
ctx.events().iter().any(|e| {
|
||||
matches!(
|
||||
e,
|
||||
runtime::recovery_recipes::RecoveryEvent::RecoveryAttempted {
|
||||
result: RecoveryResult::Recovered { steps_taken: 1 },
|
||||
..
|
||||
}
|
||||
)
|
||||
}),
|
||||
"recovery should emit structured attempt event"
|
||||
);
|
||||
|
||||
// Policy integration: recovery success + green status = merge-ready
|
||||
// (Simulating the policy check that would happen after successful recovery)
|
||||
let recovery_success = matches!(result, RecoveryResult::Recovered { .. });
|
||||
let green_level = 3; // Workspace green
|
||||
let not_stale = Duration::from_secs(30 * 60); // 30 min — fresh
|
||||
|
||||
let post_recovery_context = LaneContext::new(
|
||||
"recovered-lane",
|
||||
green_level,
|
||||
not_stale,
|
||||
LaneBlocker::None,
|
||||
ReviewStatus::Approved,
|
||||
DiffScope::Scoped,
|
||||
false,
|
||||
);
|
||||
|
||||
let policy_engine = PolicyEngine::new(vec![
|
||||
// Rule: if recovered from failure + green + approved -> merge
|
||||
PolicyRule::new(
|
||||
"merge-after-successful-recovery",
|
||||
PolicyCondition::And(vec![
|
||||
PolicyCondition::GreenAt { level: 3 },
|
||||
PolicyCondition::ReviewPassed,
|
||||
]),
|
||||
PolicyAction::MergeToDev,
|
||||
10,
|
||||
),
|
||||
]);
|
||||
|
||||
// Recovery success is a pre-condition; policy evaluates post-recovery context
|
||||
assert!(
|
||||
recovery_success,
|
||||
"recovery must succeed for lane to proceed"
|
||||
);
|
||||
let actions = policy_engine.evaluate(&post_recovery_context);
|
||||
assert_eq!(
|
||||
actions,
|
||||
vec![PolicyAction::MergeToDev],
|
||||
"post-recovery green+approved lane should be merge-ready"
|
||||
);
|
||||
}
|
||||
@@ -1,34 +0,0 @@
|
||||
[package]
|
||||
name = "rusty-claude-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"
|
||||
rustyline = "15"
|
||||
runtime = { path = "../runtime" }
|
||||
plugins = { path = "../plugins" }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json.workspace = true
|
||||
syntect = "5"
|
||||
tokio = { version = "1", features = ["rt-multi-thread", "signal", "time"] }
|
||||
tools = { path = "../tools" }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
mock-anthropic-service = { path = "../mock-anthropic-service" }
|
||||
serde_json.workspace = true
|
||||
tokio = { version = "1", features = ["rt-multi-thread"] }
|
||||
|
||||
@@ -1,64 +0,0 @@
|
||||
# Rusty Claude CLI 模块 (rusty-claude-cli)
|
||||
|
||||
本模块提供了 Claw 命令行界面的另一个功能完整的实现。它集成了对话、工具执行、插件扩展以及身份验证等核心功能。
|
||||
|
||||
## 概览
|
||||
|
||||
`rusty-claude-cli` 是一个全功能的 CLI 应用程序,其主要职责包括:
|
||||
- **用户交互**:提供交互式 REPL 和非交互式命令执行(`prompt` 子命令)。
|
||||
- **环境初始化**:处理项目初始化 (`init`) 和配置加载。
|
||||
- **身份验证**:通过本地回环服务器处理 OAuth 登录流程。
|
||||
- **状态渲染**:实现丰富的终端 UI 效果,如 Markdown 渲染、语法高亮和动态加载动画 (Spinner)。
|
||||
- **会话管理**:支持从保存的文件中恢复会话并执行追加的斜杠命令。
|
||||
|
||||
## 与 `claw-cli` 的关系
|
||||
|
||||
虽然 `rusty-claude-cli` 和 `claw-cli` 都生成名为 `claw` 的二进制文件,但 `rusty-claude-cli` 包含更复杂的集成逻辑:
|
||||
- 它直接引用了几乎所有的核心 crate(`runtime`, `api`, `tools`, `plugins`, `commands`)。
|
||||
- 它的 `main.rs` 实现非常庞大,包含了大量的业务编排逻辑。
|
||||
- 它可以作为一个独立的、集成度极高的 CLI 参考实现。
|
||||
|
||||
## 关键特性
|
||||
|
||||
- **多功能子命令**:
|
||||
- `prompt`:快速运行单次推理。
|
||||
- `login`/`logout`:OAuth 身份验证管理。
|
||||
- `init`:项目环境自举。
|
||||
- `bootstrap-plan`:查看系统的启动阶段。
|
||||
- `dump-manifests`:从上游源码中提取并显示功能清单。
|
||||
- **增强的 REPL**:
|
||||
- 支持多行输入和历史记录。
|
||||
- 集成了斜杠命令处理引擎。
|
||||
- 提供详细的消耗统计和权限模式切换报告。
|
||||
- **灵活的权限控制**:支持通过命令行参数 `--permission-mode` 或环境变量动态调整权限级别。
|
||||
|
||||
## 实现逻辑
|
||||
|
||||
### 核心子模块
|
||||
|
||||
- **`main.rs`**: 程序的入口,包含了复杂的参数解析逻辑和 REPL 循环。
|
||||
- **`render.rs`**: 封装了 `TerminalRenderer` 和 `Spinner`,负责所有的终端输出美化。
|
||||
- **`input.rs`**: 处理从标准输入读取数据及命令解析。
|
||||
- **`init.rs`**: 专注于仓库的初始化和 `.claw.md` 文件的生成。
|
||||
- **`app.rs`**: 可能包含应用程序级别的高层状态管理(取决于具体实现)。
|
||||
|
||||
### 工作流程
|
||||
|
||||
1. 程序启动,解析命令行参数。
|
||||
2. 根据参数决定是执行单次任务还是进入 REPL 模式。
|
||||
3. 在 REPL 模式下,初始化 `ConversationRuntime`。
|
||||
4. 进入循环:读取用户输入 -> 处理斜杠命令或发送给 AI -> 渲染响应 -> 执行工具 -> 循环。
|
||||
5. 会话数据根据需要保存或恢复。
|
||||
|
||||
## 使用示例
|
||||
|
||||
```bash
|
||||
# 启动交互模式
|
||||
cargo run -p rusty-claude-cli --bin claw
|
||||
|
||||
# 直接运行 Prompt
|
||||
cargo run -p rusty-claude-cli --bin claw prompt "检查代码中的内存泄漏"
|
||||
|
||||
# 恢复之前的会话并执行压缩
|
||||
cargo run -p rusty-claude-cli --bin claw --resume session.json /compact
|
||||
```
|
||||
@@ -1,38 +0,0 @@
|
||||
use std::env;
|
||||
use std::process::Command;
|
||||
|
||||
fn main() {
|
||||
// Get git SHA (short hash)
|
||||
let git_sha = Command::new("git")
|
||||
.args(["rev-parse", "--short", "HEAD"])
|
||||
.output()
|
||||
.ok()
|
||||
.and_then(|output| {
|
||||
if output.status.success() {
|
||||
String::from_utf8(output.stdout).ok()
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.map_or_else(|| "unknown".to_string(), |s| s.trim().to_string());
|
||||
|
||||
println!("cargo:rustc-env=GIT_SHA={git_sha}");
|
||||
|
||||
// TARGET is always set by Cargo during build
|
||||
let target = env::var("TARGET").unwrap_or_else(|_| "unknown".to_string());
|
||||
println!("cargo:rustc-env=TARGET={target}");
|
||||
|
||||
// Build date from SOURCE_DATE_EPOCH (reproducible builds) or current date.
|
||||
let build_date = std::env::var("SOURCE_DATE_EPOCH")
|
||||
.ok()
|
||||
.and_then(|epoch| epoch.parse::<i64>().ok())
|
||||
.map_or_else(
|
||||
|| "unknown".to_string(),
|
||||
|_| std::env::var("BUILD_DATE").unwrap_or_else(|_| "unknown".to_string()),
|
||||
);
|
||||
println!("cargo:rustc-env=BUILD_DATE={build_date}");
|
||||
|
||||
// Rerun if git state changes
|
||||
println!("cargo:rerun-if-changed=.git/HEAD");
|
||||
println!("cargo:rerun-if-changed=.git/refs");
|
||||
}
|
||||
@@ -1,330 +0,0 @@
|
||||
use std::borrow::Cow;
|
||||
use std::cell::RefCell;
|
||||
use std::collections::BTreeSet;
|
||||
use std::io::{self, IsTerminal, Write};
|
||||
|
||||
use rustyline::completion::{Completer, Pair};
|
||||
use rustyline::error::ReadlineError;
|
||||
use rustyline::highlight::{CmdKind, Highlighter};
|
||||
use rustyline::hint::Hinter;
|
||||
use rustyline::history::DefaultHistory;
|
||||
use rustyline::validate::Validator;
|
||||
use rustyline::{
|
||||
Cmd, CompletionType, Config, Context, EditMode, Editor, Helper, KeyCode, KeyEvent, Modifiers,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum ReadOutcome {
|
||||
Submit(String),
|
||||
Cancel,
|
||||
Exit,
|
||||
}
|
||||
|
||||
struct SlashCommandHelper {
|
||||
completions: Vec<String>,
|
||||
current_line: RefCell<String>,
|
||||
}
|
||||
|
||||
impl SlashCommandHelper {
|
||||
fn new(completions: Vec<String>) -> Self {
|
||||
Self {
|
||||
completions: normalize_completions(completions),
|
||||
current_line: RefCell::new(String::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn reset_current_line(&self) {
|
||||
self.current_line.borrow_mut().clear();
|
||||
}
|
||||
|
||||
fn current_line(&self) -> String {
|
||||
self.current_line.borrow().clone()
|
||||
}
|
||||
|
||||
fn set_current_line(&self, line: &str) {
|
||||
let mut current = self.current_line.borrow_mut();
|
||||
current.clear();
|
||||
current.push_str(line);
|
||||
}
|
||||
|
||||
fn set_completions(&mut self, completions: Vec<String>) {
|
||||
self.completions = normalize_completions(completions);
|
||||
}
|
||||
}
|
||||
|
||||
impl Completer for SlashCommandHelper {
|
||||
type Candidate = Pair;
|
||||
|
||||
fn complete(
|
||||
&self,
|
||||
line: &str,
|
||||
pos: usize,
|
||||
_ctx: &Context<'_>,
|
||||
) -> rustyline::Result<(usize, Vec<Self::Candidate>)> {
|
||||
let Some(prefix) = slash_command_prefix(line, pos) else {
|
||||
return Ok((0, Vec::new()));
|
||||
};
|
||||
|
||||
let matches = self
|
||||
.completions
|
||||
.iter()
|
||||
.filter(|candidate| candidate.starts_with(prefix))
|
||||
.map(|candidate| Pair {
|
||||
display: candidate.clone(),
|
||||
replacement: candidate.clone(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok((0, matches))
|
||||
}
|
||||
}
|
||||
|
||||
impl Hinter for SlashCommandHelper {
|
||||
type Hint = String;
|
||||
}
|
||||
|
||||
impl Highlighter for SlashCommandHelper {
|
||||
fn highlight<'l>(&self, line: &'l str, _pos: usize) -> Cow<'l, str> {
|
||||
self.set_current_line(line);
|
||||
Cow::Borrowed(line)
|
||||
}
|
||||
|
||||
fn highlight_char(&self, line: &str, _pos: usize, _kind: CmdKind) -> bool {
|
||||
self.set_current_line(line);
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
impl Validator for SlashCommandHelper {}
|
||||
impl Helper for SlashCommandHelper {}
|
||||
|
||||
pub struct LineEditor {
|
||||
prompt: String,
|
||||
editor: Editor<SlashCommandHelper, DefaultHistory>,
|
||||
}
|
||||
|
||||
impl LineEditor {
|
||||
#[must_use]
|
||||
pub fn new(prompt: impl Into<String>, completions: Vec<String>) -> Self {
|
||||
let config = Config::builder()
|
||||
.completion_type(CompletionType::List)
|
||||
.edit_mode(EditMode::Emacs)
|
||||
.build();
|
||||
let mut editor = Editor::<SlashCommandHelper, DefaultHistory>::with_config(config)
|
||||
.expect("rustyline editor should initialize");
|
||||
editor.set_helper(Some(SlashCommandHelper::new(completions)));
|
||||
editor.bind_sequence(KeyEvent(KeyCode::Char('J'), Modifiers::CTRL), Cmd::Newline);
|
||||
editor.bind_sequence(KeyEvent(KeyCode::Enter, Modifiers::SHIFT), Cmd::Newline);
|
||||
|
||||
Self {
|
||||
prompt: prompt.into(),
|
||||
editor,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn push_history(&mut self, entry: impl Into<String>) {
|
||||
let entry = entry.into();
|
||||
if entry.trim().is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let _ = self.editor.add_history_entry(entry);
|
||||
}
|
||||
|
||||
pub fn set_completions(&mut self, completions: Vec<String>) {
|
||||
if let Some(helper) = self.editor.helper_mut() {
|
||||
helper.set_completions(completions);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn read_line(&mut self) -> io::Result<ReadOutcome> {
|
||||
if !io::stdin().is_terminal() || !io::stdout().is_terminal() {
|
||||
return self.read_line_fallback();
|
||||
}
|
||||
|
||||
if let Some(helper) = self.editor.helper_mut() {
|
||||
helper.reset_current_line();
|
||||
}
|
||||
|
||||
match self.editor.readline(&self.prompt) {
|
||||
Ok(line) => Ok(ReadOutcome::Submit(line)),
|
||||
Err(ReadlineError::Interrupted) => {
|
||||
let has_input = !self.current_line().is_empty();
|
||||
self.finish_interrupted_read()?;
|
||||
if has_input {
|
||||
Ok(ReadOutcome::Cancel)
|
||||
} else {
|
||||
Ok(ReadOutcome::Exit)
|
||||
}
|
||||
}
|
||||
Err(ReadlineError::Eof) => {
|
||||
self.finish_interrupted_read()?;
|
||||
Ok(ReadOutcome::Exit)
|
||||
}
|
||||
Err(error) => Err(io::Error::other(error)),
|
||||
}
|
||||
}
|
||||
|
||||
fn current_line(&self) -> String {
|
||||
self.editor
|
||||
.helper()
|
||||
.map_or_else(String::new, SlashCommandHelper::current_line)
|
||||
}
|
||||
|
||||
fn finish_interrupted_read(&mut self) -> io::Result<()> {
|
||||
if let Some(helper) = self.editor.helper_mut() {
|
||||
helper.reset_current_line();
|
||||
}
|
||||
let mut stdout = io::stdout();
|
||||
writeln!(stdout)
|
||||
}
|
||||
|
||||
fn read_line_fallback(&self) -> io::Result<ReadOutcome> {
|
||||
let mut stdout = io::stdout();
|
||||
write!(stdout, "{}", self.prompt)?;
|
||||
stdout.flush()?;
|
||||
|
||||
let mut buffer = String::new();
|
||||
let bytes_read = io::stdin().read_line(&mut buffer)?;
|
||||
if bytes_read == 0 {
|
||||
return Ok(ReadOutcome::Exit);
|
||||
}
|
||||
|
||||
while matches!(buffer.chars().last(), Some('\n' | '\r')) {
|
||||
buffer.pop();
|
||||
}
|
||||
Ok(ReadOutcome::Submit(buffer))
|
||||
}
|
||||
}
|
||||
|
||||
fn slash_command_prefix(line: &str, pos: usize) -> Option<&str> {
|
||||
if pos != line.len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let prefix = &line[..pos];
|
||||
if !prefix.starts_with('/') {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(prefix)
|
||||
}
|
||||
|
||||
fn normalize_completions(completions: Vec<String>) -> Vec<String> {
|
||||
let mut seen = BTreeSet::new();
|
||||
completions
|
||||
.into_iter()
|
||||
.filter(|candidate| candidate.starts_with('/'))
|
||||
.filter(|candidate| seen.insert(candidate.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{slash_command_prefix, LineEditor, SlashCommandHelper};
|
||||
use rustyline::completion::Completer;
|
||||
use rustyline::highlight::Highlighter;
|
||||
use rustyline::history::{DefaultHistory, History};
|
||||
use rustyline::Context;
|
||||
|
||||
#[test]
|
||||
fn extracts_terminal_slash_command_prefixes_with_arguments() {
|
||||
assert_eq!(slash_command_prefix("/he", 3), Some("/he"));
|
||||
assert_eq!(slash_command_prefix("/help me", 8), Some("/help me"));
|
||||
assert_eq!(
|
||||
slash_command_prefix("/session switch ses", 19),
|
||||
Some("/session switch ses")
|
||||
);
|
||||
assert_eq!(slash_command_prefix("hello", 5), None);
|
||||
assert_eq!(slash_command_prefix("/help", 2), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completes_matching_slash_commands() {
|
||||
let helper = SlashCommandHelper::new(vec![
|
||||
"/help".to_string(),
|
||||
"/hello".to_string(),
|
||||
"/status".to_string(),
|
||||
]);
|
||||
let history = DefaultHistory::new();
|
||||
let ctx = Context::new(&history);
|
||||
let (start, matches) = helper
|
||||
.complete("/he", 3, &ctx)
|
||||
.expect("completion should work");
|
||||
|
||||
assert_eq!(start, 0);
|
||||
assert_eq!(
|
||||
matches
|
||||
.into_iter()
|
||||
.map(|candidate| candidate.replacement)
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["/help".to_string(), "/hello".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completes_matching_slash_command_arguments() {
|
||||
let helper = SlashCommandHelper::new(vec![
|
||||
"/model".to_string(),
|
||||
"/model opus".to_string(),
|
||||
"/model sonnet".to_string(),
|
||||
"/session switch alpha".to_string(),
|
||||
]);
|
||||
let history = DefaultHistory::new();
|
||||
let ctx = Context::new(&history);
|
||||
let (start, matches) = helper
|
||||
.complete("/model o", 8, &ctx)
|
||||
.expect("completion should work");
|
||||
|
||||
assert_eq!(start, 0);
|
||||
assert_eq!(
|
||||
matches
|
||||
.into_iter()
|
||||
.map(|candidate| candidate.replacement)
|
||||
.collect::<Vec<_>>(),
|
||||
vec!["/model opus".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ignores_non_slash_command_completion_requests() {
|
||||
let helper = SlashCommandHelper::new(vec!["/help".to_string()]);
|
||||
let history = DefaultHistory::new();
|
||||
let ctx = Context::new(&history);
|
||||
let (_, matches) = helper
|
||||
.complete("hello", 5, &ctx)
|
||||
.expect("completion should work");
|
||||
|
||||
assert!(matches.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracks_current_buffer_through_highlighter() {
|
||||
let helper = SlashCommandHelper::new(Vec::new());
|
||||
let _ = helper.highlight("draft", 5);
|
||||
|
||||
assert_eq!(helper.current_line(), "draft");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn push_history_ignores_blank_entries() {
|
||||
let mut editor = LineEditor::new("> ", vec!["/help".to_string()]);
|
||||
editor.push_history(" ");
|
||||
editor.push_history("/help");
|
||||
|
||||
assert_eq!(editor.editor.history().len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_completions_replaces_and_normalizes_candidates() {
|
||||
let mut editor = LineEditor::new("> ", vec!["/help".to_string()]);
|
||||
editor.set_completions(vec![
|
||||
"/model opus".to_string(),
|
||||
"/model opus".to_string(),
|
||||
"status".to_string(),
|
||||
]);
|
||||
|
||||
let helper = editor.editor.helper().expect("helper should exist");
|
||||
assert_eq!(helper.completions, vec!["/model opus".to_string()]);
|
||||
}
|
||||
}
|
||||