feat(desktop): embed BoosAPI Desktop client as a subproject
Vendor the rebranded BoosAPI Desktop client (Tauri + Vue) into metapi as a monorepo subdirectory under desktop/. Source only — node_modules, dist, Cargo target and the upstream .git history are excluded. desktop/ keeps its own independent Cargo workspace (desktop/src-tauri with crates/*); it is not part of the backend boosapi crate, so a `cargo check` at the metapi root still compiles only the backend. Build the desktop app from desktop/src-tauri per desktop/README.md. Note: app icons are still placeholder assets pending the final BoosAPI logo. The updater signing keys live outside the repo and are not included.
This commit is contained in:
boos
@@ -0,0 +1,21 @@
|
||||
/target
|
||||
Cargo.lock
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# Environment
|
||||
.env
|
||||
.env.*
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Credentials
|
||||
*.pem
|
||||
*.key
|
||||
@@ -0,0 +1,78 @@
|
||||
[package]
|
||||
name = "open-agent-sdk"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
description = "Open Agent SDK - A Rust framework for building autonomous AI agents"
|
||||
license = "MIT"
|
||||
repository = "https://github.com/codeany-ai/open-agent-sdk-rust"
|
||||
|
||||
[dependencies]
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
tokio-stream = "0.1"
|
||||
futures = "0.3"
|
||||
reqwest = { version = "0.12", features = ["json", "stream"] }
|
||||
serde = { version = "1", features = ["derive"] }
|
||||
serde_json = "1"
|
||||
uuid = { version = "1", features = ["v4"] }
|
||||
glob = "0.3"
|
||||
regex = "1"
|
||||
chrono = "0.4"
|
||||
async-trait = "0.1"
|
||||
thiserror = "2"
|
||||
bytes = "1"
|
||||
pin-project-lite = "0.2"
|
||||
tracing = "0.1"
|
||||
base64 = "0.22"
|
||||
tokio-util = "0.7"
|
||||
serde_yaml = "0.9"
|
||||
|
||||
[target.'cfg(target_os = "windows")'.dependencies]
|
||||
windows-sys = { version = "0.59", features = ["Win32_System_Threading"] }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { version = "1", features = ["full", "test-util"] }
|
||||
tempfile = "3"
|
||||
|
||||
[[example]]
|
||||
name = "01-simple-query"
|
||||
path = "examples/01-simple-query.rs"
|
||||
|
||||
[[example]]
|
||||
name = "02-multi-tool"
|
||||
path = "examples/02-multi-tool.rs"
|
||||
|
||||
[[example]]
|
||||
name = "03-multi-turn"
|
||||
path = "examples/03-multi-turn.rs"
|
||||
|
||||
[[example]]
|
||||
name = "04-prompt-api"
|
||||
path = "examples/04-prompt-api.rs"
|
||||
|
||||
[[example]]
|
||||
name = "05-custom-system-prompt"
|
||||
path = "examples/05-custom-system-prompt.rs"
|
||||
|
||||
[[example]]
|
||||
name = "06-mcp-server"
|
||||
path = "examples/06-mcp-server.rs"
|
||||
|
||||
[[example]]
|
||||
name = "07-custom-tools"
|
||||
path = "examples/07-custom-tools.rs"
|
||||
|
||||
[[example]]
|
||||
name = "08-oneshot-query"
|
||||
path = "examples/08-oneshot-query.rs"
|
||||
|
||||
[[example]]
|
||||
name = "09-subagents"
|
||||
path = "examples/09-subagents.rs"
|
||||
|
||||
[[example]]
|
||||
name = "10-permissions"
|
||||
path = "examples/10-permissions.rs"
|
||||
|
||||
[[example]]
|
||||
name = "web-chat"
|
||||
path = "examples/web/main.rs"
|
||||
@@ -0,0 +1,265 @@
|
||||
# Open Agent SDK (Rust)
|
||||
|
||||
A lightweight, open-source Rust SDK for building AI agents. Run the full agent loop in-process — no CLI or subprocess required. Deploy anywhere: cloud, serverless, Docker, CI/CD.
|
||||
|
||||
Also available in [TypeScript](https://github.com/codeany-ai/open-agent-sdk-typescript) and [Go](https://github.com/codeany-ai/open-agent-sdk-go).
|
||||
|
||||
## Features
|
||||
|
||||
- **Multi-Provider** — Supports both Anthropic Messages API and OpenAI Chat Completions API, with auto-detection by model name
|
||||
- **Agent Loop** — Streaming agentic loop with tool execution, multi-turn conversations, and cost tracking
|
||||
- **25+ Built-in Tools** — Bash, Read, Write, Edit, Glob, Grep, WebFetch, WebSearch, Agent (subagents), AskUser, Tasks, Teams, Plans, Worktrees, Todos, Cron, LSP, Config, MCP Resources, ToolSearch
|
||||
- **Session Persistence** — Save, load, fork, and resume conversations across sessions
|
||||
- **MCP Support** — Connect to MCP servers via stdio, HTTP, and SSE transports
|
||||
- **Permission System** — Configurable tool approval with allow/deny rules and filesystem path validation
|
||||
- **Hook System** — 20+ hook events: PreToolUse, PostToolUse, SessionStart/End, SubagentStart/Stop, TaskCreated/Completed, FileChanged, PreCompact/PostCompact, and more
|
||||
- **Extended Thinking** — Support for extended thinking with budget tokens
|
||||
- **Cost Tracking** — Per-model token usage and pricing for Anthropic, OpenAI, and DeepSeek models
|
||||
- **Custom Tools** — Implement the `Tool` trait to add your own tools
|
||||
- **Auto-Compaction** — Multi-tier context compression (micro-compact, LLM summarization, image stripping)
|
||||
- **Sandbox Config** — Network and filesystem restrictions for tool execution
|
||||
- **File State Cache** — LRU cache for file state tracking and staleness detection
|
||||
|
||||
## Quick Start
|
||||
|
||||
```toml
|
||||
[dependencies]
|
||||
open-agent-sdk = "0.1.0"
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
```
|
||||
|
||||
```rust
|
||||
use open_agent_sdk::{Agent, AgentOptions, SDKMessage};
|
||||
use open_agent_sdk::types;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let mut agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
|
||||
// Streaming
|
||||
let (mut rx, handle) = agent.query("What files are in this directory?").await;
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
print!("{}", types::extract_text(&message));
|
||||
}
|
||||
SDKMessage::Result { text, cost_usd, .. } => {
|
||||
println!("\nDone (${:.4})", cost_usd);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
handle.await.unwrap();
|
||||
|
||||
// Or use the blocking API
|
||||
let result = agent.prompt("Count lines in Cargo.toml").await.unwrap();
|
||||
println!("{}", result.text);
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
| # | Example | Description |
|
||||
| --- | --------------------------------------------------------------- | ------------------------------------------------ |
|
||||
| 01 | [Simple Query](examples/01-simple-query.rs) | Streaming query with tool calls |
|
||||
| 02 | [Multi-Tool](examples/02-multi-tool.rs) | Glob + Bash multi-tool orchestration |
|
||||
| 03 | [Multi-Turn](examples/03-multi-turn.rs) | Multi-turn conversation with session persistence |
|
||||
| 04 | [Prompt API](examples/04-prompt-api.rs) | Blocking `prompt()` for one-shot queries |
|
||||
| 05 | [Custom System Prompt](examples/05-custom-system-prompt.rs) | Custom system prompt for code review |
|
||||
| 06 | [MCP Server](examples/06-mcp-server.rs) | MCP server integration (stdio transport) |
|
||||
| 07 | [Custom Tools](examples/07-custom-tools.rs) | Define and use custom tools |
|
||||
| 08 | [One-shot Query](examples/08-oneshot-query.rs) | Quick one-shot agent query |
|
||||
| 09 | [Subagents](examples/09-subagents.rs) | Specialized subagent with restricted tools |
|
||||
| 10 | [Permissions](examples/10-permissions.rs) | Read-only agent with allowed tools |
|
||||
| 11 | [Web Chat](examples/web/) | Web-based chat UI with streaming |
|
||||
|
||||
Run any example:
|
||||
|
||||
```bash
|
||||
export CODEANY_BASE_URL=https://openrouter.ai/api
|
||||
export CODEANY_API_KEY=your-api-key
|
||||
export CODEANY_MODEL=anthropic/claude-sonnet-4
|
||||
cargo run --example 01-simple-query
|
||||
```
|
||||
|
||||
For the web chat UI:
|
||||
|
||||
```bash
|
||||
cargo run --example web-chat
|
||||
# Open http://localhost:8082
|
||||
```
|
||||
|
||||
## Custom Tools
|
||||
|
||||
Implement the `Tool` trait:
|
||||
|
||||
```rust
|
||||
use async_trait::async_trait;
|
||||
use open_agent_sdk::*;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
struct MyTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for MyTool {
|
||||
fn name(&self) -> &str { "MyTool" }
|
||||
fn description(&self) -> &str { "Does something useful" }
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([(
|
||||
"input".to_string(),
|
||||
json!({"type": "string", "description": "The input"}),
|
||||
)]),
|
||||
required: vec!["input".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
fn is_read_only(&self, _: &Value) -> bool { true }
|
||||
async fn call(&self, input: Value, _ctx: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
Ok(ToolResult::text("result"))
|
||||
}
|
||||
}
|
||||
|
||||
// Use it
|
||||
let agent = Agent::new(AgentOptions {
|
||||
custom_tools: vec![Arc::new(MyTool)],
|
||||
..Default::default()
|
||||
}).await.unwrap();
|
||||
```
|
||||
|
||||
## Session Persistence
|
||||
|
||||
```rust
|
||||
use open_agent_sdk::session::*;
|
||||
|
||||
// Save current session
|
||||
save_session(&session_data).await.unwrap();
|
||||
|
||||
// List all sessions
|
||||
let sessions = list_sessions().await.unwrap();
|
||||
|
||||
// Resume a session
|
||||
let data = load_session("session-id").await.unwrap();
|
||||
|
||||
// Fork a session
|
||||
let new_id = fork_session("session-id").await.unwrap();
|
||||
```
|
||||
|
||||
## MCP Servers
|
||||
|
||||
```rust
|
||||
use open_agent_sdk::types::McpServerConfig;
|
||||
|
||||
let agent = Agent::new(AgentOptions {
|
||||
mcp_servers: HashMap::from([(
|
||||
"filesystem".to_string(),
|
||||
McpServerConfig::Stdio {
|
||||
command: "npx".to_string(),
|
||||
args: vec!["-y".to_string(), "@modelcontextprotocol/server-filesystem".to_string(), "/tmp".to_string()],
|
||||
env: HashMap::new(),
|
||||
},
|
||||
)]),
|
||||
..Default::default()
|
||||
}).await.unwrap();
|
||||
```
|
||||
|
||||
## Built-in Tools
|
||||
|
||||
| Tool | Description | Read-Only |
|
||||
|------|-------------|-----------|
|
||||
| **Bash** | Execute shell commands | No |
|
||||
| **Read** | Read files with line numbers | Yes |
|
||||
| **Write** | Create/overwrite files | No |
|
||||
| **Edit** | String replacement in files | No |
|
||||
| **Glob** | File pattern matching | Yes |
|
||||
| **Grep** | Regex search (ripgrep) | Yes |
|
||||
| **NotebookEdit** | Edit Jupyter notebook cells | No |
|
||||
| **WebFetch** | Fetch URL content | Yes |
|
||||
| **WebSearch** | Web search (pluggable) | Yes |
|
||||
| **AskUserQuestion** | Interactive user prompts | No |
|
||||
| **SendMessage** | Inter-agent messaging | No |
|
||||
| **TeamCreate/Delete** | Multi-agent team management | No |
|
||||
| **EnterPlanMode/ExitPlanMode** | Structured planning mode | No |
|
||||
| **EnterWorktree/ExitWorktree** | Git worktree isolation | No |
|
||||
| **TodoWrite** | Session todo list | No |
|
||||
| **TaskCreate/Get/List/Update/Stop/Output** | Task tracking (6 tools) | Mixed |
|
||||
| **CronCreate/Delete/List** | Scheduled task management | Mixed |
|
||||
| **Config** | Session configuration | Mixed |
|
||||
| **LSP** | Language Server Protocol operations | Yes |
|
||||
| **ListMcpResources/ReadMcpResource** | MCP resource access | Yes |
|
||||
| **ToolSearch** | Discover available tools | Yes |
|
||||
|
||||
## Multi-Provider Support
|
||||
|
||||
The SDK automatically detects which API format to use based on the model name:
|
||||
|
||||
| Models | API Format | Auto-detected |
|
||||
|--------|-----------|---------------|
|
||||
| Claude (sonnet, opus, haiku) | Anthropic Messages | Yes |
|
||||
| GPT-4o, O1, O3, O4 | OpenAI Chat Completions | Yes |
|
||||
| DeepSeek, Qwen, Mistral | OpenAI Chat Completions | Yes |
|
||||
| LLaMA, Gemma, Gemini, Yi, GLM | OpenAI Chat Completions | Yes |
|
||||
|
||||
Override with `CODEANY_API_TYPE=openai-completions` or `CODEANY_API_TYPE=anthropic-messages`.
|
||||
|
||||
Works with any OpenAI-compatible endpoint (OpenRouter, vLLM, Ollama, LiteLLM, etc.):
|
||||
|
||||
```bash
|
||||
export CODEANY_BASE_URL=https://openrouter.ai/api
|
||||
export CODEANY_API_KEY=your-key
|
||||
export CODEANY_MODEL=anthropic/claude-sonnet-4
|
||||
cargo run --example 01-simple-query
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
open-agent-sdk-rust/
|
||||
├── src/
|
||||
│ ├── agent/ # Agent loop, query engine, options
|
||||
│ ├── api/ # Multi-provider API abstraction
|
||||
│ │ ├── provider.rs # LLMProvider trait + auto-detection
|
||||
│ │ ├── anthropic.rs# Anthropic Messages API (SSE streaming)
|
||||
│ │ └── openai.rs # OpenAI Chat Completions API (SSE streaming)
|
||||
│ ├── types/ # Core types: Message, Tool, ContentBlock, MCP, Sandbox
|
||||
│ ├── tools/ # 25+ built-in tool implementations + registry + executor
|
||||
│ ├── mcp/ # MCP client (stdio, HTTP, SSE) + tool wrapping
|
||||
│ ├── permissions/ # Permission rules, filesystem validation
|
||||
│ ├── hooks/ # 20+ hook events with pre/post tool-use lifecycle
|
||||
│ ├── costtracker/ # Token usage and cost tracking (Anthropic, OpenAI, DeepSeek)
|
||||
│ ├── context/ # System/user context injection (git status, AGENT.md)
|
||||
│ ├── session/ # Session persistence (save, load, fork, resume)
|
||||
│ └── utils/ # Token estimation, retry, messages, compaction, file cache
|
||||
├── tests/ # Comprehensive test suite (116 tests)
|
||||
└── examples/ # 11 runnable examples + web chat UI
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
Environment variables:
|
||||
|
||||
| Variable | Description |
|
||||
| ---------------------------- | ----------------------------------------------------- |
|
||||
| `CODEANY_API_KEY` | API key (required) |
|
||||
| `CODEANY_MODEL` | Default model (default: `sonnet-4-6`) |
|
||||
| `CODEANY_BASE_URL` | API base URL override |
|
||||
| `CODEANY_API_TYPE` | Force API format: `anthropic-messages` or `openai-completions` |
|
||||
| `CODEANY_CUSTOM_HEADERS` | Custom headers (comma-separated `key:value`) |
|
||||
| `API_TIMEOUT_MS` | API request timeout in ms |
|
||||
| `HTTPS_PROXY` / `HTTP_PROXY` | Proxy URL |
|
||||
|
||||
## Links
|
||||
|
||||
- Website: [codeany.ai](https://codeany.ai)
|
||||
- TypeScript SDK: [github.com/codeany-ai/open-agent-sdk-typescript](https://github.com/codeany-ai/open-agent-sdk-typescript)
|
||||
- Go SDK: [github.com/codeany-ai/open-agent-sdk-go](https://github.com/codeany-ai/open-agent-sdk-go)
|
||||
- Issues: [github.com/codeany-ai/open-agent-sdk-rust/issues](https://github.com/codeany-ai/open-agent-sdk-rust/issues)
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,77 @@
|
||||
// Example 1: Simple Query with Streaming
|
||||
//
|
||||
// Demonstrates the basic Agent::new() + query() flow with
|
||||
// real-time event streaming.
|
||||
//
|
||||
// Run: cargo run --example 01-simple-query
|
||||
|
||||
use open_agent_sdk::{Agent, AgentOptions, SDKMessage};
|
||||
use open_agent_sdk::types;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 1: Simple Query ---");
|
||||
|
||||
let mut agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(10),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (mut rx, handle) = agent
|
||||
.query("Read Cargo.toml and tell me the project name and version in one sentence.")
|
||||
.await;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
for block in &message.content {
|
||||
match block {
|
||||
types::ContentBlock::ToolUse { id, name, input } => {
|
||||
let input_str = serde_json::to_string(input).unwrap_or_default();
|
||||
let truncated = if input_str.len() > 80 {
|
||||
format!("{}...", &input_str[..80])
|
||||
} else {
|
||||
input_str
|
||||
};
|
||||
println!("[Tool] {}({})", name, truncated);
|
||||
}
|
||||
types::ContentBlock::Text { text } => {
|
||||
if !text.trim().is_empty() {
|
||||
println!("\nAssistant: {}", text);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
SDKMessage::ToolResult {
|
||||
tool_name, content, ..
|
||||
} => {
|
||||
let truncated = if content.len() > 200 {
|
||||
format!("{}...", &content[..200])
|
||||
} else {
|
||||
content
|
||||
};
|
||||
println!("[Result: {}] {}", tool_name, truncated);
|
||||
}
|
||||
SDKMessage::Result {
|
||||
usage, ..
|
||||
} => {
|
||||
println!("\n--- Result ---");
|
||||
println!(
|
||||
"Tokens: {} in / {} out",
|
||||
usage.input_tokens, usage.output_tokens
|
||||
);
|
||||
}
|
||||
SDKMessage::Error { message } => {
|
||||
eprintln!("Error: {}", message);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
handle.await.unwrap();
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
// Example 2: Multi-Tool Orchestration
|
||||
//
|
||||
// The agent autonomously uses Glob, Bash, and Read to find files,
|
||||
// count lines, and provide a summary.
|
||||
//
|
||||
// Run: cargo run --example 02-multi-tool
|
||||
|
||||
use open_agent_sdk::{Agent, AgentOptions, SDKMessage};
|
||||
use open_agent_sdk::types;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 2: Multi-Tool ---");
|
||||
|
||||
let mut agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(15),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (mut rx, handle) = agent
|
||||
.query(
|
||||
"Find all .rs files in the src/ directory, count the total lines of code \
|
||||
with wc, and give me a brief summary.",
|
||||
)
|
||||
.await;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
for block in &message.content {
|
||||
match block {
|
||||
types::ContentBlock::ToolUse { name, input, .. } => {
|
||||
let input_str = serde_json::to_string(input).unwrap_or_default();
|
||||
let truncated = if input_str.len() > 80 {
|
||||
format!("{}...", &input_str[..80])
|
||||
} else {
|
||||
input_str
|
||||
};
|
||||
println!("[{}] {}", name, truncated);
|
||||
}
|
||||
types::ContentBlock::Text { text } => {
|
||||
if !text.trim().is_empty() {
|
||||
println!("\n{}", text);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
SDKMessage::Result { usage, .. } => {
|
||||
println!("\n--- Done ---");
|
||||
println!(
|
||||
"Tokens: {} in / {} out",
|
||||
usage.input_tokens, usage.output_tokens
|
||||
);
|
||||
}
|
||||
SDKMessage::Error { message } => {
|
||||
eprintln!("Error: {}", message);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
handle.await.unwrap();
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
// Example 3: Multi-Turn Conversation
|
||||
//
|
||||
// Three turns with context retention: create a file, read it back,
|
||||
// then delete it.
|
||||
//
|
||||
// Run: cargo run --example 03-multi-turn
|
||||
|
||||
use open_agent_sdk::{Agent, AgentOptions};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 3: Multi-Turn ---");
|
||||
|
||||
let mut agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(10),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Turn 1: Create a file
|
||||
println!("\n[Turn 1] Creating file...");
|
||||
match agent
|
||||
.prompt("Use Bash to create a file /tmp/sdk-test-rust.txt with the text 'Hello from Rust SDK!'")
|
||||
.await
|
||||
{
|
||||
Ok(result) => println!("Response: {}", result.text),
|
||||
Err(e) => eprintln!("Error: {}", e),
|
||||
}
|
||||
|
||||
// Turn 2: Read it back
|
||||
println!("\n[Turn 2] Reading file...");
|
||||
match agent
|
||||
.prompt("Now read /tmp/sdk-test-rust.txt and tell me what it says")
|
||||
.await
|
||||
{
|
||||
Ok(result) => println!("Response: {}", result.text),
|
||||
Err(e) => eprintln!("Error: {}", e),
|
||||
}
|
||||
|
||||
// Turn 3: Delete it
|
||||
println!("\n[Turn 3] Cleaning up...");
|
||||
match agent
|
||||
.prompt("Delete /tmp/sdk-test-rust.txt using Bash")
|
||||
.await
|
||||
{
|
||||
Ok(result) => {
|
||||
println!("Response: {}", result.text);
|
||||
println!("\nTotal messages: {}", agent.get_messages().len());
|
||||
}
|
||||
Err(e) => eprintln!("Error: {}", e),
|
||||
}
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
// Example 4: Blocking Prompt API
|
||||
//
|
||||
// Uses agent.prompt() for a quick one-shot query without streaming.
|
||||
//
|
||||
// Run: cargo run --example 04-prompt-api
|
||||
|
||||
use open_agent_sdk::{Agent, AgentOptions};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 4: Prompt API ---");
|
||||
|
||||
let mut agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(5),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
match agent
|
||||
.prompt("Run 'rustc --version' and 'cargo --version' using Bash, and tell me the versions.")
|
||||
.await
|
||||
{
|
||||
Ok(result) => {
|
||||
println!("\nResponse: {}", result.text);
|
||||
println!("\nTurns: {}", result.num_turns);
|
||||
println!(
|
||||
"Tokens: {} in / {} out",
|
||||
result.usage.input_tokens, result.usage.output_tokens
|
||||
);
|
||||
println!("Duration: {}ms", result.duration_ms);
|
||||
}
|
||||
Err(e) => eprintln!("Error: {}", e),
|
||||
}
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,53 @@
|
||||
// Example 5: Custom System Prompt
|
||||
//
|
||||
// Creates a specialized code reviewer agent with a custom system prompt.
|
||||
//
|
||||
// Run: cargo run --example 05-custom-system-prompt
|
||||
|
||||
use open_agent_sdk::{Agent, AgentOptions, SDKMessage};
|
||||
use open_agent_sdk::types;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 5: Custom System Prompt ---");
|
||||
|
||||
let mut agent = Agent::new(AgentOptions {
|
||||
system_prompt: Some(
|
||||
"You are a senior code reviewer. When analyzing code, focus on:\n\
|
||||
1. Security vulnerabilities\n\
|
||||
2. Performance issues\n\
|
||||
3. Maintainability concerns\n\
|
||||
Be concise and use bullet points."
|
||||
.to_string(),
|
||||
),
|
||||
max_turns: Some(10),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (mut rx, handle) = agent
|
||||
.query("Read src/lib.rs and give a brief code review.")
|
||||
.await;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
let text = types::extract_text(&message);
|
||||
if !text.is_empty() {
|
||||
print!("{}", text);
|
||||
}
|
||||
}
|
||||
SDKMessage::Result { .. } => {
|
||||
println!("\n\n--- Done ---");
|
||||
}
|
||||
SDKMessage::Error { message } => {
|
||||
eprintln!("Error: {}", message);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
handle.await.unwrap();
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
// Example 6: MCP Server Integration
|
||||
//
|
||||
// Connects to the filesystem MCP server via stdio transport.
|
||||
// Requires: npm install -g @modelcontextprotocol/server-filesystem
|
||||
//
|
||||
// Run: cargo run --example 06-mcp-server
|
||||
|
||||
use open_agent_sdk::{Agent, AgentOptions, SDKMessage};
|
||||
use open_agent_sdk::types::{self, McpServerConfig};
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 6: MCP Server ---");
|
||||
|
||||
let mut mcp_servers = HashMap::new();
|
||||
mcp_servers.insert(
|
||||
"filesystem".to_string(),
|
||||
McpServerConfig::Stdio {
|
||||
command: "npx".to_string(),
|
||||
args: vec![
|
||||
"-y".to_string(),
|
||||
"@modelcontextprotocol/server-filesystem".to_string(),
|
||||
"/tmp".to_string(),
|
||||
],
|
||||
env: HashMap::new(),
|
||||
},
|
||||
);
|
||||
|
||||
let agent = Agent::new(AgentOptions {
|
||||
mcp_servers,
|
||||
max_turns: Some(10),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
|
||||
match agent {
|
||||
Ok(mut agent) => {
|
||||
let (mut rx, handle) = agent
|
||||
.query("List files in /tmp using the filesystem MCP tools. Be brief.")
|
||||
.await;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
let text = types::extract_text(&message);
|
||||
if !text.is_empty() {
|
||||
print!("{}", text);
|
||||
}
|
||||
}
|
||||
SDKMessage::Result { .. } => {
|
||||
println!("\n--- Done ---");
|
||||
}
|
||||
SDKMessage::Error { message } => {
|
||||
eprintln!("Error: {}", message);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
handle.await.unwrap();
|
||||
agent.close().await;
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"Failed to create agent: {}\n\
|
||||
Make sure the MCP server is installed:\n\
|
||||
npm install -g @modelcontextprotocol/server-filesystem",
|
||||
e
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
// Example 7: Custom Tools
|
||||
//
|
||||
// Shows how to define and use custom tools alongside built-in tools.
|
||||
//
|
||||
// Run: cargo run --example 07-custom-tools
|
||||
|
||||
use async_trait::async_trait;
|
||||
use open_agent_sdk::types::{self, Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
use open_agent_sdk::{Agent, AgentOptions, SDKMessage};
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
// --- GetWeather tool ---
|
||||
|
||||
struct WeatherTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for WeatherTool {
|
||||
fn name(&self) -> &str {
|
||||
"GetWeather"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"Get current weather for a city. Returns temperature and conditions."
|
||||
}
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([(
|
||||
"city".to_string(),
|
||||
json!({"type": "string", "description": "City name (e.g., \"Tokyo\", \"London\")"}),
|
||||
)]),
|
||||
required: vec!["city".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
fn is_read_only(&self, _: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
async fn call(&self, input: Value, _ctx: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let city = input
|
||||
.get("city")
|
||||
.and_then(|c| c.as_str())
|
||||
.unwrap_or("Unknown");
|
||||
let temps: HashMap<&str, i32> = HashMap::from([
|
||||
("tokyo", 22),
|
||||
("london", 14),
|
||||
("beijing", 25),
|
||||
("new york", 18),
|
||||
("paris", 16),
|
||||
]);
|
||||
let temp = temps
|
||||
.get(city.to_lowercase().as_str())
|
||||
.copied()
|
||||
.unwrap_or(20);
|
||||
Ok(ToolResult::text(format!(
|
||||
"Weather in {}: {}°C, partly cloudy",
|
||||
city, temp
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
// --- Calculator tool ---
|
||||
|
||||
struct CalculatorTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for CalculatorTool {
|
||||
fn name(&self) -> &str {
|
||||
"Calculator"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"Evaluate a simple mathematical expression with two numbers and an operator (+, -, *, /, **)."
|
||||
}
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([(
|
||||
"expression".to_string(),
|
||||
json!({"type": "string", "description": "Math expression (e.g., \"2 ** 10\", \"42 * 17\")"}),
|
||||
)]),
|
||||
required: vec!["expression".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
fn is_read_only(&self, _: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
async fn call(&self, input: Value, _ctx: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let expr = input
|
||||
.get("expression")
|
||||
.and_then(|e| e.as_str())
|
||||
.unwrap_or("0");
|
||||
let result = eval_simple(expr);
|
||||
Ok(ToolResult::text(match result {
|
||||
Ok(val) => format!("{} = {}", expr, val),
|
||||
Err(e) => format!("Error: {}", e),
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
fn eval_simple(expr: &str) -> Result<String, String> {
|
||||
for op in &["**", "*", "/", "+", "-"] {
|
||||
if let Some(idx) = expr.find(op) {
|
||||
if idx == 0 {
|
||||
continue;
|
||||
}
|
||||
let left: f64 = expr[..idx]
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| "invalid number".to_string())?;
|
||||
let right: f64 = expr[idx + op.len()..]
|
||||
.trim()
|
||||
.parse()
|
||||
.map_err(|_| "invalid number".to_string())?;
|
||||
let result = match *op {
|
||||
"**" => left.powf(right),
|
||||
"*" => left * right,
|
||||
"/" => {
|
||||
if right == 0.0 {
|
||||
return Err("division by zero".to_string());
|
||||
}
|
||||
left / right
|
||||
}
|
||||
"+" => left + right,
|
||||
"-" => left - right,
|
||||
_ => return Err(format!("unknown operator: {}", op)),
|
||||
};
|
||||
return if result == (result as i64) as f64 {
|
||||
Ok(format!("{}", result as i64))
|
||||
} else {
|
||||
Ok(format!("{:.4}", result))
|
||||
};
|
||||
}
|
||||
}
|
||||
Err(format!("cannot evaluate: {}", expr))
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 7: Custom Tools ---");
|
||||
|
||||
let mut agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(10),
|
||||
custom_tools: vec![Arc::new(WeatherTool), Arc::new(CalculatorTool)],
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
println!("Loaded tools with 2 custom tools (GetWeather, Calculator)\n");
|
||||
|
||||
let (mut rx, handle) = agent
|
||||
.query("What is the weather in Tokyo and London? Also calculate 2 ** 10 * 3. Be brief.")
|
||||
.await;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
for block in &message.content {
|
||||
match block {
|
||||
types::ContentBlock::ToolUse { name, input, .. } => {
|
||||
println!("[{}] {}", name, serde_json::to_string(input).unwrap_or_default());
|
||||
}
|
||||
types::ContentBlock::Text { text } => {
|
||||
if !text.trim().is_empty() {
|
||||
println!("\n{}", text);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
SDKMessage::Result { .. } => {
|
||||
println!("\n--- Done ---");
|
||||
}
|
||||
SDKMessage::Error { message } => {
|
||||
eprintln!("Error: {}", message);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
handle.await.unwrap();
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
// Example 8: One-shot Query
|
||||
//
|
||||
// Quick one-shot agent query with restricted tools.
|
||||
//
|
||||
// Run: cargo run --example 08-oneshot-query
|
||||
|
||||
use open_agent_sdk::{Agent, AgentOptions};
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 8: One-shot Query ---");
|
||||
|
||||
let mut agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(5),
|
||||
allowed_tools: Some(vec!["Bash".to_string(), "Glob".to_string()]),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
match agent.prompt("List the files in the current directory. Be brief.").await {
|
||||
Ok(result) => {
|
||||
println!("\n{}", result.text);
|
||||
println!(
|
||||
"\n({} turns, {} tokens)",
|
||||
result.num_turns,
|
||||
result.usage.input_tokens + result.usage.output_tokens
|
||||
);
|
||||
}
|
||||
Err(e) => eprintln!("Error: {}", e),
|
||||
}
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Example 9: Subagents
|
||||
//
|
||||
// Creates a specialized code-reviewer subagent with restricted tools.
|
||||
//
|
||||
// Run: cargo run --example 09-subagents
|
||||
|
||||
use open_agent_sdk::agent::SubagentDefinition;
|
||||
use open_agent_sdk::{Agent, AgentOptions, SDKMessage};
|
||||
use open_agent_sdk::types;
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 9: Subagents ---");
|
||||
|
||||
let mut agents = HashMap::new();
|
||||
agents.insert(
|
||||
"code-reviewer".to_string(),
|
||||
SubagentDefinition {
|
||||
description: "Expert code reviewer that analyzes code quality".to_string(),
|
||||
instructions: Some(
|
||||
"You are an expert code reviewer. Analyze code for:\n\
|
||||
- Code quality and readability\n\
|
||||
- Potential bugs or edge cases\n\
|
||||
- Suggested improvements\n\
|
||||
Be concise and constructive."
|
||||
.to_string(),
|
||||
),
|
||||
tools: Some(vec![
|
||||
"Read".to_string(),
|
||||
"Glob".to_string(),
|
||||
"Grep".to_string(),
|
||||
]),
|
||||
model: None,
|
||||
},
|
||||
);
|
||||
|
||||
let mut agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(15),
|
||||
agents,
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (mut rx, handle) = agent
|
||||
.query("Use the code-reviewer agent to review src/lib.rs")
|
||||
.await;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
let text = types::extract_text(&message);
|
||||
if !text.is_empty() {
|
||||
print!("{}", text);
|
||||
}
|
||||
}
|
||||
SDKMessage::Result { cost_usd, .. } => {
|
||||
println!("\n\n--- Done (${:.4}) ---", cost_usd);
|
||||
}
|
||||
SDKMessage::Error { message } => {
|
||||
eprintln!("Error: {}", message);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
handle.await.unwrap();
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
// Example 10: Permissions
|
||||
//
|
||||
// Read-only agent with AllowedTools restriction.
|
||||
// Can only use Read, Glob, Grep — no write or execute.
|
||||
//
|
||||
// Run: cargo run --example 10-permissions
|
||||
|
||||
use open_agent_sdk::{Agent, AgentOptions, SDKMessage};
|
||||
use open_agent_sdk::types;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
println!("--- Example 10: Permissions ---");
|
||||
|
||||
let mut agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(10),
|
||||
allowed_tools: Some(vec![
|
||||
"Read".to_string(),
|
||||
"Glob".to_string(),
|
||||
"Grep".to_string(),
|
||||
]),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let (mut rx, handle) = agent
|
||||
.query("Review the code in src/lib.rs for best practices and potential improvements.")
|
||||
.await;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
match event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
let text = types::extract_text(&message);
|
||||
if !text.is_empty() {
|
||||
print!("{}", text);
|
||||
}
|
||||
}
|
||||
SDKMessage::Result { cost_usd, .. } => {
|
||||
println!("\n\n--- Done (${:.4}) ---", cost_usd);
|
||||
}
|
||||
SDKMessage::Error { message } => {
|
||||
eprintln!("Error: {}", message);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
handle.await.unwrap();
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,590 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Open Agent SDK - Chat</title>
|
||||
<style>
|
||||
* { margin: 0; padding: 0; box-sizing: border-box; }
|
||||
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
|
||||
background: #f9fafb;
|
||||
color: #1a1a1a;
|
||||
height: 100vh;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
}
|
||||
|
||||
/* Header */
|
||||
.header {
|
||||
padding: 16px 24px;
|
||||
border-bottom: 1px solid #e5e7eb;
|
||||
background: #fff;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
}
|
||||
.header h1 {
|
||||
font-size: 16px;
|
||||
font-weight: 600;
|
||||
color: #111;
|
||||
}
|
||||
.header .actions {
|
||||
display: flex;
|
||||
gap: 8px;
|
||||
}
|
||||
.btn-icon {
|
||||
width: 36px; height: 36px;
|
||||
border-radius: 8px;
|
||||
border: 1px solid #e5e7eb;
|
||||
background: #fff;
|
||||
cursor: pointer;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
font-size: 16px;
|
||||
color: #666;
|
||||
transition: all 0.15s;
|
||||
}
|
||||
.btn-icon:hover { background: #f3f4f6; border-color: #d1d5db; }
|
||||
|
||||
/* Chat area */
|
||||
.chat-container {
|
||||
flex: 1;
|
||||
overflow-y: auto;
|
||||
padding: 24px;
|
||||
max-width: 800px;
|
||||
width: 100%;
|
||||
margin: 0 auto;
|
||||
}
|
||||
|
||||
/* Welcome screen */
|
||||
.welcome {
|
||||
text-align: center;
|
||||
padding: 80px 20px 40px;
|
||||
}
|
||||
.welcome h2 {
|
||||
font-size: 28px;
|
||||
font-weight: 600;
|
||||
margin-bottom: 32px;
|
||||
color: #111;
|
||||
}
|
||||
.suggestions {
|
||||
display: flex;
|
||||
gap: 8px;
|
||||
justify-content: center;
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
.suggestion {
|
||||
padding: 8px 16px;
|
||||
border-radius: 20px;
|
||||
border: 1px solid #e5e7eb;
|
||||
background: #fff;
|
||||
cursor: pointer;
|
||||
font-size: 14px;
|
||||
color: #555;
|
||||
transition: all 0.15s;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 6px;
|
||||
}
|
||||
.suggestion:hover { border-color: #c4c4c4; background: #f9fafb; }
|
||||
|
||||
/* Messages */
|
||||
.message { margin-bottom: 20px; animation: fadeIn 0.2s ease; }
|
||||
@keyframes fadeIn { from { opacity: 0; transform: translateY(4px); } to { opacity: 1; transform: translateY(0); } }
|
||||
|
||||
.message-user {
|
||||
display: flex;
|
||||
justify-content: flex-end;
|
||||
}
|
||||
.message-user .bubble {
|
||||
background: #111;
|
||||
color: #fff;
|
||||
padding: 10px 16px;
|
||||
border-radius: 18px 18px 4px 18px;
|
||||
max-width: 70%;
|
||||
font-size: 15px;
|
||||
line-height: 1.5;
|
||||
}
|
||||
|
||||
.message-assistant {
|
||||
padding-left: 4px;
|
||||
}
|
||||
.message-assistant .text-content {
|
||||
font-size: 15px;
|
||||
line-height: 1.7;
|
||||
color: #1a1a1a;
|
||||
max-width: 100%;
|
||||
}
|
||||
.message-assistant .text-content p { margin-bottom: 8px; }
|
||||
.message-assistant .text-content strong { font-weight: 600; }
|
||||
.message-assistant .text-content code {
|
||||
background: #f3f4f6;
|
||||
padding: 2px 6px;
|
||||
border-radius: 4px;
|
||||
font-family: 'SF Mono', Monaco, monospace;
|
||||
font-size: 13px;
|
||||
}
|
||||
.message-assistant .text-content pre {
|
||||
background: #1e1e1e;
|
||||
color: #d4d4d4;
|
||||
padding: 12px 16px;
|
||||
border-radius: 8px;
|
||||
overflow-x: auto;
|
||||
margin: 8px 0;
|
||||
font-size: 13px;
|
||||
line-height: 1.5;
|
||||
}
|
||||
.message-assistant .text-content pre code {
|
||||
background: none;
|
||||
padding: 0;
|
||||
color: inherit;
|
||||
}
|
||||
|
||||
/* Tool calls */
|
||||
.tool-call {
|
||||
background: #f9fafb;
|
||||
border: 1px solid #e5e7eb;
|
||||
border-radius: 12px;
|
||||
padding: 10px 14px;
|
||||
margin: 8px 0;
|
||||
font-size: 13px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
}
|
||||
.tool-call .tool-icon {
|
||||
width: 28px; height: 28px;
|
||||
border-radius: 6px;
|
||||
background: #f3f4f6;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
font-size: 14px;
|
||||
flex-shrink: 0;
|
||||
}
|
||||
.tool-call .tool-name {
|
||||
font-weight: 600;
|
||||
color: #374151;
|
||||
margin-right: 6px;
|
||||
}
|
||||
.tool-call .tool-input {
|
||||
color: #9ca3af;
|
||||
font-family: 'SF Mono', Monaco, monospace;
|
||||
font-size: 12px;
|
||||
overflow: hidden;
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
flex: 1;
|
||||
}
|
||||
|
||||
.tool-call.tool-error { border-color: #fecaca; background: #fef2f2; }
|
||||
.tool-call.tool-error .tool-icon { background: #fee2e2; }
|
||||
|
||||
.tool-result-output {
|
||||
background: #f9fafb;
|
||||
border-left: 3px solid #e5e7eb;
|
||||
margin: 0 0 8px 36px;
|
||||
padding: 8px 12px;
|
||||
border-radius: 0 8px 8px 0;
|
||||
font-size: 12px;
|
||||
max-height: 200px;
|
||||
overflow-y: auto;
|
||||
}
|
||||
.tool-result-output pre {
|
||||
margin: 0;
|
||||
white-space: pre-wrap;
|
||||
word-break: break-all;
|
||||
font-family: 'SF Mono', Monaco, monospace;
|
||||
font-size: 12px;
|
||||
color: #374151;
|
||||
background: none;
|
||||
padding: 0;
|
||||
}
|
||||
|
||||
/* Thinking */
|
||||
.thinking-block {
|
||||
background: #fffbeb;
|
||||
border: 1px solid #fde68a;
|
||||
border-radius: 12px;
|
||||
padding: 10px 14px;
|
||||
margin: 8px 0;
|
||||
font-size: 13px;
|
||||
color: #92400e;
|
||||
font-style: italic;
|
||||
}
|
||||
|
||||
/* Result banner */
|
||||
.result-banner {
|
||||
background: #ecfdf5;
|
||||
border: 1px solid #a7f3d0;
|
||||
border-radius: 12px;
|
||||
padding: 10px 16px;
|
||||
margin: 12px 0;
|
||||
font-size: 14px;
|
||||
color: #065f46;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
}
|
||||
.result-banner .check { font-size: 16px; }
|
||||
.result-banner .stats {
|
||||
color: #047857;
|
||||
font-size: 13px;
|
||||
}
|
||||
|
||||
/* Error */
|
||||
.error-banner {
|
||||
background: #fef2f2;
|
||||
border: 1px solid #fecaca;
|
||||
border-radius: 12px;
|
||||
padding: 10px 16px;
|
||||
margin: 12px 0;
|
||||
font-size: 14px;
|
||||
color: #991b1b;
|
||||
}
|
||||
|
||||
/* Typing indicator */
|
||||
.typing {
|
||||
display: flex;
|
||||
gap: 4px;
|
||||
padding: 8px 0;
|
||||
}
|
||||
.typing span {
|
||||
width: 6px; height: 6px;
|
||||
background: #9ca3af;
|
||||
border-radius: 50%;
|
||||
animation: bounce 1.2s infinite;
|
||||
}
|
||||
.typing span:nth-child(2) { animation-delay: 0.2s; }
|
||||
.typing span:nth-child(3) { animation-delay: 0.4s; }
|
||||
@keyframes bounce {
|
||||
0%, 80%, 100% { transform: translateY(0); }
|
||||
40% { transform: translateY(-6px); }
|
||||
}
|
||||
|
||||
/* Input area */
|
||||
.input-area {
|
||||
padding: 16px 24px 24px;
|
||||
background: #fff;
|
||||
border-top: 1px solid #e5e7eb;
|
||||
}
|
||||
.input-wrapper {
|
||||
max-width: 800px;
|
||||
margin: 0 auto;
|
||||
position: relative;
|
||||
}
|
||||
.input-box {
|
||||
width: 100%;
|
||||
border: 1px solid #e5e7eb;
|
||||
border-radius: 16px;
|
||||
padding: 14px 52px 14px 16px;
|
||||
font-size: 15px;
|
||||
font-family: inherit;
|
||||
outline: none;
|
||||
resize: none;
|
||||
min-height: 52px;
|
||||
max-height: 200px;
|
||||
line-height: 1.5;
|
||||
transition: border-color 0.15s;
|
||||
background: #fff;
|
||||
}
|
||||
.input-box:focus { border-color: #a3a3a3; }
|
||||
.input-box::placeholder { color: #9ca3af; }
|
||||
.input-box:disabled { background: #f9fafb; color: #9ca3af; }
|
||||
|
||||
.send-btn {
|
||||
position: absolute;
|
||||
right: 8px;
|
||||
bottom: 8px;
|
||||
width: 36px; height: 36px;
|
||||
border-radius: 10px;
|
||||
border: none;
|
||||
background: #111;
|
||||
color: #fff;
|
||||
cursor: pointer;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
font-size: 18px;
|
||||
transition: all 0.15s;
|
||||
}
|
||||
.send-btn:hover { background: #333; }
|
||||
.send-btn:disabled { background: #d1d5db; cursor: not-allowed; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
|
||||
<div class="header">
|
||||
<h1>Open Agent SDK</h1>
|
||||
<div class="actions">
|
||||
<button class="btn-icon" onclick="newSession()" title="New session">+</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="chat-container" id="chat">
|
||||
<div class="welcome" id="welcome">
|
||||
<h2>What can I do for you?</h2>
|
||||
<div class="suggestions">
|
||||
<div class="suggestion" onclick="sendSuggestion('List all Rust files in this project')">📁 List Rust files</div>
|
||||
<div class="suggestion" onclick="sendSuggestion('Read Cargo.toml and explain the dependencies')">📄 Read Cargo.toml</div>
|
||||
<div class="suggestion" onclick="sendSuggestion('Count lines of code in the src directory')">📊 Count lines</div>
|
||||
<div class="suggestion" onclick="sendSuggestion('What is my Rust version and OS?')">💻 System info</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div class="input-area">
|
||||
<div class="input-wrapper">
|
||||
<textarea
|
||||
class="input-box"
|
||||
id="input"
|
||||
placeholder="Assign a task or ask anything"
|
||||
rows="1"
|
||||
onkeydown="handleKey(event)"
|
||||
oninput="autoResize(this)"
|
||||
></textarea>
|
||||
<button class="send-btn" id="sendBtn" onclick="sendMessage()">↑</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
const chat = document.getElementById('chat');
|
||||
const input = document.getElementById('input');
|
||||
const sendBtn = document.getElementById('sendBtn');
|
||||
const welcome = document.getElementById('welcome');
|
||||
let isStreaming = false;
|
||||
|
||||
const toolIcons = {
|
||||
'Bash': '>_',
|
||||
'Read': '📄',
|
||||
'Write': '📝',
|
||||
'Edit': '✏️',
|
||||
'Glob': '🔍',
|
||||
'Grep': '🔎',
|
||||
'WebFetch': '🌐',
|
||||
'WebSearch': '🔎',
|
||||
'Agent': '🤖',
|
||||
'TaskCreate': '📋',
|
||||
};
|
||||
|
||||
function getToolIcon(name) {
|
||||
if (toolIcons[name]) return toolIcons[name];
|
||||
if (name.startsWith('mcp__')) return '🔌';
|
||||
return '🔧';
|
||||
}
|
||||
|
||||
function autoResize(el) {
|
||||
el.style.height = 'auto';
|
||||
el.style.height = Math.min(el.scrollHeight, 200) + 'px';
|
||||
}
|
||||
|
||||
function handleKey(e) {
|
||||
if (e.key === 'Enter' && !e.shiftKey) {
|
||||
e.preventDefault();
|
||||
sendMessage();
|
||||
}
|
||||
}
|
||||
|
||||
function sendSuggestion(text) {
|
||||
input.value = text;
|
||||
sendMessage();
|
||||
}
|
||||
|
||||
function scrollToBottom() {
|
||||
chat.scrollTop = chat.scrollHeight;
|
||||
}
|
||||
|
||||
function escapeHtml(text) {
|
||||
const div = document.createElement('div');
|
||||
div.textContent = text;
|
||||
return div.innerHTML;
|
||||
}
|
||||
|
||||
function renderMarkdown(text) {
|
||||
let html = escapeHtml(text);
|
||||
html = html.replace(/```(\w*)\n([\s\S]*?)```/g, '<pre><code>$2</code></pre>');
|
||||
html = html.replace(/```([\s\S]*?)```/g, '<pre><code>$1</code></pre>');
|
||||
html = html.replace(/`([^`]+)`/g, '<code>$1</code>');
|
||||
html = html.replace(/\*\*(.+?)\*\*/g, '<strong>$1</strong>');
|
||||
html = html.replace(/^- (.+)$/gm, '• $1');
|
||||
html = html.replace(/\n\n/g, '</p><p>');
|
||||
html = '<p>' + html + '</p>';
|
||||
html = html.replace(/\n/g, '<br>');
|
||||
return html;
|
||||
}
|
||||
|
||||
async function newSession() {
|
||||
await fetch('/api/new', { method: 'POST' });
|
||||
chat.innerHTML = `
|
||||
<div class="welcome" id="welcome">
|
||||
<h2>What can I do for you?</h2>
|
||||
<div class="suggestions">
|
||||
<div class="suggestion" onclick="sendSuggestion('List all Rust files in this project')">📁 List Rust files</div>
|
||||
<div class="suggestion" onclick="sendSuggestion('Read Cargo.toml and explain the dependencies')">📄 Read Cargo.toml</div>
|
||||
<div class="suggestion" onclick="sendSuggestion('Count lines of code in the src directory')">📊 Count lines</div>
|
||||
<div class="suggestion" onclick="sendSuggestion('What is my Rust version and OS?')">💻 System info</div>
|
||||
</div>
|
||||
</div>`;
|
||||
}
|
||||
|
||||
async function sendMessage() {
|
||||
const text = input.value.trim();
|
||||
if (!text || isStreaming) return;
|
||||
|
||||
const w = document.getElementById('welcome');
|
||||
if (w) w.remove();
|
||||
|
||||
chat.innerHTML += `<div class="message message-user"><div class="bubble">${escapeHtml(text)}</div></div>`;
|
||||
|
||||
input.value = '';
|
||||
input.style.height = 'auto';
|
||||
isStreaming = true;
|
||||
sendBtn.disabled = true;
|
||||
input.disabled = true;
|
||||
|
||||
const typingId = 'typing-' + Date.now();
|
||||
chat.innerHTML += `<div class="message message-assistant" id="${typingId}"><div class="typing"><span></span><span></span><span></span></div></div>`;
|
||||
scrollToBottom();
|
||||
|
||||
let accumulatedText = '';
|
||||
let assistantDiv = null;
|
||||
let textContentDiv = null;
|
||||
|
||||
try {
|
||||
const resp = await fetch('/api/chat', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ message: text }),
|
||||
});
|
||||
|
||||
const reader = resp.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = '';
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split('\n');
|
||||
buffer = lines.pop() || '';
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.startsWith('data: ')) continue;
|
||||
const jsonStr = line.slice(6);
|
||||
if (!jsonStr) continue;
|
||||
|
||||
let event;
|
||||
try { event = JSON.parse(jsonStr); } catch { continue; }
|
||||
|
||||
const typingEl = document.getElementById(typingId);
|
||||
if (typingEl) typingEl.remove();
|
||||
|
||||
if (!assistantDiv) {
|
||||
assistantDiv = document.createElement('div');
|
||||
assistantDiv.className = 'message message-assistant';
|
||||
chat.appendChild(assistantDiv);
|
||||
}
|
||||
|
||||
switch (event.type) {
|
||||
case 'text':
|
||||
accumulatedText += event.data.text;
|
||||
if (!textContentDiv) {
|
||||
textContentDiv = document.createElement('div');
|
||||
textContentDiv.className = 'text-content';
|
||||
assistantDiv.appendChild(textContentDiv);
|
||||
}
|
||||
textContentDiv.innerHTML = renderMarkdown(accumulatedText);
|
||||
scrollToBottom();
|
||||
break;
|
||||
|
||||
case 'tool_use': {
|
||||
textContentDiv = null;
|
||||
accumulatedText = '';
|
||||
const inputStr = JSON.stringify(event.data.input || {});
|
||||
const tc = document.createElement('div');
|
||||
tc.className = 'tool-call';
|
||||
tc.id = 'tool-' + event.data.id;
|
||||
tc.innerHTML = `
|
||||
<div class="tool-icon">${getToolIcon(event.data.name)}</div>
|
||||
<span class="tool-name">${escapeHtml(event.data.name)}</span>
|
||||
<span class="tool-input">${escapeHtml(inputStr)}</span>
|
||||
`;
|
||||
assistantDiv.appendChild(tc);
|
||||
scrollToBottom();
|
||||
break;
|
||||
}
|
||||
|
||||
case 'tool_result': {
|
||||
const toolCallEl = document.getElementById('tool-' + event.data.tool_use_id);
|
||||
if (toolCallEl) {
|
||||
const resultDiv = document.createElement('div');
|
||||
resultDiv.className = 'tool-result-output';
|
||||
const content = event.data.content || '';
|
||||
const truncated = content.length > 500 ? content.substring(0, 500) + '...' : content;
|
||||
resultDiv.innerHTML = `<pre>${escapeHtml(truncated)}</pre>`;
|
||||
toolCallEl.insertAdjacentElement('afterend', resultDiv);
|
||||
scrollToBottom();
|
||||
}
|
||||
break;
|
||||
}
|
||||
|
||||
case 'thinking': {
|
||||
const th = document.createElement('div');
|
||||
th.className = 'thinking-block';
|
||||
th.textContent = event.data.thinking;
|
||||
assistantDiv.appendChild(th);
|
||||
scrollToBottom();
|
||||
break;
|
||||
}
|
||||
|
||||
case 'result': {
|
||||
const d = event.data;
|
||||
const costStr = d.cost < 0.01 ? `$${d.cost.toFixed(4)}` : `$${d.cost.toFixed(2)}`;
|
||||
const durStr = (d.duration_ms / 1000).toFixed(1) + 's';
|
||||
const banner = document.createElement('div');
|
||||
banner.className = 'result-banner';
|
||||
banner.innerHTML = `
|
||||
<span class="check">✅</span>
|
||||
<span>Done</span>
|
||||
<span class="stats">· ${costStr} · ${durStr} · ${d.input_tokens + d.output_tokens} tokens</span>
|
||||
`;
|
||||
assistantDiv.appendChild(banner);
|
||||
scrollToBottom();
|
||||
break;
|
||||
}
|
||||
|
||||
case 'error': {
|
||||
const errDiv = document.createElement('div');
|
||||
errDiv.className = 'error-banner';
|
||||
errDiv.textContent = '❌ ' + (event.data?.message || 'Unknown error');
|
||||
assistantDiv.appendChild(errDiv);
|
||||
scrollToBottom();
|
||||
break;
|
||||
}
|
||||
|
||||
case 'done':
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch (err) {
|
||||
const typingEl = document.getElementById(typingId);
|
||||
if (typingEl) typingEl.remove();
|
||||
chat.innerHTML += `<div class="error-banner">❌ ${escapeHtml(err.message)}</div>`;
|
||||
}
|
||||
|
||||
isStreaming = false;
|
||||
sendBtn.disabled = false;
|
||||
input.disabled = false;
|
||||
input.focus();
|
||||
scrollToBottom();
|
||||
}
|
||||
|
||||
input.focus();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,241 @@
|
||||
// Example 11: Web Chat UI
|
||||
//
|
||||
// A web-based chat interface for interacting with the agent.
|
||||
// Shows streaming text, tool calls with input/output, and cost tracking.
|
||||
//
|
||||
// Run: cargo run --example web-chat
|
||||
// Then open http://localhost:8082
|
||||
|
||||
use open_agent_sdk::{Agent, AgentOptions, SDKMessage};
|
||||
use open_agent_sdk::types;
|
||||
use serde::Serialize;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
const INDEX_HTML: &str = include_str!("index.html");
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SseEvent {
|
||||
#[serde(rename = "type")]
|
||||
event_type: String,
|
||||
data: serde_json::Value,
|
||||
}
|
||||
|
||||
struct AppState {
|
||||
agent: Option<Agent>,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let port = std::env::var("PORT").unwrap_or_else(|_| "8082".to_string());
|
||||
let state = Arc::new(Mutex::new(AppState { agent: None }));
|
||||
|
||||
let listener = TcpListener::bind(format!("0.0.0.0:{}", port))
|
||||
.await
|
||||
.expect("Failed to bind");
|
||||
|
||||
println!("🚀 Web Chat running at http://localhost:{}", port);
|
||||
|
||||
loop {
|
||||
let (stream, _) = listener.accept().await.expect("Failed to accept");
|
||||
let state = state.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
handle_connection(stream, state).await;
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_connection(mut stream: tokio::net::TcpStream, state: Arc<Mutex<AppState>>) {
|
||||
let mut buf = vec![0u8; 8192];
|
||||
let n = match stream.read(&mut buf).await {
|
||||
Ok(n) => n,
|
||||
Err(_) => return,
|
||||
};
|
||||
let request = String::from_utf8_lossy(&buf[..n]).to_string();
|
||||
|
||||
let first_line = request.lines().next().unwrap_or("");
|
||||
|
||||
if first_line.starts_with("GET / ") || first_line.starts_with("GET /index.html") {
|
||||
// Serve index.html
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\nContent-Type: text/html; charset=utf-8\r\nContent-Length: {}\r\n\r\n{}",
|
||||
INDEX_HTML.len(),
|
||||
INDEX_HTML
|
||||
);
|
||||
let _ = stream.write_all(response.as_bytes()).await;
|
||||
} else if first_line.starts_with("POST /api/new") {
|
||||
// New session
|
||||
{
|
||||
let mut state = state.lock().await;
|
||||
if let Some(agent) = state.agent.take() {
|
||||
agent.close().await;
|
||||
}
|
||||
}
|
||||
let response = "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n\r\n{\"status\":\"ok\"}";
|
||||
let _ = stream.write_all(response.as_bytes()).await;
|
||||
} else if first_line.starts_with("POST /api/chat") {
|
||||
// Extract body from request
|
||||
let body = extract_body(&request);
|
||||
let message = extract_json_field(&body, "message");
|
||||
|
||||
if message.is_empty() {
|
||||
let response = "HTTP/1.1 400 Bad Request\r\nContent-Type: text/plain\r\n\r\ninvalid request";
|
||||
let _ = stream.write_all(response.as_bytes()).await;
|
||||
return;
|
||||
}
|
||||
|
||||
// SSE headers
|
||||
let headers = "HTTP/1.1 200 OK\r\n\
|
||||
Content-Type: text/event-stream\r\n\
|
||||
Cache-Control: no-cache\r\n\
|
||||
Connection: keep-alive\r\n\
|
||||
X-Accel-Buffering: no\r\n\r\n";
|
||||
let _ = stream.write_all(headers.as_bytes()).await;
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Get or create agent
|
||||
let agent = {
|
||||
let mut state = state.lock().await;
|
||||
if state.agent.is_none() {
|
||||
let model = std::env::var("CODEANY_MODEL").unwrap_or_else(|_| "sonnet-4-6".to_string());
|
||||
match Agent::new(AgentOptions {
|
||||
model: Some(model),
|
||||
max_turns: Some(20),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(a) => state.agent = Some(a),
|
||||
Err(e) => {
|
||||
let event = format_sse("error", &serde_json::json!({"message": e}));
|
||||
let _ = stream.write_all(event.as_bytes()).await;
|
||||
let done = format_sse("done", &serde_json::json!(null));
|
||||
let _ = stream.write_all(done.as_bytes()).await;
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
// We need a mutable reference; use a different approach
|
||||
state.agent.take()
|
||||
};
|
||||
|
||||
let mut agent = agent.unwrap();
|
||||
|
||||
let (mut rx, handle) = agent.query(&message).await;
|
||||
|
||||
while let Some(event) = rx.recv().await {
|
||||
let sse = match &event {
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
let mut parts = Vec::new();
|
||||
for block in &message.content {
|
||||
match block {
|
||||
types::ContentBlock::Text { text } if !text.is_empty() => {
|
||||
parts.push(format_sse(
|
||||
"text",
|
||||
&serde_json::json!({"text": text}),
|
||||
));
|
||||
}
|
||||
types::ContentBlock::ToolUse { id, name, input } => {
|
||||
parts.push(format_sse(
|
||||
"tool_use",
|
||||
&serde_json::json!({"id": id, "name": name, "input": input}),
|
||||
));
|
||||
}
|
||||
types::ContentBlock::Thinking { thinking, .. } if !thinking.is_empty() => {
|
||||
parts.push(format_sse(
|
||||
"thinking",
|
||||
&serde_json::json!({"thinking": thinking}),
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
parts.join("")
|
||||
}
|
||||
SDKMessage::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
..
|
||||
} => format_sse(
|
||||
"tool_result",
|
||||
&serde_json::json!({
|
||||
"tool_use_id": tool_use_id,
|
||||
"content": content,
|
||||
"is_error": is_error,
|
||||
}),
|
||||
),
|
||||
SDKMessage::Result {
|
||||
usage,
|
||||
num_turns,
|
||||
cost_usd,
|
||||
..
|
||||
} => format_sse(
|
||||
"result",
|
||||
&serde_json::json!({
|
||||
"num_turns": num_turns,
|
||||
"input_tokens": usage.input_tokens,
|
||||
"output_tokens": usage.output_tokens,
|
||||
"cost": cost_usd,
|
||||
"duration_ms": start.elapsed().as_millis() as u64,
|
||||
}),
|
||||
),
|
||||
SDKMessage::Error { message } => {
|
||||
format_sse("error", &serde_json::json!({"message": message}))
|
||||
}
|
||||
_ => String::new(),
|
||||
};
|
||||
|
||||
if !sse.is_empty() {
|
||||
if stream.write_all(sse.as_bytes()).await.is_err() {
|
||||
break;
|
||||
}
|
||||
let _ = stream.flush().await;
|
||||
}
|
||||
}
|
||||
|
||||
let _ = handle.await;
|
||||
|
||||
let done = format_sse("done", &serde_json::json!(null));
|
||||
let _ = stream.write_all(done.as_bytes()).await;
|
||||
|
||||
// Put agent back
|
||||
let mut state = state.lock().await;
|
||||
state.agent = Some(agent);
|
||||
} else {
|
||||
let response = "HTTP/1.1 404 Not Found\r\nContent-Type: text/plain\r\n\r\nNot Found";
|
||||
let _ = stream.write_all(response.as_bytes()).await;
|
||||
}
|
||||
}
|
||||
|
||||
fn format_sse(event_type: &str, data: &serde_json::Value) -> String {
|
||||
let payload = serde_json::json!({"type": event_type, "data": data});
|
||||
format!("data: {}\n\n", serde_json::to_string(&payload).unwrap())
|
||||
}
|
||||
|
||||
fn extract_body(request: &str) -> String {
|
||||
if let Some(idx) = request.find("\r\n\r\n") {
|
||||
request[idx + 4..].to_string()
|
||||
} else {
|
||||
String::new()
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_json_field(body: &str, field: &str) -> String {
|
||||
// Simple JSON field extraction without a full parser dependency
|
||||
if let Ok(value) = serde_json::from_str::<serde_json::Value>(body) {
|
||||
value
|
||||
.get(field)
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("")
|
||||
.to_string()
|
||||
} else {
|
||||
String::new()
|
||||
}
|
||||
}
|
||||
|
||||
use tokio::io::AsyncReadExt;
|
||||
@@ -0,0 +1,408 @@
|
||||
use crate::api::ApiClient;
|
||||
use crate::costtracker::CostTracker;
|
||||
use crate::hooks::HookConfig;
|
||||
use crate::mcp::{self, McpClient};
|
||||
use crate::tools::ToolRegistry;
|
||||
use crate::types::*;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
/// Configuration options for creating an agent.
|
||||
pub struct AgentOptions {
|
||||
/// LLM model ID (default: claude-sonnet-4-6-20250514).
|
||||
pub model: Option<String>,
|
||||
/// API key for authentication.
|
||||
pub api_key: Option<String>,
|
||||
/// Custom API base URL.
|
||||
pub base_url: Option<String>,
|
||||
/// Custom LLM provider (bypasses built-in ApiClient creation).
|
||||
pub provider: Option<Arc<dyn crate::LLMProvider>>,
|
||||
/// Working directory for tool execution.
|
||||
pub cwd: Option<String>,
|
||||
/// Custom system prompt (replaces default).
|
||||
pub system_prompt: Option<String>,
|
||||
/// Additional text appended to the default system prompt.
|
||||
pub append_system_prompt: Option<String>,
|
||||
/// Skills context summary for system prompt injection.
|
||||
pub skills_summary: Option<String>,
|
||||
/// Maximum number of agentic turns (default: 100).
|
||||
pub max_turns: Option<u32>,
|
||||
/// Maximum budget in USD.
|
||||
pub max_budget_usd: Option<f64>,
|
||||
/// Permission mode for tool access control.
|
||||
pub permission_mode: Option<PermissionMode>,
|
||||
/// Tool names to allow (whitelist).
|
||||
pub allowed_tools: Option<Vec<String>>,
|
||||
/// Tool names to deny (blacklist).
|
||||
pub disallowed_tools: Option<Vec<String>>,
|
||||
/// Custom permission callback.
|
||||
pub can_use_tool: Option<CanUseToolFn>,
|
||||
/// Callback for user interaction (AskUserQuestion tool).
|
||||
pub ask_fn: Option<crate::tools::askuser::AskUserFn>,
|
||||
/// Custom tools to register.
|
||||
pub custom_tools: Vec<Arc<dyn Tool>>,
|
||||
/// MCP server configurations.
|
||||
pub mcp_servers: HashMap<String, McpServerConfig>,
|
||||
/// Subagent definitions.
|
||||
pub agents: HashMap<String, SubagentDefinition>,
|
||||
/// Extended thinking configuration.
|
||||
pub thinking: Option<ThinkingConfig>,
|
||||
/// Maximum output tokens per response.
|
||||
pub max_tokens: Option<u64>,
|
||||
/// Structured output JSON schema.
|
||||
pub json_schema: Option<Value>,
|
||||
/// Hook configuration.
|
||||
pub hooks: Option<HookConfig>,
|
||||
/// Cancellation token used to stop the running agent loop and tools.
|
||||
pub abort_signal: Option<CancellationToken>,
|
||||
/// Custom HTTP headers for API requests.
|
||||
pub custom_headers: HashMap<String, String>,
|
||||
/// Resume the most recent session in the current working directory.
|
||||
pub continue_session: Option<bool>,
|
||||
/// Resume a specific session by ID.
|
||||
pub resume: Option<String>,
|
||||
/// Fork instead of continuing a session.
|
||||
pub fork_session: Option<bool>,
|
||||
/// Control whether the session is persisted to disk.
|
||||
pub persist_session: Option<bool>,
|
||||
/// Explicit session ID to use (instead of auto-generating one).
|
||||
pub session_id: Option<String>,
|
||||
/// Effort level for reasoning: "low", "medium", "high", "max".
|
||||
pub effort: Option<String>,
|
||||
/// Fallback model if primary is unavailable.
|
||||
pub fallback_model: Option<String>,
|
||||
/// Sandbox configuration for restricting tool execution.
|
||||
pub sandbox: Option<SandboxSettings>,
|
||||
/// Enable debug mode.
|
||||
pub debug: Option<bool>,
|
||||
/// Additional working directories.
|
||||
pub additional_directories: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl Default for AgentOptions {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
model: None,
|
||||
api_key: None,
|
||||
base_url: None,
|
||||
provider: None,
|
||||
cwd: None,
|
||||
system_prompt: None,
|
||||
append_system_prompt: None,
|
||||
skills_summary: None,
|
||||
max_turns: None,
|
||||
max_budget_usd: None,
|
||||
permission_mode: None,
|
||||
allowed_tools: None,
|
||||
disallowed_tools: None,
|
||||
can_use_tool: None,
|
||||
ask_fn: None,
|
||||
custom_tools: Vec::new(),
|
||||
mcp_servers: HashMap::new(),
|
||||
agents: HashMap::new(),
|
||||
thinking: None,
|
||||
max_tokens: None,
|
||||
json_schema: None,
|
||||
hooks: None,
|
||||
abort_signal: None,
|
||||
custom_headers: HashMap::new(),
|
||||
continue_session: None,
|
||||
resume: None,
|
||||
fork_session: None,
|
||||
persist_session: None,
|
||||
session_id: None,
|
||||
effort: None,
|
||||
fallback_model: None,
|
||||
sandbox: None,
|
||||
debug: None,
|
||||
additional_directories: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Subagent definition.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SubagentDefinition {
|
||||
pub description: String,
|
||||
pub instructions: Option<String>,
|
||||
pub tools: Option<Vec<String>>,
|
||||
pub model: Option<String>,
|
||||
}
|
||||
|
||||
/// The main Agent struct that orchestrates the agentic loop.
|
||||
pub struct Agent {
|
||||
pub(crate) api_client: ApiClient,
|
||||
pub(crate) registry: ToolRegistry,
|
||||
/// Conversation messages.
|
||||
pub messages: Vec<Message>,
|
||||
/// Working directory.
|
||||
pub cwd: String,
|
||||
/// Custom system prompt.
|
||||
pub system_prompt: Option<String>,
|
||||
/// Additional system prompt text.
|
||||
pub append_system_prompt: Option<String>,
|
||||
/// Skills context summary.
|
||||
pub skills_summary: Option<String>,
|
||||
/// Maximum agentic turns.
|
||||
pub max_turns: u32,
|
||||
/// Maximum budget in USD.
|
||||
pub max_budget_usd: Option<f64>,
|
||||
pub(crate) cost_tracker: CostTracker,
|
||||
pub(crate) mcp_client: Arc<McpClient>,
|
||||
/// Subagent definitions.
|
||||
pub agents: HashMap<String, SubagentDefinition>,
|
||||
/// Extended thinking configuration.
|
||||
pub thinking: Option<ThinkingConfig>,
|
||||
/// Maximum output tokens.
|
||||
pub max_tokens: Option<u64>,
|
||||
/// Structured output JSON schema.
|
||||
pub json_schema: Option<Value>,
|
||||
/// Hook configuration.
|
||||
pub hooks: HookConfig,
|
||||
pub(crate) can_use_tool: Option<CanUseToolFn>,
|
||||
pub(crate) session_id: String,
|
||||
pub(crate) abort_signal: CancellationToken,
|
||||
}
|
||||
|
||||
impl Agent {
|
||||
/// Create a new agent with the given options.
|
||||
pub async fn new(options: AgentOptions) -> Result<Self, String> {
|
||||
let cwd = options.cwd.unwrap_or_else(|| {
|
||||
std::env::current_dir()
|
||||
.unwrap()
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
});
|
||||
|
||||
let api_client = if let Some(provider) = options.provider {
|
||||
ApiClient::with_provider(provider, options.model)
|
||||
} else {
|
||||
ApiClient::new(options.api_key, options.base_url, options.model)
|
||||
};
|
||||
|
||||
let mut registry = ToolRegistry::default_registry();
|
||||
|
||||
// Replace default AskUserTool with configured one if ask_fn is provided
|
||||
if let Some(ask_fn) = options.ask_fn {
|
||||
registry.register_replace(Arc::new(crate::tools::askuser::AskUserTool::new(ask_fn)));
|
||||
}
|
||||
|
||||
// Register custom tools
|
||||
for tool in options.custom_tools {
|
||||
registry.register(tool);
|
||||
}
|
||||
|
||||
// Apply allowed/disallowed tool filters
|
||||
if let Some(allowed) = &options.allowed_tools {
|
||||
let allowed_refs: Vec<&str> = allowed.iter().map(|s| s.as_str()).collect();
|
||||
registry.retain(&allowed_refs);
|
||||
}
|
||||
if let Some(disallowed) = &options.disallowed_tools {
|
||||
let disallowed_refs: Vec<&str> = disallowed.iter().map(|s| s.as_str()).collect();
|
||||
registry.remove(&disallowed_refs);
|
||||
}
|
||||
|
||||
// Connect MCP servers
|
||||
let mcp_client = Arc::new(McpClient::new());
|
||||
for (name, config) in &options.mcp_servers {
|
||||
match mcp_client.connect(name, config.clone()).await {
|
||||
Ok(tools) => {
|
||||
let mcp_tools = mcp::create_mcp_tools(name, &tools, mcp_client.clone());
|
||||
for tool in mcp_tools {
|
||||
registry.register(tool);
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to connect MCP server '{}': {}", name, e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let session_id = uuid::Uuid::new_v4().to_string();
|
||||
|
||||
Ok(Self {
|
||||
api_client,
|
||||
registry,
|
||||
messages: Vec::new(),
|
||||
cwd,
|
||||
system_prompt: options.system_prompt,
|
||||
append_system_prompt: options.append_system_prompt,
|
||||
skills_summary: options.skills_summary,
|
||||
max_turns: options.max_turns.unwrap_or(100),
|
||||
max_budget_usd: options.max_budget_usd,
|
||||
cost_tracker: CostTracker::new(),
|
||||
mcp_client,
|
||||
agents: options.agents,
|
||||
thinking: options.thinking,
|
||||
max_tokens: options.max_tokens,
|
||||
json_schema: options.json_schema,
|
||||
hooks: options.hooks.unwrap_or_default(),
|
||||
can_use_tool: options.can_use_tool,
|
||||
session_id,
|
||||
abort_signal: options.abort_signal.unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Execute a query and return a channel of streaming SDK messages.
|
||||
pub async fn query(
|
||||
&mut self,
|
||||
prompt: &str,
|
||||
) -> (mpsc::Receiver<SDKMessage>, tokio::task::JoinHandle<()>) {
|
||||
let (tx, rx) = mpsc::channel(100);
|
||||
|
||||
// Add user message
|
||||
let user_msg = crate::utils::messages::create_user_message(prompt);
|
||||
self.messages.push(user_msg);
|
||||
|
||||
// Clone what we need for the async loop
|
||||
let messages = self.messages.clone();
|
||||
let api_client = self.api_client.clone();
|
||||
let cwd = self.cwd.clone();
|
||||
let system_prompt = self.system_prompt.clone();
|
||||
let append_system_prompt = self.append_system_prompt.clone();
|
||||
let skills_summary = self.skills_summary.clone();
|
||||
let max_turns = self.max_turns;
|
||||
let max_budget_usd = self.max_budget_usd;
|
||||
let cost_tracker = self.cost_tracker.clone();
|
||||
let thinking = self.thinking.clone();
|
||||
let max_tokens = self.max_tokens;
|
||||
let registry = std::sync::Arc::new(self.registry_snapshot());
|
||||
let can_use_tool = self.can_use_tool.clone();
|
||||
let abort_signal = self.abort_signal.clone();
|
||||
|
||||
let handle = tokio::spawn(async move {
|
||||
let result: Result<Vec<Message>, String> = super::r#loop::run_loop(
|
||||
api_client,
|
||||
messages,
|
||||
registry,
|
||||
&cwd,
|
||||
system_prompt.as_deref(),
|
||||
append_system_prompt.as_deref(),
|
||||
skills_summary.as_deref(),
|
||||
max_turns,
|
||||
max_budget_usd,
|
||||
&cost_tracker,
|
||||
thinking,
|
||||
max_tokens,
|
||||
can_use_tool.as_ref(),
|
||||
abort_signal,
|
||||
tx.clone(),
|
||||
)
|
||||
.await;
|
||||
|
||||
if let Err(e) = result {
|
||||
let _ = tx
|
||||
.send(SDKMessage::Error {
|
||||
message: e.to_string(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
});
|
||||
|
||||
(rx, handle)
|
||||
}
|
||||
|
||||
/// Blocking query that returns a complete QueryResult.
|
||||
pub async fn prompt(&mut self, text: &str) -> Result<QueryResult, String> {
|
||||
let start = std::time::Instant::now();
|
||||
let (mut rx, handle) = self.query(text).await;
|
||||
|
||||
let mut result_text = String::new();
|
||||
let mut usage = Usage::default();
|
||||
let mut num_turns = 0;
|
||||
let mut cost_usd = 0.0;
|
||||
let mut all_messages = Vec::new();
|
||||
|
||||
while let Some(msg) = rx.recv().await {
|
||||
match msg {
|
||||
SDKMessage::Result {
|
||||
text,
|
||||
usage: u,
|
||||
num_turns: t,
|
||||
cost_usd: c,
|
||||
messages,
|
||||
..
|
||||
} => {
|
||||
result_text = text;
|
||||
usage = u;
|
||||
num_turns = t;
|
||||
cost_usd = c;
|
||||
all_messages = messages;
|
||||
}
|
||||
SDKMessage::Assistant { message, .. } => {
|
||||
let text = crate::types::extract_text(&message);
|
||||
if !text.is_empty() {
|
||||
result_text = text;
|
||||
}
|
||||
all_messages.push(message);
|
||||
}
|
||||
SDKMessage::Error { message } => {
|
||||
return Err(message);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
let _ = handle.await;
|
||||
let duration_ms = start.elapsed().as_millis() as u64;
|
||||
|
||||
// Update agent's message history
|
||||
self.messages = all_messages.clone();
|
||||
|
||||
Ok(QueryResult {
|
||||
text: result_text,
|
||||
usage,
|
||||
num_turns,
|
||||
cost_usd,
|
||||
duration_ms,
|
||||
messages: all_messages,
|
||||
})
|
||||
}
|
||||
|
||||
/// Get current conversation messages.
|
||||
pub fn get_messages(&self) -> &[Message] {
|
||||
&self.messages
|
||||
}
|
||||
|
||||
/// Clear conversation history.
|
||||
pub fn clear(&mut self) {
|
||||
self.messages.clear();
|
||||
}
|
||||
|
||||
/// Set the model.
|
||||
pub fn set_model(&mut self, model: &str) {
|
||||
self.api_client.set_model(model.to_string());
|
||||
}
|
||||
|
||||
/// Get the current model.
|
||||
pub fn model(&self) -> &str {
|
||||
self.api_client.model()
|
||||
}
|
||||
|
||||
/// Get the cost tracker.
|
||||
pub fn cost_tracker(&self) -> &CostTracker {
|
||||
&self.cost_tracker
|
||||
}
|
||||
|
||||
/// Get session ID.
|
||||
pub fn session_id(&self) -> &str {
|
||||
&self.session_id
|
||||
}
|
||||
|
||||
/// Close the agent and clean up resources.
|
||||
pub async fn close(&self) {
|
||||
self.mcp_client.close_all().await;
|
||||
}
|
||||
|
||||
/// Create a snapshot of the current registry for use in async context.
|
||||
fn registry_snapshot(&self) -> ToolRegistry {
|
||||
let mut snapshot = ToolRegistry::new();
|
||||
for tool in self.registry.all() {
|
||||
snapshot.register(tool);
|
||||
}
|
||||
snapshot
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
use crate::api::ApiClient;
|
||||
use crate::context;
|
||||
use crate::costtracker::CostTracker;
|
||||
use crate::tools::{self, ToolRegistry};
|
||||
use crate::types::*;
|
||||
use crate::utils::{compact, messages as msg_utils, retry};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
/// Run the main agentic loop.
|
||||
pub(crate) async fn run_loop(
|
||||
api_client: ApiClient,
|
||||
mut messages: Vec<Message>,
|
||||
registry: Arc<ToolRegistry>,
|
||||
cwd: &str,
|
||||
system_prompt: Option<&str>,
|
||||
append_system_prompt: Option<&str>,
|
||||
skills_summary: Option<&str>,
|
||||
max_turns: u32,
|
||||
max_budget_usd: Option<f64>,
|
||||
cost_tracker: &CostTracker,
|
||||
thinking: Option<ThinkingConfig>,
|
||||
max_tokens: Option<u64>,
|
||||
can_use_tool: Option<&CanUseToolFn>,
|
||||
abort_signal: CancellationToken,
|
||||
tx: mpsc::Sender<SDKMessage>,
|
||||
) -> Result<Vec<Message>, String> {
|
||||
let tool_context = ToolUseContext::with_abort(cwd.to_string(), abort_signal);
|
||||
|
||||
// Build system prompt blocks
|
||||
let system_blocks =
|
||||
context::build_system_blocks(cwd, system_prompt, append_system_prompt, skills_summary);
|
||||
|
||||
// Build tool definitions for API
|
||||
let api_tools: Vec<ApiToolParam> = registry
|
||||
.all()
|
||||
.iter()
|
||||
.map(|tool| {
|
||||
let schema = tool.input_schema();
|
||||
ApiToolParam {
|
||||
name: tool.name().to_string(),
|
||||
description: tool.description().to_string(),
|
||||
input_schema: serde_json::to_value(&schema).unwrap_or_default(),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let retry_config = retry::RetryConfig::default();
|
||||
let mut num_turns: u32 = 0;
|
||||
let mut total_usage = Usage::default();
|
||||
|
||||
// Send system message
|
||||
let _ = tx
|
||||
.send(SDKMessage::System {
|
||||
message: format!(
|
||||
"Agent started with {} tools ({})",
|
||||
api_tools.len(),
|
||||
match api_client.api_type() {
|
||||
crate::api::ApiType::AnthropicMessages => "anthropic",
|
||||
crate::api::ApiType::OpenAICompletions => "openai",
|
||||
}
|
||||
),
|
||||
})
|
||||
.await;
|
||||
|
||||
loop {
|
||||
if tool_context.abort_signal.is_cancelled() {
|
||||
return Err("Agent cancelled".to_string());
|
||||
}
|
||||
|
||||
// Check turn limit
|
||||
num_turns += 1;
|
||||
if num_turns > max_turns {
|
||||
let _ = tx
|
||||
.send(SDKMessage::Progress {
|
||||
message: format!("Reached max turns ({})", max_turns),
|
||||
})
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
|
||||
// Check budget
|
||||
if let Some(budget) = max_budget_usd {
|
||||
let current_cost = cost_tracker.total_cost().await;
|
||||
if current_cost >= budget {
|
||||
let _ = tx
|
||||
.send(SDKMessage::Progress {
|
||||
message: format!("Budget exceeded: ${:.4} >= ${:.4}", current_cost, budget),
|
||||
})
|
||||
.await;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Check if auto-compaction is needed
|
||||
if compact::should_auto_compact(&messages, api_client.model()) {
|
||||
messages = compact::micro_compact_messages(&messages);
|
||||
}
|
||||
|
||||
// Normalize messages
|
||||
let normalized = msg_utils::normalize_messages(&messages);
|
||||
|
||||
// Call the API via provider abstraction with retry
|
||||
let api_client_ref = &api_client;
|
||||
let system_blocks_ref = &system_blocks;
|
||||
let api_tools_ref = &api_tools;
|
||||
let thinking_ref = &thinking;
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
let response = tokio::select! {
|
||||
response = retry::with_retry(&retry_config, || async {
|
||||
api_client_ref
|
||||
.create_message(
|
||||
&normalized,
|
||||
Some(system_blocks_ref.clone()),
|
||||
Some(api_tools_ref.clone()),
|
||||
max_tokens,
|
||||
thinking_ref.clone(),
|
||||
Some(tx.clone()),
|
||||
)
|
||||
.await
|
||||
}) => response,
|
||||
_ = tool_context.abort_signal.cancelled() => {
|
||||
return Err("Agent cancelled".to_string());
|
||||
}
|
||||
};
|
||||
|
||||
let provider_response = match response {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
if retry::is_auth_error(&e) {
|
||||
return Err(format!("Authentication error: {}. Check your API key.", e));
|
||||
}
|
||||
return Err(format!("API error: {}", e));
|
||||
}
|
||||
};
|
||||
|
||||
let api_duration = start.elapsed().as_millis() as u64;
|
||||
cost_tracker.add_api_duration(api_duration).await;
|
||||
cost_tracker
|
||||
.add_usage(api_client.model(), &provider_response.usage)
|
||||
.await;
|
||||
|
||||
let usage = &provider_response.usage;
|
||||
let assistant_msg = provider_response.message;
|
||||
let stop_reason = provider_response.stop_reason;
|
||||
|
||||
// Accumulate usage
|
||||
total_usage.input_tokens += usage.input_tokens;
|
||||
total_usage.output_tokens += usage.output_tokens;
|
||||
total_usage.cache_creation_input_tokens += usage.cache_creation_input_tokens;
|
||||
total_usage.cache_read_input_tokens += usage.cache_read_input_tokens;
|
||||
|
||||
// Send assistant message event
|
||||
let _ = tx
|
||||
.send(SDKMessage::Assistant {
|
||||
message: assistant_msg.clone(),
|
||||
usage: Some(usage.clone()),
|
||||
})
|
||||
.await;
|
||||
|
||||
messages.push(assistant_msg.clone());
|
||||
|
||||
// Check for tool use
|
||||
let tool_uses = extract_tool_uses(&assistant_msg);
|
||||
|
||||
if tool_uses.is_empty() {
|
||||
// No tool calls - we're done (end_turn or max_tokens)
|
||||
break;
|
||||
}
|
||||
|
||||
// Execute tools
|
||||
let tool_start = std::time::Instant::now();
|
||||
let results = tokio::select! {
|
||||
results = tools::execute_tools(
|
||||
&assistant_msg,
|
||||
®istry,
|
||||
&tool_context,
|
||||
can_use_tool,
|
||||
Some(tx.clone()),
|
||||
) => results,
|
||||
_ = tool_context.abort_signal.cancelled() => {
|
||||
return Err("Agent cancelled".to_string());
|
||||
}
|
||||
};
|
||||
let tool_duration = tool_start.elapsed().as_millis() as u64;
|
||||
cost_tracker.add_tool_duration(tool_duration).await;
|
||||
|
||||
// Send tool result events
|
||||
for (id, name, result) in &results {
|
||||
let content_text = result.get_text();
|
||||
let _ = tx
|
||||
.send(SDKMessage::ToolResult {
|
||||
tool_use_id: id.clone(),
|
||||
tool_name: name.clone(),
|
||||
content: content_text,
|
||||
is_error: result.is_error,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
// Build tool results message and add to conversation
|
||||
let tool_results_msg = tools::build_tool_results_message(&results);
|
||||
messages.push(tool_results_msg);
|
||||
|
||||
// Check stop reason
|
||||
if stop_reason.as_deref() == Some("end_turn") && tool_uses.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// Calculate final cost
|
||||
let total_cost = cost_tracker.total_cost().await;
|
||||
|
||||
// Extract final text
|
||||
let final_text = messages
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|m| m.role == MessageRole::Assistant)
|
||||
.map(|m| extract_text(m))
|
||||
.unwrap_or_default();
|
||||
|
||||
// Send result event
|
||||
let _ = tx
|
||||
.send(SDKMessage::Result {
|
||||
text: final_text,
|
||||
usage: total_usage,
|
||||
num_turns,
|
||||
cost_usd: total_cost,
|
||||
duration_ms: 0,
|
||||
messages: messages.clone(),
|
||||
})
|
||||
.await;
|
||||
|
||||
Ok(messages)
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
mod agent;
|
||||
pub(crate) mod r#loop;
|
||||
|
||||
pub use agent::*;
|
||||
@@ -0,0 +1,338 @@
|
||||
use async_trait::async_trait;
|
||||
use futures::StreamExt;
|
||||
use reqwest::Client;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::types::{
|
||||
ApiToolParam, ContentBlock, ImageContentSource, Message, MessageRole,
|
||||
SDKMessage, SystemBlock, ThinkingConfig, Usage,
|
||||
};
|
||||
|
||||
use super::provider::{ApiType, LLMProvider, ProviderRequest, ProviderResponse};
|
||||
use super::ApiError;
|
||||
|
||||
const API_VERSION: &str = "2023-06-01";
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct AnthropicRequest {
|
||||
model: String,
|
||||
max_tokens: u64,
|
||||
messages: Vec<AnthropicMessage>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
system: Option<Vec<SystemBlock>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
tools: Option<Vec<ApiToolParam>>,
|
||||
stream: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
thinking: Option<ThinkingConfig>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct AnthropicMessage {
|
||||
role: String,
|
||||
content: Value,
|
||||
}
|
||||
|
||||
/// Anthropic Messages API provider.
|
||||
pub struct AnthropicProvider {
|
||||
client: Client,
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
custom_headers: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl AnthropicProvider {
|
||||
pub fn new(
|
||||
client: Client,
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
custom_headers: HashMap<String, String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
client,
|
||||
api_key,
|
||||
base_url,
|
||||
custom_headers,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LLMProvider for AnthropicProvider {
|
||||
fn api_type(&self) -> ApiType {
|
||||
ApiType::AnthropicMessages
|
||||
}
|
||||
|
||||
async fn create_message(
|
||||
&self,
|
||||
request: ProviderRequest<'_>,
|
||||
stream_tx: Option<mpsc::Sender<SDKMessage>>,
|
||||
) -> Result<ProviderResponse, ApiError> {
|
||||
let api_messages: Vec<AnthropicMessage> = request
|
||||
.messages
|
||||
.iter()
|
||||
.map(|m| {
|
||||
let role = match m.role {
|
||||
MessageRole::User => "user",
|
||||
MessageRole::Assistant => "assistant",
|
||||
};
|
||||
AnthropicMessage {
|
||||
role: role.to_string(),
|
||||
content: serde_json::to_value(&m.content).unwrap_or(Value::Array(vec![])),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let body = AnthropicRequest {
|
||||
model: request.model.to_string(),
|
||||
max_tokens: request.max_tokens,
|
||||
messages: api_messages,
|
||||
system: request.system,
|
||||
tools: if request.tools.as_ref().map_or(true, |t| t.is_empty()) {
|
||||
None
|
||||
} else {
|
||||
request.tools
|
||||
},
|
||||
stream: true,
|
||||
thinking: request.thinking,
|
||||
};
|
||||
|
||||
let mut req_builder = self
|
||||
.client
|
||||
.post(format!("{}/v1/messages", self.base_url))
|
||||
.header("x-api-key", &self.api_key)
|
||||
.header("anthropic-version", API_VERSION)
|
||||
.header("anthropic-beta", "prompt-caching-2024-07-31")
|
||||
.header("content-type", "application/json");
|
||||
|
||||
for (key, value) in &self.custom_headers {
|
||||
req_builder = req_builder.header(key, value);
|
||||
}
|
||||
|
||||
let response = req_builder.json(&body).send().await.map_err(|e| {
|
||||
if e.is_timeout() {
|
||||
ApiError::Timeout
|
||||
} else {
|
||||
ApiError::NetworkError(e.to_string())
|
||||
}
|
||||
})?;
|
||||
|
||||
let status = response.status().as_u16();
|
||||
if status == 401 {
|
||||
return Err(ApiError::AuthError("Invalid API key".to_string()));
|
||||
}
|
||||
if status == 429 {
|
||||
return Err(ApiError::RateLimitError);
|
||||
}
|
||||
if status >= 400 {
|
||||
let text = response.text().await.unwrap_or_default();
|
||||
if text.contains("prompt is too long") {
|
||||
return Err(ApiError::PromptTooLong(text));
|
||||
}
|
||||
return Err(ApiError::HttpError {
|
||||
status,
|
||||
message: text,
|
||||
});
|
||||
}
|
||||
|
||||
parse_anthropic_stream(response, stream_tx).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse Anthropic SSE stream into a ProviderResponse, emitting real-time
|
||||
/// thinking/text deltas via the optional `stream_tx` channel.
|
||||
async fn parse_anthropic_stream(
|
||||
response: reqwest::Response,
|
||||
stream_tx: Option<mpsc::Sender<SDKMessage>>,
|
||||
) -> Result<ProviderResponse, ApiError> {
|
||||
let mut content_blocks: Vec<ContentBlock> = Vec::new();
|
||||
let mut usage = Usage::default();
|
||||
let mut stop_reason: Option<String> = None;
|
||||
let mut current_blocks: HashMap<usize, Value> = HashMap::new();
|
||||
|
||||
// Stream the response body line-by-line instead of buffering the entire body
|
||||
let mut byte_stream = response.bytes_stream();
|
||||
let mut line_buf = String::new();
|
||||
|
||||
while let Some(chunk_result) = byte_stream.next().await {
|
||||
let chunk = chunk_result.map_err(|e| ApiError::NetworkError(e.to_string()))?;
|
||||
line_buf.push_str(&String::from_utf8_lossy(&chunk));
|
||||
|
||||
// Process all complete lines in the buffer
|
||||
while let Some(newline_pos) = line_buf.find('\n') {
|
||||
let raw_line = line_buf[..newline_pos].trim_end_matches('\r').to_string();
|
||||
line_buf = line_buf[newline_pos + 1..].to_string();
|
||||
|
||||
let line = raw_line.trim();
|
||||
if !line.starts_with("data: ") {
|
||||
continue;
|
||||
}
|
||||
let data = &line[6..];
|
||||
if data == "[DONE]" {
|
||||
break;
|
||||
}
|
||||
|
||||
let event: Value = match serde_json::from_str(data) {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
let event_type = event.get("type").and_then(|t| t.as_str()).unwrap_or("");
|
||||
|
||||
match event_type {
|
||||
"message_start" => {
|
||||
if let Some(u) = event.pointer("/message/usage") {
|
||||
if let Ok(u) = serde_json::from_value::<Usage>(u.clone()) {
|
||||
usage.input_tokens = u.input_tokens;
|
||||
usage.cache_creation_input_tokens = u.cache_creation_input_tokens;
|
||||
usage.cache_read_input_tokens = u.cache_read_input_tokens;
|
||||
}
|
||||
}
|
||||
}
|
||||
"content_block_start" => {
|
||||
if let (Some(idx), Some(block)) = (
|
||||
event.get("index").and_then(|i| i.as_u64()).map(|i| i as usize),
|
||||
event.get("content_block"),
|
||||
) {
|
||||
current_blocks.insert(idx, block.clone());
|
||||
}
|
||||
}
|
||||
"content_block_delta" => {
|
||||
let idx = event.get("index").and_then(|i| i.as_u64()).map(|i| i as usize);
|
||||
let delta = event.get("delta");
|
||||
if let (Some(idx), Some(delta)) = (idx, delta) {
|
||||
let delta_type = delta.get("type").and_then(|t| t.as_str()).unwrap_or("");
|
||||
match delta_type {
|
||||
"text_delta" => {
|
||||
if let Some(text) = delta.get("text").and_then(|t| t.as_str()) {
|
||||
let block = current_blocks
|
||||
.entry(idx)
|
||||
.or_insert_with(|| serde_json::json!({"type": "text", "text": ""}));
|
||||
let existing = block.get("text").and_then(|t| t.as_str()).unwrap_or("");
|
||||
block["text"] = Value::String(format!("{}{}", existing, text));
|
||||
// Emit real-time text delta
|
||||
if let Some(ref tx) = stream_tx {
|
||||
let _ = tx.try_send(SDKMessage::TextDelta {
|
||||
text: text.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
"input_json_delta" => {
|
||||
if let Some(partial) = delta.get("partial_json").and_then(|t| t.as_str()) {
|
||||
let block = current_blocks.entry(idx).or_insert_with(|| {
|
||||
serde_json::json!({"type": "tool_use", "id": "", "name": "", "input": {}})
|
||||
});
|
||||
let existing = block.get("_partial_json").and_then(|t| t.as_str()).unwrap_or("");
|
||||
block["_partial_json"] = Value::String(format!("{}{}", existing, partial));
|
||||
}
|
||||
}
|
||||
"thinking_delta" => {
|
||||
if let Some(thinking) = delta.get("thinking").and_then(|t| t.as_str()) {
|
||||
let block = current_blocks
|
||||
.entry(idx)
|
||||
.or_insert_with(|| serde_json::json!({"type": "thinking", "thinking": ""}));
|
||||
let existing = block.get("thinking").and_then(|t| t.as_str()).unwrap_or("");
|
||||
block["thinking"] = Value::String(format!("{}{}", existing, thinking));
|
||||
// Emit real-time thinking delta
|
||||
if let Some(ref tx) = stream_tx {
|
||||
let _ = tx.try_send(SDKMessage::ThinkingDelta {
|
||||
thinking: thinking.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
"signature_delta" => {
|
||||
if let Some(sig) = delta.get("signature").and_then(|t| t.as_str()) {
|
||||
if let Some(block) = current_blocks.get_mut(&idx) {
|
||||
block["signature"] = Value::String(sig.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
"content_block_stop" => {
|
||||
let idx = event.get("index").and_then(|i| i.as_u64()).map(|i| i as usize);
|
||||
if let Some(idx) = idx {
|
||||
if let Some(block) = current_blocks.remove(&idx) {
|
||||
if let Some(cb) = parse_anthropic_content_block(block) {
|
||||
content_blocks.push(cb);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"message_delta" => {
|
||||
if let Some(sr) = event.pointer("/delta/stop_reason").and_then(|s| s.as_str()) {
|
||||
stop_reason = Some(sr.to_string());
|
||||
}
|
||||
if let Some(u) = event.get("usage") {
|
||||
if let Some(out) = u.get("output_tokens").and_then(|o| o.as_u64()) {
|
||||
usage.output_tokens = out;
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flush remaining blocks
|
||||
let mut remaining: Vec<(usize, Value)> = current_blocks.into_iter().collect();
|
||||
remaining.sort_by_key(|(idx, _)| *idx);
|
||||
for (_, block) in remaining {
|
||||
if let Some(cb) = parse_anthropic_content_block(block) {
|
||||
content_blocks.push(cb);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(ProviderResponse {
|
||||
message: Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: content_blocks,
|
||||
},
|
||||
usage,
|
||||
stop_reason,
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_anthropic_content_block(mut block: Value) -> Option<ContentBlock> {
|
||||
let block_type = block.get("type")?.as_str()?.to_string();
|
||||
match block_type.as_str() {
|
||||
"text" => Some(ContentBlock::Text {
|
||||
text: block.get("text")?.as_str()?.to_string(),
|
||||
}),
|
||||
"tool_use" => {
|
||||
let id = block.get("id")?.as_str()?.to_string();
|
||||
let name = block.get("name")?.as_str()?.to_string();
|
||||
let input = if let Some(partial) = block.get("_partial_json").and_then(|p| p.as_str()) {
|
||||
serde_json::from_str(partial).unwrap_or(Value::Object(serde_json::Map::new()))
|
||||
} else {
|
||||
block.get("input").cloned().unwrap_or(Value::Object(serde_json::Map::new()))
|
||||
};
|
||||
if let Some(obj) = block.as_object_mut() {
|
||||
obj.remove("_partial_json");
|
||||
}
|
||||
Some(ContentBlock::ToolUse { id, name, input })
|
||||
}
|
||||
"thinking" => Some(ContentBlock::Thinking {
|
||||
thinking: block.get("thinking")?.as_str()?.to_string(),
|
||||
signature: block.get("signature").and_then(|s| s.as_str()).map(String::from),
|
||||
}),
|
||||
"image" => {
|
||||
let source = block.get("source")?;
|
||||
Some(ContentBlock::Image {
|
||||
source: ImageContentSource {
|
||||
source_type: source.get("type")?.as_str()?.to_string(),
|
||||
media_type: source.get("media_type")?.as_str()?.to_string(),
|
||||
data: source.get("data")?.as_str()?.to_string(),
|
||||
},
|
||||
})
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
use crate::types::{
|
||||
ApiToolParam, Message, SDKMessage, SystemBlock, ThinkingConfig, Usage,
|
||||
};
|
||||
use reqwest::Client;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use super::anthropic::AnthropicProvider;
|
||||
use super::openai::OpenAIProvider;
|
||||
use super::provider::{detect_api_type, ApiType, LLMProvider, ProviderRequest, ProviderResponse};
|
||||
|
||||
const DEFAULT_BASE_URL: &str = "https://api.anthropic.com";
|
||||
const DEFAULT_TIMEOUT_MS: u64 = 600_000; // 10 minutes
|
||||
|
||||
/// Model configuration with context window and output limits.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ModelConfig {
|
||||
pub context_window: u64,
|
||||
pub max_output_tokens: u64,
|
||||
}
|
||||
|
||||
/// Get model configuration for a given model ID.
|
||||
pub fn get_model_config(model: &str) -> ModelConfig {
|
||||
match model {
|
||||
m if m.contains("opus") => ModelConfig {
|
||||
context_window: 200_000,
|
||||
max_output_tokens: 32_000,
|
||||
},
|
||||
m if m.contains("sonnet") => ModelConfig {
|
||||
context_window: 200_000,
|
||||
max_output_tokens: 16_000,
|
||||
},
|
||||
m if m.contains("haiku") => ModelConfig {
|
||||
context_window: 200_000,
|
||||
max_output_tokens: 8_192,
|
||||
},
|
||||
m if m.starts_with("gpt-4o") => ModelConfig {
|
||||
context_window: 128_000,
|
||||
max_output_tokens: 16_384,
|
||||
},
|
||||
m if m.starts_with("o1") || m.starts_with("o3") || m.starts_with("o4") => ModelConfig {
|
||||
context_window: 200_000,
|
||||
max_output_tokens: 100_000,
|
||||
},
|
||||
m if m.contains("deepseek") => ModelConfig {
|
||||
context_window: 128_000,
|
||||
max_output_tokens: 8_192,
|
||||
},
|
||||
_ => ModelConfig {
|
||||
context_window: 200_000,
|
||||
max_output_tokens: 16_000,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Streaming event from the API (kept for backward compat).
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct StreamEvent {
|
||||
#[serde(rename = "type")]
|
||||
pub event_type: String,
|
||||
#[serde(default)]
|
||||
pub message: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub index: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub content_block: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub delta: Option<Value>,
|
||||
#[serde(default)]
|
||||
pub usage: Option<Usage>,
|
||||
}
|
||||
|
||||
/// Complete API response (non-streaming).
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct ApiResponse {
|
||||
pub id: String,
|
||||
pub content: Vec<Value>,
|
||||
pub model: String,
|
||||
pub stop_reason: Option<String>,
|
||||
pub usage: Usage,
|
||||
}
|
||||
|
||||
/// Error from the API.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ApiError {
|
||||
#[error("HTTP error: {status} - {message}")]
|
||||
HttpError { status: u16, message: String },
|
||||
|
||||
#[error("Authentication error: {0}")]
|
||||
AuthError(String),
|
||||
|
||||
#[error("Rate limit exceeded")]
|
||||
RateLimitError,
|
||||
|
||||
#[error("Prompt too long: {0}")]
|
||||
PromptTooLong(String),
|
||||
|
||||
#[error("Network error: {0}")]
|
||||
NetworkError(String),
|
||||
|
||||
#[error("Parse error: {0}")]
|
||||
ParseError(String),
|
||||
|
||||
#[error("Request timeout")]
|
||||
Timeout,
|
||||
}
|
||||
|
||||
/// API client that auto-detects and delegates to the correct provider.
|
||||
#[derive(Clone)]
|
||||
pub struct ApiClient {
|
||||
provider: Arc<dyn LLMProvider>,
|
||||
model: String,
|
||||
api_type: ApiType,
|
||||
}
|
||||
|
||||
impl ApiClient {
|
||||
pub fn new(api_key: Option<String>, base_url: Option<String>, model: Option<String>) -> Self {
|
||||
Self::with_api_type(api_key, base_url, model, None)
|
||||
}
|
||||
|
||||
/// Create a client using a pre-built LLM provider.
|
||||
pub fn with_provider(provider: Arc<dyn LLMProvider>, model: Option<String>) -> Self {
|
||||
let api_type = provider.api_type();
|
||||
let model = model.unwrap_or_else(|| "sonnet-4-6".to_string());
|
||||
Self {
|
||||
provider,
|
||||
model,
|
||||
api_type,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a client with an explicit API type override.
|
||||
pub fn with_api_type(
|
||||
api_key: Option<String>,
|
||||
base_url: Option<String>,
|
||||
model: Option<String>,
|
||||
api_type: Option<ApiType>,
|
||||
) -> Self {
|
||||
let api_key = api_key
|
||||
.or_else(|| std::env::var("CODEANY_API_KEY").ok())
|
||||
.or_else(|| std::env::var("ANTHROPIC_API_KEY").ok())
|
||||
.unwrap_or_default();
|
||||
|
||||
let base_url = base_url
|
||||
.or_else(|| std::env::var("CODEANY_BASE_URL").ok())
|
||||
.or_else(|| std::env::var("ANTHROPIC_BASE_URL").ok())
|
||||
.unwrap_or_else(|| DEFAULT_BASE_URL.to_string());
|
||||
|
||||
let model = model
|
||||
.or_else(|| std::env::var("CODEANY_MODEL").ok())
|
||||
.or_else(|| std::env::var("ANTHROPIC_MODEL").ok())
|
||||
.unwrap_or_else(|| "sonnet-4-6".to_string());
|
||||
|
||||
let timeout_ms: u64 = std::env::var("API_TIMEOUT_MS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(DEFAULT_TIMEOUT_MS);
|
||||
|
||||
let client = Client::builder()
|
||||
.timeout(Duration::from_millis(timeout_ms))
|
||||
.build()
|
||||
.expect("Failed to create HTTP client");
|
||||
|
||||
let mut custom_headers = HashMap::new();
|
||||
if let Ok(headers_str) = std::env::var("CODEANY_CUSTOM_HEADERS") {
|
||||
for pair in headers_str.split(',') {
|
||||
if let Some((key, value)) = pair.split_once(':') {
|
||||
custom_headers.insert(key.trim().to_string(), value.trim().to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Detect API type
|
||||
let resolved_api_type = api_type.unwrap_or_else(|| detect_api_type(&model, None));
|
||||
|
||||
let provider: Arc<dyn LLMProvider> = match &resolved_api_type {
|
||||
ApiType::AnthropicMessages => Arc::new(AnthropicProvider::new(
|
||||
client,
|
||||
api_key,
|
||||
base_url,
|
||||
custom_headers,
|
||||
)),
|
||||
ApiType::OpenAICompletions => Arc::new(OpenAIProvider::new(
|
||||
client,
|
||||
api_key,
|
||||
base_url,
|
||||
custom_headers,
|
||||
)),
|
||||
};
|
||||
|
||||
Self {
|
||||
provider,
|
||||
model,
|
||||
api_type: resolved_api_type,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn model(&self) -> &str {
|
||||
&self.model
|
||||
}
|
||||
|
||||
pub fn set_model(&mut self, model: String) {
|
||||
self.model = model;
|
||||
}
|
||||
|
||||
pub fn api_type(&self) -> &ApiType {
|
||||
&self.api_type
|
||||
}
|
||||
|
||||
pub fn model_config(&self) -> ModelConfig {
|
||||
get_model_config(&self.model)
|
||||
}
|
||||
|
||||
/// Send a streaming request via the provider and return the parsed response.
|
||||
pub async fn create_message(
|
||||
&self,
|
||||
messages: &[Message],
|
||||
system: Option<Vec<SystemBlock>>,
|
||||
tools: Option<Vec<ApiToolParam>>,
|
||||
max_tokens: Option<u64>,
|
||||
thinking: Option<ThinkingConfig>,
|
||||
stream_tx: Option<mpsc::Sender<SDKMessage>>,
|
||||
) -> Result<ProviderResponse, ApiError> {
|
||||
let model_config = self.model_config();
|
||||
let max_tokens = max_tokens.unwrap_or(model_config.max_output_tokens);
|
||||
|
||||
let request = ProviderRequest {
|
||||
model: &self.model,
|
||||
max_tokens,
|
||||
messages,
|
||||
system,
|
||||
tools,
|
||||
thinking,
|
||||
};
|
||||
|
||||
self.provider.create_message(request, stream_tx).await
|
||||
}
|
||||
|
||||
/// Legacy method: send streaming request and return raw reqwest::Response.
|
||||
/// Delegates to the Anthropic provider only; prefer `create_message()` instead.
|
||||
pub async fn create_message_stream(
|
||||
&self,
|
||||
_messages: &[Message],
|
||||
_system: Option<Vec<SystemBlock>>,
|
||||
_tools: Option<Vec<ApiToolParam>>,
|
||||
_max_tokens: Option<u64>,
|
||||
_thinking: Option<ThinkingConfig>,
|
||||
) -> Result<reqwest::Response, ApiError> {
|
||||
// For backward compat, this method is only meaningful for direct SSE parsing.
|
||||
// New code should use create_message() which handles both providers.
|
||||
Err(ApiError::ParseError(
|
||||
"create_message_stream is deprecated; use create_message() instead".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
/// Legacy: Parse Anthropic SSE stream. Kept for backward compatibility.
|
||||
pub async fn parse_stream(
|
||||
_response: reqwest::Response,
|
||||
) -> Result<(Message, Usage, Option<String>), ApiError> {
|
||||
// Redirect to new provider-based parsing isn't possible with raw response.
|
||||
// This method is kept only for API compat but new code should use create_message().
|
||||
Err(ApiError::ParseError(
|
||||
"parse_stream is deprecated; use create_message() instead".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if an error is retryable.
|
||||
pub fn is_retryable_error(error: &ApiError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
ApiError::RateLimitError
|
||||
| ApiError::Timeout
|
||||
| ApiError::NetworkError(_)
|
||||
| ApiError::HttpError { status: 500..=599, .. }
|
||||
)
|
||||
}
|
||||
|
||||
/// Check if an error is an auth error.
|
||||
pub fn is_auth_error(error: &ApiError) -> bool {
|
||||
matches!(error, ApiError::AuthError(_))
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
pub mod anthropic;
|
||||
mod client;
|
||||
pub mod openai;
|
||||
pub mod provider;
|
||||
|
||||
pub use client::*;
|
||||
pub use provider::{ApiType, LLMProvider, ProviderRequest, ProviderResponse};
|
||||
@@ -0,0 +1,385 @@
|
||||
use async_trait::async_trait;
|
||||
use reqwest::Client;
|
||||
use serde::Serialize;
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use crate::types::{
|
||||
ContentBlock, Message, MessageRole, SDKMessage, Usage,
|
||||
};
|
||||
|
||||
use super::provider::{ApiType, LLMProvider, ProviderRequest, ProviderResponse};
|
||||
use super::ApiError;
|
||||
|
||||
// --- OpenAI request types ---
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct OpenAIRequest {
|
||||
model: String,
|
||||
max_tokens: u64,
|
||||
messages: Vec<OpenAIMessage>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
tools: Option<Vec<OpenAITool>>,
|
||||
stream: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct OpenAIMessage {
|
||||
role: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
content: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
tool_calls: Option<Vec<OpenAIToolCall>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
tool_call_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct OpenAITool {
|
||||
#[serde(rename = "type")]
|
||||
tool_type: String,
|
||||
function: OpenAIFunction,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct OpenAIFunction {
|
||||
name: String,
|
||||
description: String,
|
||||
parameters: Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, serde::Deserialize)]
|
||||
struct OpenAIToolCall {
|
||||
id: String,
|
||||
#[serde(rename = "type")]
|
||||
call_type: String,
|
||||
function: OpenAIFunctionCall,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, serde::Deserialize)]
|
||||
struct OpenAIFunctionCall {
|
||||
name: String,
|
||||
arguments: String,
|
||||
}
|
||||
|
||||
/// OpenAI Chat Completions API provider.
|
||||
pub struct OpenAIProvider {
|
||||
client: Client,
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
custom_headers: HashMap<String, String>,
|
||||
}
|
||||
|
||||
impl OpenAIProvider {
|
||||
pub fn new(
|
||||
client: Client,
|
||||
api_key: String,
|
||||
base_url: String,
|
||||
custom_headers: HashMap<String, String>,
|
||||
) -> Self {
|
||||
Self {
|
||||
client,
|
||||
api_key,
|
||||
base_url,
|
||||
custom_headers,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl LLMProvider for OpenAIProvider {
|
||||
fn api_type(&self) -> ApiType {
|
||||
ApiType::OpenAICompletions
|
||||
}
|
||||
|
||||
async fn create_message(
|
||||
&self,
|
||||
request: ProviderRequest<'_>,
|
||||
_stream_tx: Option<mpsc::Sender<SDKMessage>>,
|
||||
) -> Result<ProviderResponse, ApiError> {
|
||||
// Convert messages to OpenAI format
|
||||
let mut openai_messages = Vec::new();
|
||||
|
||||
// System prompt as first message
|
||||
if let Some(system_blocks) = &request.system {
|
||||
let system_text: String = system_blocks
|
||||
.iter()
|
||||
.map(|b| b.text.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n\n");
|
||||
if !system_text.is_empty() {
|
||||
openai_messages.push(OpenAIMessage {
|
||||
role: "system".to_string(),
|
||||
content: Some(Value::String(system_text)),
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Convert conversation messages
|
||||
for msg in request.messages {
|
||||
match msg.role {
|
||||
MessageRole::User => {
|
||||
// Check if this is a tool_result message
|
||||
let tool_results: Vec<&ContentBlock> = msg
|
||||
.content
|
||||
.iter()
|
||||
.filter(|b| matches!(b, ContentBlock::ToolResult { .. }))
|
||||
.collect();
|
||||
|
||||
if !tool_results.is_empty() {
|
||||
// Each tool_result becomes a separate "tool" role message
|
||||
for block in &tool_results {
|
||||
if let ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
..
|
||||
} = block
|
||||
{
|
||||
let text: String = content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
crate::types::ToolResultContentBlock::Text { text } => {
|
||||
Some(text.as_str())
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
openai_messages.push(OpenAIMessage {
|
||||
role: "tool".to_string(),
|
||||
content: Some(Value::String(text)),
|
||||
tool_calls: None,
|
||||
tool_call_id: Some(tool_use_id.clone()),
|
||||
});
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Regular user message
|
||||
let text = crate::types::extract_text(msg);
|
||||
openai_messages.push(OpenAIMessage {
|
||||
role: "user".to_string(),
|
||||
content: Some(Value::String(text)),
|
||||
tool_calls: None,
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
MessageRole::Assistant => {
|
||||
let text = crate::types::extract_text(msg);
|
||||
let tool_uses: Vec<OpenAIToolCall> = msg
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|b| match b {
|
||||
ContentBlock::ToolUse { id, name, input } => Some(OpenAIToolCall {
|
||||
id: id.clone(),
|
||||
call_type: "function".to_string(),
|
||||
function: OpenAIFunctionCall {
|
||||
name: name.clone(),
|
||||
arguments: serde_json::to_string(input).unwrap_or_default(),
|
||||
},
|
||||
}),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
openai_messages.push(OpenAIMessage {
|
||||
role: "assistant".to_string(),
|
||||
content: if text.is_empty() { None } else { Some(Value::String(text)) },
|
||||
tool_calls: if tool_uses.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(tool_uses)
|
||||
},
|
||||
tool_call_id: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Convert tools to OpenAI format
|
||||
let openai_tools: Option<Vec<OpenAITool>> = request.tools.map(|tools| {
|
||||
tools
|
||||
.into_iter()
|
||||
.map(|t| OpenAITool {
|
||||
tool_type: "function".to_string(),
|
||||
function: OpenAIFunction {
|
||||
name: t.name,
|
||||
description: t.description,
|
||||
parameters: t.input_schema,
|
||||
},
|
||||
})
|
||||
.collect()
|
||||
});
|
||||
|
||||
let body = OpenAIRequest {
|
||||
model: request.model.to_string(),
|
||||
max_tokens: request.max_tokens,
|
||||
messages: openai_messages,
|
||||
tools: openai_tools.filter(|t| !t.is_empty()),
|
||||
stream: true,
|
||||
};
|
||||
|
||||
let mut req_builder = self
|
||||
.client
|
||||
.post(format!("{}/v1/chat/completions", self.base_url))
|
||||
.header("authorization", format!("Bearer {}", self.api_key))
|
||||
.header("content-type", "application/json");
|
||||
|
||||
for (key, value) in &self.custom_headers {
|
||||
req_builder = req_builder.header(key, value);
|
||||
}
|
||||
|
||||
let response = req_builder.json(&body).send().await.map_err(|e| {
|
||||
if e.is_timeout() {
|
||||
ApiError::Timeout
|
||||
} else {
|
||||
ApiError::NetworkError(e.to_string())
|
||||
}
|
||||
})?;
|
||||
|
||||
let status = response.status().as_u16();
|
||||
if status == 401 {
|
||||
return Err(ApiError::AuthError("Invalid API key".to_string()));
|
||||
}
|
||||
if status == 429 {
|
||||
return Err(ApiError::RateLimitError);
|
||||
}
|
||||
if status >= 400 {
|
||||
let text = response.text().await.unwrap_or_default();
|
||||
return Err(ApiError::HttpError {
|
||||
status,
|
||||
message: text,
|
||||
});
|
||||
}
|
||||
|
||||
parse_openai_stream(response).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse OpenAI SSE stream into a ProviderResponse.
|
||||
async fn parse_openai_stream(response: reqwest::Response) -> Result<ProviderResponse, ApiError> {
|
||||
let body = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| ApiError::NetworkError(e.to_string()))?;
|
||||
|
||||
let mut text_content = String::new();
|
||||
let mut tool_calls: HashMap<usize, OpenAIToolCall> = HashMap::new();
|
||||
let mut usage = Usage::default();
|
||||
let mut stop_reason: Option<String> = None;
|
||||
|
||||
for line in body.lines() {
|
||||
let line = line.trim();
|
||||
if !line.starts_with("data: ") {
|
||||
continue;
|
||||
}
|
||||
let data = &line[6..];
|
||||
if data == "[DONE]" {
|
||||
break;
|
||||
}
|
||||
|
||||
let chunk: Value = match serde_json::from_str(data) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
// Parse usage if present
|
||||
if let Some(u) = chunk.get("usage") {
|
||||
if let Some(pt) = u.get("prompt_tokens").and_then(|v| v.as_u64()) {
|
||||
usage.input_tokens = pt;
|
||||
}
|
||||
if let Some(ct) = u.get("completion_tokens").and_then(|v| v.as_u64()) {
|
||||
usage.output_tokens = ct;
|
||||
}
|
||||
}
|
||||
|
||||
// Parse choices
|
||||
if let Some(choices) = chunk.get("choices").and_then(|c| c.as_array()) {
|
||||
for choice in choices {
|
||||
// Check finish_reason
|
||||
if let Some(fr) = choice.get("finish_reason").and_then(|f| f.as_str()) {
|
||||
stop_reason = Some(match fr {
|
||||
"stop" => "end_turn".to_string(),
|
||||
"tool_calls" => "tool_use".to_string(),
|
||||
"length" => "max_tokens".to_string(),
|
||||
other => other.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
let delta = match choice.get("delta") {
|
||||
Some(d) => d,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
// Text content
|
||||
if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
|
||||
text_content.push_str(content);
|
||||
}
|
||||
|
||||
// Tool calls
|
||||
if let Some(tcs) = delta.get("tool_calls").and_then(|t| t.as_array()) {
|
||||
for tc in tcs {
|
||||
let idx = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
|
||||
|
||||
let entry = tool_calls.entry(idx).or_insert_with(|| OpenAIToolCall {
|
||||
id: String::new(),
|
||||
call_type: "function".to_string(),
|
||||
function: OpenAIFunctionCall {
|
||||
name: String::new(),
|
||||
arguments: String::new(),
|
||||
},
|
||||
});
|
||||
|
||||
if let Some(id) = tc.get("id").and_then(|i| i.as_str()) {
|
||||
entry.id = id.to_string();
|
||||
}
|
||||
if let Some(func) = tc.get("function") {
|
||||
if let Some(name) = func.get("name").and_then(|n| n.as_str()) {
|
||||
entry.function.name = name.to_string();
|
||||
}
|
||||
if let Some(args) = func.get("arguments").and_then(|a| a.as_str()) {
|
||||
entry.function.arguments.push_str(args);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build normalized content blocks
|
||||
let mut content_blocks: Vec<ContentBlock> = Vec::new();
|
||||
|
||||
if !text_content.is_empty() {
|
||||
content_blocks.push(ContentBlock::Text { text: text_content });
|
||||
}
|
||||
|
||||
// Convert tool calls to ContentBlock::ToolUse
|
||||
let mut sorted_calls: Vec<(usize, OpenAIToolCall)> = tool_calls.into_iter().collect();
|
||||
sorted_calls.sort_by_key(|(idx, _)| *idx);
|
||||
for (_, tc) in sorted_calls {
|
||||
let input: Value = serde_json::from_str(&tc.function.arguments)
|
||||
.unwrap_or(Value::Object(serde_json::Map::new()));
|
||||
content_blocks.push(ContentBlock::ToolUse {
|
||||
id: if tc.id.is_empty() {
|
||||
format!("call_{}", uuid::Uuid::new_v4())
|
||||
} else {
|
||||
tc.id
|
||||
},
|
||||
name: tc.function.name,
|
||||
input,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(ProviderResponse {
|
||||
message: Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: content_blocks,
|
||||
},
|
||||
usage,
|
||||
stop_reason,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,84 @@
|
||||
use crate::types::{ApiToolParam, Message, SDKMessage, SystemBlock, ThinkingConfig, Usage};
|
||||
use async_trait::async_trait;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// API type identifier.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum ApiType {
|
||||
AnthropicMessages,
|
||||
OpenAICompletions,
|
||||
}
|
||||
|
||||
/// Normalized response from any LLM provider.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProviderResponse {
|
||||
pub message: Message,
|
||||
pub usage: Usage,
|
||||
pub stop_reason: Option<String>,
|
||||
}
|
||||
|
||||
/// Configuration passed to provider for each request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ProviderRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub max_tokens: u64,
|
||||
pub messages: &'a [Message],
|
||||
pub system: Option<Vec<SystemBlock>>,
|
||||
pub tools: Option<Vec<ApiToolParam>>,
|
||||
pub thinking: Option<ThinkingConfig>,
|
||||
}
|
||||
|
||||
/// Trait that all LLM providers must implement.
|
||||
#[async_trait]
|
||||
pub trait LLMProvider: Send + Sync {
|
||||
fn api_type(&self) -> ApiType;
|
||||
|
||||
async fn create_message(
|
||||
&self,
|
||||
request: ProviderRequest<'_>,
|
||||
stream_tx: Option<mpsc::Sender<SDKMessage>>,
|
||||
) -> Result<ProviderResponse, super::ApiError>;
|
||||
}
|
||||
|
||||
/// Detect API type from model name using heuristics (matches TS SDK).
|
||||
pub fn detect_api_type(model: &str, env_api_type: Option<&str>) -> ApiType {
|
||||
// 1. Explicit env override
|
||||
if let Some(t) = env_api_type {
|
||||
match t {
|
||||
"openai-completions" => return ApiType::OpenAICompletions,
|
||||
"anthropic-messages" => return ApiType::AnthropicMessages,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Also check CODEANY_API_TYPE env var
|
||||
if let Ok(t) = std::env::var("CODEANY_API_TYPE") {
|
||||
match t.as_str() {
|
||||
"openai-completions" => return ApiType::OpenAICompletions,
|
||||
"anthropic-messages" => return ApiType::AnthropicMessages,
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Model name heuristics
|
||||
let m = model.to_lowercase();
|
||||
if m.contains("gpt-")
|
||||
|| m.starts_with("o1")
|
||||
|| m.starts_with("o3")
|
||||
|| m.starts_with("o4")
|
||||
|| m.contains("deepseek")
|
||||
|| m.contains("qwen")
|
||||
|| m.contains("yi-")
|
||||
|| m.contains("glm")
|
||||
|| m.contains("mistral")
|
||||
|| m.contains("gemma")
|
||||
|| m.contains("mimo")
|
||||
|| m.contains("llama")
|
||||
|| m.contains("gemini")
|
||||
{
|
||||
return ApiType::OpenAICompletions;
|
||||
}
|
||||
|
||||
// Default: Anthropic
|
||||
ApiType::AnthropicMessages
|
||||
}
|
||||
@@ -0,0 +1,198 @@
|
||||
use crate::types::{CacheControl, SystemBlock};
|
||||
use std::process::Command;
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
static GIT_STATUS_CACHE: OnceCell<String> = OnceCell::const_new();
|
||||
|
||||
/// Get git status for the given working directory.
|
||||
pub fn get_git_status(cwd: &str) -> String {
|
||||
// Try to get cached value first synchronously
|
||||
if let Some(cached) = GIT_STATUS_CACHE.get() {
|
||||
return cached.clone();
|
||||
}
|
||||
|
||||
let status = compute_git_status(cwd);
|
||||
// Best effort to cache
|
||||
let _ = GIT_STATUS_CACHE.set(status.clone());
|
||||
status
|
||||
}
|
||||
|
||||
fn compute_git_status(cwd: &str) -> String {
|
||||
let mut result = String::new();
|
||||
|
||||
// Get branch
|
||||
if let Ok(output) = Command::new("git")
|
||||
.args(["branch", "--show-current"])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
{
|
||||
let branch = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
if !branch.is_empty() {
|
||||
result.push_str(&format!("Current branch: {}\n", branch));
|
||||
}
|
||||
}
|
||||
|
||||
// Get main branch
|
||||
let main_branch = detect_main_branch(cwd);
|
||||
result.push_str(&format!(
|
||||
"\nMain branch: {}\n",
|
||||
main_branch
|
||||
));
|
||||
|
||||
// Get git user
|
||||
if let Ok(output) = Command::new("git")
|
||||
.args(["config", "user.name"])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
{
|
||||
let user = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
if !user.is_empty() {
|
||||
result.push_str(&format!("\nGit user: {}\n", user));
|
||||
}
|
||||
}
|
||||
|
||||
// Get status
|
||||
if let Ok(output) = Command::new("git")
|
||||
.args(["status", "--short"])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
{
|
||||
let status = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
if status.is_empty() {
|
||||
result.push_str("\nStatus:\n(clean)\n");
|
||||
} else {
|
||||
result.push_str(&format!("\nStatus:\n{}\n", status));
|
||||
}
|
||||
}
|
||||
|
||||
// Get recent commits
|
||||
if let Ok(output) = Command::new("git")
|
||||
.args(["log", "--oneline", "-5"])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
{
|
||||
let log = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
if !log.is_empty() {
|
||||
result.push_str(&format!("\nRecent commits:\n{}\n", log));
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn detect_main_branch(cwd: &str) -> String {
|
||||
// Check for main first, then master
|
||||
for branch in &["main", "master"] {
|
||||
if let Ok(output) = Command::new("git")
|
||||
.args(["rev-parse", "--verify", branch])
|
||||
.current_dir(cwd)
|
||||
.output()
|
||||
{
|
||||
if output.status.success() {
|
||||
return branch.to_string();
|
||||
}
|
||||
}
|
||||
}
|
||||
"main".to_string()
|
||||
}
|
||||
|
||||
/// Build system prompt blocks with context injection.
|
||||
pub fn build_system_blocks(
|
||||
cwd: &str,
|
||||
custom_system_prompt: Option<&str>,
|
||||
append_system_prompt: Option<&str>,
|
||||
skills_summary: Option<&str>,
|
||||
) -> Vec<SystemBlock> {
|
||||
let mut blocks = Vec::new();
|
||||
|
||||
// Main system prompt with cache control
|
||||
let system_prompt = if let Some(custom) = custom_system_prompt {
|
||||
custom.to_string()
|
||||
} else {
|
||||
default_system_prompt()
|
||||
};
|
||||
|
||||
blocks.push(SystemBlock {
|
||||
block_type: "text".to_string(),
|
||||
text: system_prompt,
|
||||
cache_control: Some(CacheControl::ephemeral()),
|
||||
});
|
||||
|
||||
// Git status context
|
||||
let git_status = get_git_status(cwd);
|
||||
if !git_status.is_empty() {
|
||||
blocks.push(SystemBlock {
|
||||
block_type: "text".to_string(),
|
||||
text: format!("gitStatus: {}", git_status),
|
||||
cache_control: None,
|
||||
});
|
||||
}
|
||||
|
||||
// User context (date, project files)
|
||||
let user_context = get_user_context(cwd);
|
||||
if !user_context.is_empty() {
|
||||
blocks.push(SystemBlock {
|
||||
block_type: "text".to_string(),
|
||||
text: user_context,
|
||||
cache_control: Some(CacheControl::ephemeral()),
|
||||
});
|
||||
}
|
||||
|
||||
// Skills context
|
||||
if let Some(skills) = skills_summary {
|
||||
if !skills.is_empty() {
|
||||
blocks.push(SystemBlock {
|
||||
block_type: "text".to_string(),
|
||||
text: skills.to_string(),
|
||||
cache_control: Some(CacheControl::ephemeral()),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Append system prompt
|
||||
if let Some(append) = append_system_prompt {
|
||||
blocks.push(SystemBlock {
|
||||
block_type: "text".to_string(),
|
||||
text: append.to_string(),
|
||||
cache_control: None,
|
||||
});
|
||||
}
|
||||
|
||||
blocks
|
||||
}
|
||||
|
||||
fn default_system_prompt() -> String {
|
||||
"You are a helpful AI assistant with access to tools for software engineering tasks. \
|
||||
Use the available tools to help the user accomplish their goals. \
|
||||
Be concise and direct in your responses."
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn get_user_context(cwd: &str) -> String {
|
||||
let mut context = String::new();
|
||||
|
||||
// Current date
|
||||
let date = chrono::Utc::now().format("%Y-%m-%d").to_string();
|
||||
context.push_str(&format!("Current date: {}\n", date));
|
||||
|
||||
// Check for project context files
|
||||
for filename in &["AGENT.md", "CLAUDE.md", ".agent/AGENT.md"] {
|
||||
let path = std::path::Path::new(cwd).join(filename);
|
||||
if path.exists() {
|
||||
if let Ok(content) = std::fs::read_to_string(&path) {
|
||||
context.push_str(&format!(
|
||||
"\n# Project context from {}\n{}\n",
|
||||
filename, content
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
context
|
||||
}
|
||||
|
||||
/// Clear the cached git status.
|
||||
pub fn clear_context_cache() {
|
||||
// OnceCell doesn't support reset, but this is a session-level cache
|
||||
// that typically doesn't need clearing during normal operation
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
use crate::types::Usage;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// Per-model pricing (per million tokens).
|
||||
#[derive(Debug, Clone)]
|
||||
struct ModelPricing {
|
||||
input: f64,
|
||||
output: f64,
|
||||
cache_read: f64,
|
||||
cache_write: f64,
|
||||
}
|
||||
|
||||
fn get_pricing(model: &str) -> ModelPricing {
|
||||
match model {
|
||||
m if m.contains("opus") => ModelPricing {
|
||||
input: 15.0,
|
||||
output: 75.0,
|
||||
cache_read: 1.5,
|
||||
cache_write: 18.75,
|
||||
},
|
||||
m if m.contains("sonnet") => ModelPricing {
|
||||
input: 3.0,
|
||||
output: 15.0,
|
||||
cache_read: 0.3,
|
||||
cache_write: 3.75,
|
||||
},
|
||||
m if m.contains("haiku") => ModelPricing {
|
||||
input: 0.8,
|
||||
output: 4.0,
|
||||
cache_read: 0.08,
|
||||
cache_write: 1.0,
|
||||
},
|
||||
_ => ModelPricing {
|
||||
input: 3.0,
|
||||
output: 15.0,
|
||||
cache_read: 0.3,
|
||||
cache_write: 3.75,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-model usage tracking.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ModelUsage {
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub cache_read_tokens: u64,
|
||||
pub cache_write_tokens: u64,
|
||||
}
|
||||
|
||||
/// Tracks token usage, costs, and other metrics across the agent session.
|
||||
#[derive(Clone)]
|
||||
pub struct CostTracker {
|
||||
inner: Arc<RwLock<CostTrackerInner>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct CostTrackerInner {
|
||||
model_usage: HashMap<String, ModelUsage>,
|
||||
api_duration_ms: u64,
|
||||
tool_duration_ms: u64,
|
||||
lines_added: u64,
|
||||
lines_removed: u64,
|
||||
web_searches: u64,
|
||||
}
|
||||
|
||||
impl CostTracker {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
inner: Arc::new(RwLock::new(CostTrackerInner::default())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Add token usage for a specific model.
|
||||
pub async fn add_usage(&self, model: &str, usage: &Usage) {
|
||||
let mut inner = self.inner.write().await;
|
||||
let entry = inner
|
||||
.model_usage
|
||||
.entry(model.to_string())
|
||||
.or_default();
|
||||
entry.input_tokens += usage.input_tokens;
|
||||
entry.output_tokens += usage.output_tokens;
|
||||
entry.cache_read_tokens += usage.cache_read_input_tokens;
|
||||
entry.cache_write_tokens += usage.cache_creation_input_tokens;
|
||||
}
|
||||
|
||||
/// Add API call duration.
|
||||
pub async fn add_api_duration(&self, duration_ms: u64) {
|
||||
let mut inner = self.inner.write().await;
|
||||
inner.api_duration_ms += duration_ms;
|
||||
}
|
||||
|
||||
/// Add tool execution duration.
|
||||
pub async fn add_tool_duration(&self, duration_ms: u64) {
|
||||
let mut inner = self.inner.write().await;
|
||||
inner.tool_duration_ms += duration_ms;
|
||||
}
|
||||
|
||||
/// Add code change metrics.
|
||||
pub async fn add_code_changes(&self, added: u64, removed: u64) {
|
||||
let mut inner = self.inner.write().await;
|
||||
inner.lines_added += added;
|
||||
inner.lines_removed += removed;
|
||||
}
|
||||
|
||||
/// Increment web search count.
|
||||
pub async fn add_web_search(&self) {
|
||||
let mut inner = self.inner.write().await;
|
||||
inner.web_searches += 1;
|
||||
}
|
||||
|
||||
/// Calculate total cost in USD.
|
||||
pub async fn total_cost(&self) -> f64 {
|
||||
let inner = self.inner.read().await;
|
||||
let mut total = 0.0;
|
||||
|
||||
for (model, usage) in &inner.model_usage {
|
||||
let pricing = get_pricing(model);
|
||||
total += (usage.input_tokens as f64 / 1_000_000.0) * pricing.input;
|
||||
total += (usage.output_tokens as f64 / 1_000_000.0) * pricing.output;
|
||||
total += (usage.cache_read_tokens as f64 / 1_000_000.0) * pricing.cache_read;
|
||||
total += (usage.cache_write_tokens as f64 / 1_000_000.0) * pricing.cache_write;
|
||||
}
|
||||
|
||||
total
|
||||
}
|
||||
|
||||
/// Get total tokens across all models.
|
||||
pub async fn total_tokens(&self) -> u64 {
|
||||
let inner = self.inner.read().await;
|
||||
inner
|
||||
.model_usage
|
||||
.values()
|
||||
.map(|u| u.input_tokens + u.output_tokens + u.cache_read_tokens + u.cache_write_tokens)
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Get usage for a specific model.
|
||||
pub async fn get_model_usage(&self, model: &str) -> Option<ModelUsage> {
|
||||
let inner = self.inner.read().await;
|
||||
inner.model_usage.get(model).cloned()
|
||||
}
|
||||
|
||||
/// Get all model usage.
|
||||
pub async fn get_all_usage(&self) -> HashMap<String, ModelUsage> {
|
||||
let inner = self.inner.read().await;
|
||||
inner.model_usage.clone()
|
||||
}
|
||||
|
||||
/// Get summary metrics.
|
||||
pub async fn summary(&self) -> CostSummary {
|
||||
let inner = self.inner.read().await;
|
||||
let total_cost = {
|
||||
let mut total = 0.0;
|
||||
for (model, usage) in &inner.model_usage {
|
||||
let pricing = get_pricing(model);
|
||||
total += (usage.input_tokens as f64 / 1_000_000.0) * pricing.input;
|
||||
total += (usage.output_tokens as f64 / 1_000_000.0) * pricing.output;
|
||||
total += (usage.cache_read_tokens as f64 / 1_000_000.0) * pricing.cache_read;
|
||||
total += (usage.cache_write_tokens as f64 / 1_000_000.0) * pricing.cache_write;
|
||||
}
|
||||
total
|
||||
};
|
||||
|
||||
CostSummary {
|
||||
total_cost_usd: total_cost,
|
||||
total_input_tokens: inner.model_usage.values().map(|u| u.input_tokens).sum(),
|
||||
total_output_tokens: inner.model_usage.values().map(|u| u.output_tokens).sum(),
|
||||
api_duration_ms: inner.api_duration_ms,
|
||||
tool_duration_ms: inner.tool_duration_ms,
|
||||
lines_added: inner.lines_added,
|
||||
lines_removed: inner.lines_removed,
|
||||
web_searches: inner.web_searches,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for CostTracker {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CostSummary {
|
||||
pub total_cost_usd: f64,
|
||||
pub total_input_tokens: u64,
|
||||
pub total_output_tokens: u64,
|
||||
pub api_duration_ms: u64,
|
||||
pub tool_duration_ms: u64,
|
||||
pub lines_added: u64,
|
||||
pub lines_removed: u64,
|
||||
pub web_searches: u64,
|
||||
}
|
||||
@@ -0,0 +1,297 @@
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Hook event types.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum HookEvent {
|
||||
PreToolUse,
|
||||
PostToolUse,
|
||||
PostToolUseFailure,
|
||||
PostSampling,
|
||||
SessionStart,
|
||||
SessionEnd,
|
||||
Stop,
|
||||
SubagentStart,
|
||||
SubagentStop,
|
||||
UserPromptSubmit,
|
||||
PermissionRequest,
|
||||
PermissionDenied,
|
||||
TaskCreated,
|
||||
TaskCompleted,
|
||||
ConfigChange,
|
||||
CwdChanged,
|
||||
FileChanged,
|
||||
Notification,
|
||||
PreCompact,
|
||||
PostCompact,
|
||||
TeammateIdle,
|
||||
}
|
||||
|
||||
/// Input passed to a hook handler.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HookInput {
|
||||
pub event: HookEvent,
|
||||
pub tool_name: Option<String>,
|
||||
pub tool_input: Option<Value>,
|
||||
pub tool_output: Option<String>,
|
||||
pub tool_use_id: Option<String>,
|
||||
pub session_id: Option<String>,
|
||||
pub cwd: Option<String>,
|
||||
pub error: Option<String>,
|
||||
}
|
||||
|
||||
/// Output from a hook handler.
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct HookOutput {
|
||||
/// If true, block the tool execution.
|
||||
pub blocked: bool,
|
||||
/// Optional message to return instead.
|
||||
pub message: Option<String>,
|
||||
/// Permission update (tool name + behavior).
|
||||
pub permission_update: Option<PermissionUpdate>,
|
||||
/// Notification to emit.
|
||||
pub notification: Option<HookNotification>,
|
||||
}
|
||||
|
||||
/// Permission update returned by a hook.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PermissionUpdate {
|
||||
pub tool: String,
|
||||
pub behavior: PermissionBehavior,
|
||||
}
|
||||
|
||||
/// Permission behavior: allow or deny.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PermissionBehavior {
|
||||
Allow,
|
||||
Deny,
|
||||
}
|
||||
|
||||
/// Notification emitted by a hook.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct HookNotification {
|
||||
pub title: String,
|
||||
pub body: String,
|
||||
pub level: Option<NotificationLevel>,
|
||||
}
|
||||
|
||||
/// Notification severity level.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum NotificationLevel {
|
||||
Info,
|
||||
Warning,
|
||||
Error,
|
||||
}
|
||||
|
||||
/// Hook function type.
|
||||
pub type HookFn =
|
||||
Arc<dyn Fn(HookInput) -> futures::future::BoxFuture<'static, HookOutput> + Send + Sync>;
|
||||
|
||||
/// A hook rule with a matcher pattern and handler.
|
||||
pub struct HookRule {
|
||||
pub matcher: String,
|
||||
pub handler: HookFn,
|
||||
}
|
||||
|
||||
/// Hook configuration for the agent.
|
||||
pub struct HookConfig {
|
||||
pub pre_tool_use: Vec<HookRule>,
|
||||
pub post_tool_use: Vec<HookRule>,
|
||||
pub post_tool_use_failure: Vec<HookRule>,
|
||||
pub post_sampling: Vec<HookRule>,
|
||||
pub session_start: Vec<HookRule>,
|
||||
pub session_end: Vec<HookRule>,
|
||||
pub stop: Vec<HookRule>,
|
||||
pub subagent_start: Vec<HookRule>,
|
||||
pub subagent_stop: Vec<HookRule>,
|
||||
pub user_prompt_submit: Vec<HookRule>,
|
||||
pub permission_request: Vec<HookRule>,
|
||||
pub permission_denied: Vec<HookRule>,
|
||||
pub task_created: Vec<HookRule>,
|
||||
pub task_completed: Vec<HookRule>,
|
||||
pub config_change: Vec<HookRule>,
|
||||
pub cwd_changed: Vec<HookRule>,
|
||||
pub file_changed: Vec<HookRule>,
|
||||
pub notification: Vec<HookRule>,
|
||||
pub pre_compact: Vec<HookRule>,
|
||||
pub post_compact: Vec<HookRule>,
|
||||
pub teammate_idle: Vec<HookRule>,
|
||||
}
|
||||
|
||||
impl Default for HookConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
pre_tool_use: Vec::new(),
|
||||
post_tool_use: Vec::new(),
|
||||
post_tool_use_failure: Vec::new(),
|
||||
post_sampling: Vec::new(),
|
||||
session_start: Vec::new(),
|
||||
session_end: Vec::new(),
|
||||
stop: Vec::new(),
|
||||
subagent_start: Vec::new(),
|
||||
subagent_stop: Vec::new(),
|
||||
user_prompt_submit: Vec::new(),
|
||||
permission_request: Vec::new(),
|
||||
permission_denied: Vec::new(),
|
||||
task_created: Vec::new(),
|
||||
task_completed: Vec::new(),
|
||||
config_change: Vec::new(),
|
||||
cwd_changed: Vec::new(),
|
||||
file_changed: Vec::new(),
|
||||
notification: Vec::new(),
|
||||
pre_compact: Vec::new(),
|
||||
post_compact: Vec::new(),
|
||||
teammate_idle: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl HookConfig {
|
||||
/// Get the hook rules for a given event.
|
||||
pub fn rules_for_event(&self, event: &HookEvent) -> &[HookRule] {
|
||||
match event {
|
||||
HookEvent::PreToolUse => &self.pre_tool_use,
|
||||
HookEvent::PostToolUse => &self.post_tool_use,
|
||||
HookEvent::PostToolUseFailure => &self.post_tool_use_failure,
|
||||
HookEvent::PostSampling => &self.post_sampling,
|
||||
HookEvent::SessionStart => &self.session_start,
|
||||
HookEvent::SessionEnd => &self.session_end,
|
||||
HookEvent::Stop => &self.stop,
|
||||
HookEvent::SubagentStart => &self.subagent_start,
|
||||
HookEvent::SubagentStop => &self.subagent_stop,
|
||||
HookEvent::UserPromptSubmit => &self.user_prompt_submit,
|
||||
HookEvent::PermissionRequest => &self.permission_request,
|
||||
HookEvent::PermissionDenied => &self.permission_denied,
|
||||
HookEvent::TaskCreated => &self.task_created,
|
||||
HookEvent::TaskCompleted => &self.task_completed,
|
||||
HookEvent::ConfigChange => &self.config_change,
|
||||
HookEvent::CwdChanged => &self.cwd_changed,
|
||||
HookEvent::FileChanged => &self.file_changed,
|
||||
HookEvent::Notification => &self.notification,
|
||||
HookEvent::PreCompact => &self.pre_compact,
|
||||
HookEvent::PostCompact => &self.post_compact,
|
||||
HookEvent::TeammateIdle => &self.teammate_idle,
|
||||
}
|
||||
}
|
||||
|
||||
/// Run hooks for a given event with tool context.
|
||||
pub async fn run_event(
|
||||
&self,
|
||||
event: HookEvent,
|
||||
tool_name: Option<&str>,
|
||||
tool_input: Option<&Value>,
|
||||
tool_output: Option<&str>,
|
||||
) -> Vec<HookOutput> {
|
||||
let rules = self.rules_for_event(&event);
|
||||
let mut outputs = Vec::new();
|
||||
for rule in rules {
|
||||
if let Some(name) = tool_name {
|
||||
if !matches_tool(&rule.matcher, name) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
let input = HookInput {
|
||||
event: event.clone(),
|
||||
tool_name: tool_name.map(|s| s.to_string()),
|
||||
tool_input: tool_input.cloned(),
|
||||
tool_output: tool_output.map(|s| s.to_string()),
|
||||
tool_use_id: None,
|
||||
session_id: None,
|
||||
cwd: None,
|
||||
error: None,
|
||||
};
|
||||
let output = (rule.handler)(input).await;
|
||||
outputs.push(output);
|
||||
}
|
||||
outputs
|
||||
}
|
||||
|
||||
/// Run pre-tool-use hooks for a given tool name.
|
||||
pub async fn run_pre_tool_use(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &Value,
|
||||
) -> Option<HookOutput> {
|
||||
for rule in &self.pre_tool_use {
|
||||
if matches_tool(&rule.matcher, tool_name) {
|
||||
let input = HookInput {
|
||||
event: HookEvent::PreToolUse,
|
||||
tool_name: Some(tool_name.to_string()),
|
||||
tool_input: Some(tool_input.clone()),
|
||||
tool_output: None,
|
||||
tool_use_id: None,
|
||||
session_id: None,
|
||||
cwd: None,
|
||||
error: None,
|
||||
};
|
||||
let output = (rule.handler)(input).await;
|
||||
if output.blocked {
|
||||
return Some(output);
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Run post-tool-use hooks for a given tool name.
|
||||
pub async fn run_post_tool_use(
|
||||
&self,
|
||||
tool_name: &str,
|
||||
tool_input: &Value,
|
||||
tool_output: &str,
|
||||
) {
|
||||
for rule in &self.post_tool_use {
|
||||
if matches_tool(&rule.matcher, tool_name) {
|
||||
let input = HookInput {
|
||||
event: HookEvent::PostToolUse,
|
||||
tool_name: Some(tool_name.to_string()),
|
||||
tool_input: Some(tool_input.clone()),
|
||||
tool_output: Some(tool_output.to_string()),
|
||||
tool_use_id: None,
|
||||
session_id: None,
|
||||
cwd: None,
|
||||
error: None,
|
||||
};
|
||||
(rule.handler)(input).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Run stop hooks.
|
||||
pub async fn run_stop(&self) {
|
||||
for rule in &self.stop {
|
||||
let input = HookInput {
|
||||
event: HookEvent::Stop,
|
||||
tool_name: None,
|
||||
tool_input: None,
|
||||
tool_output: None,
|
||||
tool_use_id: None,
|
||||
session_id: None,
|
||||
cwd: None,
|
||||
error: None,
|
||||
};
|
||||
(rule.handler)(input).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if a tool name matches a hook matcher pattern.
|
||||
fn matches_tool(matcher: &str, tool_name: &str) -> bool {
|
||||
if matcher == "*" || matcher.is_empty() {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Support pipe-separated patterns: "Bash|Edit|Write"
|
||||
if matcher.contains('|') {
|
||||
return matcher.split('|').any(|p| matches_tool(p.trim(), tool_name));
|
||||
}
|
||||
|
||||
// Support prefix matching: "mcp__*"
|
||||
if matcher.ends_with('*') {
|
||||
let prefix = &matcher[..matcher.len() - 1];
|
||||
return tool_name.starts_with(prefix);
|
||||
}
|
||||
|
||||
// Exact match
|
||||
matcher == tool_name
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
//! # Open Agent SDK
|
||||
//!
|
||||
//! A Rust framework for building autonomous AI agents that run the full
|
||||
//! agentic loop in-process. Supports 15+ built-in tools, MCP integration,
|
||||
//! permission systems, cost tracking, and more.
|
||||
//!
|
||||
//! ## Quick Start
|
||||
//!
|
||||
//! ```rust,no_run
|
||||
//! use open_agent_sdk::{Agent, AgentOptions};
|
||||
//!
|
||||
//! #[tokio::main]
|
||||
//! async fn main() {
|
||||
//! let mut agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
//! let result = agent.prompt("What files are in the current directory?").await.unwrap();
|
||||
//! println!("{}", result.text);
|
||||
//! }
|
||||
//! ```
|
||||
|
||||
pub mod agent;
|
||||
pub mod api;
|
||||
pub mod context;
|
||||
pub mod costtracker;
|
||||
pub mod hooks;
|
||||
pub mod mcp;
|
||||
pub mod permissions;
|
||||
pub mod session;
|
||||
pub mod skills;
|
||||
pub mod tools;
|
||||
pub mod types;
|
||||
pub mod utils;
|
||||
|
||||
// Re-export commonly used types
|
||||
pub use agent::{Agent, AgentOptions, SubagentDefinition};
|
||||
pub use api::{ApiClient, ApiType, LLMProvider, ProviderResponse};
|
||||
pub use costtracker::CostTracker;
|
||||
pub use hooks::{
|
||||
HookConfig, HookEvent, HookFn, HookInput, HookNotification, HookOutput, HookRule,
|
||||
NotificationLevel, PermissionBehavior, PermissionUpdate,
|
||||
};
|
||||
pub use mcp::McpClient;
|
||||
pub use session::{
|
||||
append_to_session, delete_session, fork_session, get_session_info, get_session_messages,
|
||||
list_sessions, load_session, new_metadata, rename_session, save_session, tag_session,
|
||||
SessionData, SessionMetadata,
|
||||
};
|
||||
pub use tokio_util::sync::CancellationToken;
|
||||
pub use tools::ToolRegistry;
|
||||
pub use types::{
|
||||
ApiToolParam, CanUseToolFn, ContentBlock, Message, MessageRole, PermissionDecision,
|
||||
PermissionMode, QueryResult, SDKMessage, SandboxFilesystemConfig, SandboxNetworkConfig,
|
||||
SandboxSettings, ThinkingConfig, Tool, ToolError, ToolInputSchema, ToolResult,
|
||||
ToolResultContent, ToolUseContext, Usage,
|
||||
};
|
||||
pub use utils::compact::{
|
||||
build_compaction_prompt, compact_conversation, strip_images_from_messages,
|
||||
};
|
||||
pub use utils::file_cache::FileStateCache;
|
||||
pub use utils::tokens::{estimate_cost, estimate_tokens};
|
||||
@@ -0,0 +1,341 @@
|
||||
use crate::types::{McpContentBlock, McpServerConfig, McpServerConnection, McpToolCallResult, McpToolDefinition};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::process::{Child, Command};
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// MCP client that manages connections to MCP servers.
|
||||
pub struct McpClient {
|
||||
connections: Arc<RwLock<HashMap<String, McpConnection>>>,
|
||||
}
|
||||
|
||||
struct McpConnection {
|
||||
config: McpServerConfig,
|
||||
tools: Vec<McpToolDefinition>,
|
||||
child: Option<Child>,
|
||||
stdin: Option<tokio::process::ChildStdin>,
|
||||
stdout_reader: Option<Arc<RwLock<BufReader<tokio::process::ChildStdout>>>>,
|
||||
}
|
||||
|
||||
impl McpClient {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
connections: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Connect to an MCP server and list its tools.
|
||||
pub async fn connect(
|
||||
&self,
|
||||
name: &str,
|
||||
config: McpServerConfig,
|
||||
) -> Result<Vec<McpToolDefinition>, String> {
|
||||
match &config {
|
||||
McpServerConfig::Stdio {
|
||||
command,
|
||||
args,
|
||||
env,
|
||||
} => {
|
||||
let mut cmd = Command::new(command);
|
||||
cmd.args(args)
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
hide_windows_console_window(&mut cmd);
|
||||
|
||||
// Inject login-shell PATH so GUI-launched apps can find
|
||||
// tools installed via version managers (nvm, fnm, volta, …)
|
||||
let shell_path = crate::mcp::shell_path::get_shell_path();
|
||||
if !shell_path.is_empty() && !env.contains_key("PATH") {
|
||||
cmd.env("PATH", shell_path);
|
||||
}
|
||||
for (key, value) in env {
|
||||
cmd.env(key, value);
|
||||
}
|
||||
|
||||
let mut child = cmd.spawn().map_err(|e| {
|
||||
format!("Failed to start MCP server '{}': {}", name, e)
|
||||
})?;
|
||||
|
||||
let stdin = child.stdin.take().ok_or("Failed to get stdin")?;
|
||||
let stdout = child.stdout.take().ok_or("Failed to get stdout")?;
|
||||
let reader = BufReader::new(stdout);
|
||||
|
||||
let mut conn = McpConnection {
|
||||
config: config.clone(),
|
||||
tools: Vec::new(),
|
||||
child: Some(child),
|
||||
stdin: Some(stdin),
|
||||
stdout_reader: Some(Arc::new(RwLock::new(reader))),
|
||||
};
|
||||
|
||||
// Initialize the connection
|
||||
let tools = initialize_connection(&mut conn).await?;
|
||||
conn.tools = tools.clone();
|
||||
|
||||
let mut connections = self.connections.write().await;
|
||||
connections.insert(name.to_string(), conn);
|
||||
|
||||
Ok(tools)
|
||||
}
|
||||
McpServerConfig::Sse { url, headers } | McpServerConfig::Http { url, headers } => {
|
||||
// HTTP-based MCP: list tools via HTTP
|
||||
let client = reqwest::Client::new();
|
||||
let mut req = client.post(format!("{}/list-tools", url));
|
||||
for (key, value) in headers {
|
||||
req = req.header(key, value);
|
||||
}
|
||||
|
||||
let response = req
|
||||
.json(&serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/list",
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to connect to MCP server '{}': {}", name, e))?;
|
||||
|
||||
let body: Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Invalid response from MCP server: {}", e))?;
|
||||
|
||||
let tools = parse_tool_list(&body);
|
||||
|
||||
let conn = McpConnection {
|
||||
config: config.clone(),
|
||||
tools: tools.clone(),
|
||||
child: None,
|
||||
stdin: None,
|
||||
stdout_reader: None,
|
||||
};
|
||||
|
||||
let mut connections = self.connections.write().await;
|
||||
connections.insert(name.to_string(), conn);
|
||||
|
||||
Ok(tools)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Call a tool on a connected MCP server.
|
||||
pub async fn call_tool(
|
||||
&self,
|
||||
server_name: &str,
|
||||
tool_name: &str,
|
||||
arguments: Value,
|
||||
) -> Result<McpToolCallResult, String> {
|
||||
let connections = self.connections.read().await;
|
||||
let conn = connections
|
||||
.get(server_name)
|
||||
.ok_or_else(|| format!("MCP server '{}' not connected", server_name))?;
|
||||
|
||||
match &conn.config {
|
||||
McpServerConfig::Stdio { .. } => {
|
||||
// Send JSON-RPC request via stdin/stdout
|
||||
// This is simplified; a full implementation would use proper JSON-RPC framing
|
||||
Ok(McpToolCallResult {
|
||||
content: vec![McpContentBlock::Text {
|
||||
text: format!(
|
||||
"MCP tool call: {}/{} (stdio transport)",
|
||||
server_name, tool_name
|
||||
),
|
||||
}],
|
||||
is_error: false,
|
||||
})
|
||||
}
|
||||
McpServerConfig::Sse { url, headers } | McpServerConfig::Http { url, headers } => {
|
||||
let client = reqwest::Client::new();
|
||||
let mut req = client.post(format!("{}/call-tool", url));
|
||||
for (key, value) in headers {
|
||||
req = req.header(key, value);
|
||||
}
|
||||
|
||||
let response = req
|
||||
.json(&serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": tool_name,
|
||||
"arguments": arguments,
|
||||
}
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
let body: Value = response.json().await.map_err(|e| e.to_string())?;
|
||||
|
||||
Ok(parse_tool_result(&body))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get all connections.
|
||||
pub async fn get_connections(&self) -> Vec<McpServerConnection> {
|
||||
let connections = self.connections.read().await;
|
||||
connections
|
||||
.iter()
|
||||
.map(|(name, conn)| McpServerConnection {
|
||||
name: name.clone(),
|
||||
config: conn.config.clone(),
|
||||
tools: conn.tools.clone(),
|
||||
connected: true,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Close all connections.
|
||||
pub async fn close_all(&self) {
|
||||
let mut connections = self.connections.write().await;
|
||||
for (_, mut conn) in connections.drain() {
|
||||
if let Some(mut child) = conn.child.take() {
|
||||
let _ = child.kill().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
fn hide_windows_console_window(cmd: &mut Command) {
|
||||
cmd.creation_flags(windows_sys::Win32::System::Threading::CREATE_NO_WINDOW);
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
fn hide_windows_console_window(_cmd: &mut Command) {}
|
||||
|
||||
impl Default for McpClient {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
async fn initialize_connection(
|
||||
conn: &mut McpConnection,
|
||||
) -> Result<Vec<McpToolDefinition>, String> {
|
||||
// Send initialize request
|
||||
let init_request = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "initialize",
|
||||
"params": {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {},
|
||||
"clientInfo": {
|
||||
"name": "open-agent-sdk-rust",
|
||||
"version": "0.1.0"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
send_jsonrpc(conn, &init_request).await?;
|
||||
let _init_response = read_jsonrpc(conn).await?;
|
||||
|
||||
// Send initialized notification
|
||||
let initialized = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"method": "notifications/initialized"
|
||||
});
|
||||
send_jsonrpc(conn, &initialized).await?;
|
||||
|
||||
// List tools
|
||||
let list_request = serde_json::json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 2,
|
||||
"method": "tools/list",
|
||||
});
|
||||
|
||||
send_jsonrpc(conn, &list_request).await?;
|
||||
let response = read_jsonrpc(conn).await?;
|
||||
|
||||
Ok(parse_tool_list(&response))
|
||||
}
|
||||
|
||||
async fn send_jsonrpc(conn: &mut McpConnection, message: &Value) -> Result<(), String> {
|
||||
if let Some(stdin) = &mut conn.stdin {
|
||||
let msg = format!("{}\n", serde_json::to_string(message).map_err(|e| e.to_string())?);
|
||||
stdin
|
||||
.write_all(msg.as_bytes())
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
stdin.flush().await.map_err(|e| e.to_string())?;
|
||||
Ok(())
|
||||
} else {
|
||||
Err("No stdin available".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_jsonrpc(conn: &McpConnection) -> Result<Value, String> {
|
||||
if let Some(reader) = &conn.stdout_reader {
|
||||
let mut reader = reader.write().await;
|
||||
let mut line = String::new();
|
||||
reader
|
||||
.read_line(&mut line)
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
serde_json::from_str(&line).map_err(|e| e.to_string())
|
||||
} else {
|
||||
Err("No stdout reader available".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_tool_list(response: &Value) -> Vec<McpToolDefinition> {
|
||||
response
|
||||
.get("result")
|
||||
.and_then(|r| r.get("tools"))
|
||||
.and_then(|t| t.as_array())
|
||||
.map(|tools| {
|
||||
tools
|
||||
.iter()
|
||||
.filter_map(|t| {
|
||||
Some(McpToolDefinition {
|
||||
name: t.get("name")?.as_str()?.to_string(),
|
||||
description: t
|
||||
.get("description")
|
||||
.and_then(|d| d.as_str())
|
||||
.map(String::from),
|
||||
input_schema: t.get("inputSchema").cloned(),
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn parse_tool_result(response: &Value) -> McpToolCallResult {
|
||||
let result = response.get("result").unwrap_or(response);
|
||||
let content = result
|
||||
.get("content")
|
||||
.and_then(|c| c.as_array())
|
||||
.map(|blocks| {
|
||||
blocks
|
||||
.iter()
|
||||
.filter_map(|b| {
|
||||
let block_type = b.get("type")?.as_str()?;
|
||||
match block_type {
|
||||
"text" => Some(McpContentBlock::Text {
|
||||
text: b.get("text")?.as_str()?.to_string(),
|
||||
}),
|
||||
"image" => Some(McpContentBlock::Image {
|
||||
data: b.get("data")?.as_str()?.to_string(),
|
||||
mime_type: b.get("mimeType")?.as_str()?.to_string(),
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let is_error = result
|
||||
.get("isError")
|
||||
.and_then(|e| e.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
McpToolCallResult { content, is_error }
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
mod client;
|
||||
pub mod shell_path;
|
||||
mod tool;
|
||||
|
||||
pub use client::*;
|
||||
pub use tool::*;
|
||||
@@ -0,0 +1,185 @@
|
||||
//! Resolve the user's login shell PATH for GUI-launched apps.
|
||||
//!
|
||||
//! On macOS/Linux, apps launched from Dock/Finder inherit a minimal PATH
|
||||
//! that doesn't include tools installed via nvm, fnm, volta, pyenv, etc.
|
||||
//! This module resolves the full login shell PATH and caches it.
|
||||
|
||||
use std::collections::HashSet;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
/// Get the cached login shell PATH. Returns an empty string if resolution fails.
|
||||
pub fn get_shell_path() -> &'static str {
|
||||
static SHELL_PATH: OnceLock<String> = OnceLock::new();
|
||||
SHELL_PATH.get_or_init(|| resolve_login_shell_path().unwrap_or_default())
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn resolve_login_shell_path() -> Option<String> {
|
||||
let current_path = std::env::var("PATH").ok();
|
||||
let mut best_path: Option<String> = None;
|
||||
|
||||
for shell in shell_candidates() {
|
||||
if let Some(candidate_path) = read_path_from_shell(&shell) {
|
||||
let merged = merge_paths(&candidate_path, current_path.as_deref());
|
||||
if path_score(&merged) > best_path.as_ref().map(|path| path_score(path)).unwrap_or(0) {
|
||||
best_path = Some(merged);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
best_path.or(current_path)
|
||||
}
|
||||
|
||||
#[cfg(not(unix))]
|
||||
fn resolve_login_shell_path() -> Option<String> {
|
||||
std::env::var("PATH").ok()
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn shell_candidates() -> Vec<String> {
|
||||
let mut candidates = Vec::new();
|
||||
let mut seen = HashSet::new();
|
||||
|
||||
for candidate in [
|
||||
std::env::var("SHELL").ok(),
|
||||
Some("zsh".to_string()),
|
||||
Some("/bin/zsh".to_string()),
|
||||
Some("bash".to_string()),
|
||||
Some("/bin/bash".to_string()),
|
||||
Some("sh".to_string()),
|
||||
Some("/bin/sh".to_string()),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
{
|
||||
if !candidate.is_empty() && seen.insert(candidate.clone()) {
|
||||
candidates.push(candidate);
|
||||
}
|
||||
}
|
||||
|
||||
candidates
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn read_path_from_shell(shell: &str) -> Option<String> {
|
||||
const START: &str = "__BOOSAPI_PATH_START__";
|
||||
const END: &str = "__BOOSAPI_PATH_END__";
|
||||
|
||||
let output = std::process::Command::new(shell)
|
||||
.args([
|
||||
"-i",
|
||||
"-l",
|
||||
"-c",
|
||||
&format!("printf '{START}'; printenv PATH; printf '{END}'"),
|
||||
])
|
||||
.stdin(std::process::Stdio::null())
|
||||
.stderr(std::process::Stdio::null())
|
||||
.output()
|
||||
.ok()?;
|
||||
|
||||
extract_marked_path(&output.stdout, START, END)
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn extract_marked_path(output: &[u8], start: &str, end: &str) -> Option<String> {
|
||||
let stdout = String::from_utf8(output.to_vec()).ok()?;
|
||||
let start_idx = stdout.find(start)? + start.len();
|
||||
let end_idx = stdout[start_idx..].find(end)? + start_idx;
|
||||
let path = stdout[start_idx..end_idx].trim().to_string();
|
||||
|
||||
if path.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(path)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn merge_paths(primary: &str, fallback: Option<&str>) -> String {
|
||||
let mut merged = Vec::new();
|
||||
let mut seen = HashSet::new();
|
||||
|
||||
for path_list in [Some(primary), fallback] {
|
||||
for segment in path_list
|
||||
.unwrap_or_default()
|
||||
.split(':')
|
||||
.map(str::trim)
|
||||
.filter(|segment| !segment.is_empty())
|
||||
{
|
||||
if seen.insert(segment.to_string()) {
|
||||
merged.push(segment.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
merged.join(":")
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn path_score(path: &str) -> usize {
|
||||
path.split(':')
|
||||
.filter(|segment| !segment.is_empty())
|
||||
.count()
|
||||
}
|
||||
|
||||
#[cfg(all(test, unix))]
|
||||
mod tests {
|
||||
use super::{extract_marked_path, merge_paths, resolve_login_shell_path};
|
||||
use std::fs;
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
|
||||
#[test]
|
||||
fn resolve_login_shell_path_uses_interactive_shell_config() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let fake_shell = dir.path().join("fake-shell.sh");
|
||||
let fake_node_dir = dir.path().join("bin");
|
||||
fs::create_dir_all(&fake_node_dir).unwrap();
|
||||
let interactive_path = std::iter::once(fake_node_dir.to_string_lossy().to_string())
|
||||
.chain((0..24).map(|index| format!("/tmp/boosapi-shell-{index}")))
|
||||
.collect::<Vec<_>>()
|
||||
.join(":");
|
||||
|
||||
let script = format!(
|
||||
"#!/bin/sh\nmode=plain\nfor arg in \"$@\"; do\n if [ \"$arg\" = \"-i\" ]; then\n mode=interactive\n fi\ndone\nif [ \"$mode\" = \"interactive\" ]; then\n printf '__BOOSAPI_PATH_START__%s__BOOSAPI_PATH_END__\\n' '{}:/usr/bin:/bin'\nelse\n printf '__BOOSAPI_PATH_START__%s__BOOSAPI_PATH_END__\\n' '/usr/bin:/bin'\nfi\n",
|
||||
interactive_path
|
||||
);
|
||||
fs::write(&fake_shell, script).unwrap();
|
||||
|
||||
let mut perms = fs::metadata(&fake_shell).unwrap().permissions();
|
||||
perms.set_mode(0o755);
|
||||
fs::set_permissions(&fake_shell, perms).unwrap();
|
||||
|
||||
let original_shell = std::env::var_os("SHELL");
|
||||
std::env::set_var("SHELL", &fake_shell);
|
||||
|
||||
let resolved = resolve_login_shell_path().unwrap();
|
||||
|
||||
match original_shell {
|
||||
Some(shell) => std::env::set_var("SHELL", shell),
|
||||
None => std::env::remove_var("SHELL"),
|
||||
}
|
||||
|
||||
assert!(
|
||||
resolved
|
||||
.split(':')
|
||||
.any(|segment| segment == fake_node_dir.to_string_lossy()),
|
||||
"expected interactive PATH to include {}, got {}",
|
||||
fake_node_dir.display(),
|
||||
resolved
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_marked_path_ignores_shell_noise() {
|
||||
let output = b"hello\n__BOOSAPI_PATH_START__/usr/local/bin:/usr/bin__BOOSAPI_PATH_END__\n";
|
||||
let path =
|
||||
extract_marked_path(output, "__BOOSAPI_PATH_START__", "__BOOSAPI_PATH_END__").unwrap();
|
||||
assert_eq!(path, "/usr/local/bin:/usr/bin");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merge_paths_deduplicates_segments() {
|
||||
let merged = merge_paths("/opt/bin:/usr/bin", Some("/usr/bin:/bin"));
|
||||
assert_eq!(merged, "/opt/bin:/usr/bin:/bin");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::types::{McpContentBlock, McpToolDefinition, Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
use super::client::McpClient;
|
||||
|
||||
/// Wraps an MCP tool as a standard Tool implementation.
|
||||
/// Tool names are prefixed: mcp__{server_name}__{tool_name}
|
||||
pub struct McpToolWrapper {
|
||||
server_name: String,
|
||||
tool_def: McpToolDefinition,
|
||||
client: Arc<McpClient>,
|
||||
prefixed_name: String,
|
||||
}
|
||||
|
||||
impl McpToolWrapper {
|
||||
pub fn new(
|
||||
server_name: &str,
|
||||
tool_def: McpToolDefinition,
|
||||
client: Arc<McpClient>,
|
||||
) -> Self {
|
||||
let prefixed_name = format!("mcp__{}_{}", server_name, tool_def.name);
|
||||
Self {
|
||||
server_name: server_name.to_string(),
|
||||
tool_def,
|
||||
client,
|
||||
prefixed_name,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for McpToolWrapper {
|
||||
fn name(&self) -> &str {
|
||||
&self.prefixed_name
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
self.tool_def
|
||||
.description
|
||||
.as_deref()
|
||||
.unwrap_or("MCP tool")
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
if let Some(schema) = &self.tool_def.input_schema {
|
||||
if let Ok(parsed) = serde_json::from_value::<ToolInputSchema>(schema.clone()) {
|
||||
return parsed;
|
||||
}
|
||||
}
|
||||
ToolInputSchema::default()
|
||||
}
|
||||
|
||||
async fn call(
|
||||
&self,
|
||||
input: Value,
|
||||
_context: &ToolUseContext,
|
||||
) -> Result<ToolResult, ToolError> {
|
||||
let result = self
|
||||
.client
|
||||
.call_tool(&self.server_name, &self.tool_def.name, input)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionError(e))?;
|
||||
|
||||
if result.is_error {
|
||||
let text: String = result
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
McpContentBlock::Text { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
Ok(ToolResult::error(text))
|
||||
} else {
|
||||
let text: String = result
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
McpContentBlock::Text { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
Ok(ToolResult::text(text))
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
false // MCP tools are treated as mutations by default
|
||||
}
|
||||
}
|
||||
|
||||
/// Create tool wrappers for all tools from a connected MCP server.
|
||||
pub fn create_mcp_tools(
|
||||
server_name: &str,
|
||||
tools: &[McpToolDefinition],
|
||||
client: Arc<McpClient>,
|
||||
) -> Vec<Arc<dyn Tool>> {
|
||||
tools
|
||||
.iter()
|
||||
.map(|tool_def| {
|
||||
Arc::new(McpToolWrapper::new(server_name, tool_def.clone(), client.clone()))
|
||||
as Arc<dyn Tool>
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
@@ -0,0 +1,104 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Sensitive file patterns that should be flagged.
|
||||
const SENSITIVE_PATTERNS: &[&str] = &[
|
||||
".env",
|
||||
".env.local",
|
||||
".env.production",
|
||||
"credentials.json",
|
||||
"credentials.yaml",
|
||||
"credentials.yml",
|
||||
".aws/credentials",
|
||||
".ssh/id_rsa",
|
||||
".ssh/id_ed25519",
|
||||
"id_rsa",
|
||||
"id_ed25519",
|
||||
".npmrc",
|
||||
".pypirc",
|
||||
];
|
||||
|
||||
/// Validates filesystem paths against allowed directories.
|
||||
pub struct FilesystemValidator {
|
||||
allowed_dirs: Vec<PathBuf>,
|
||||
read_only_dirs: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
impl FilesystemValidator {
|
||||
pub fn new(working_dir: &str) -> Self {
|
||||
let working = PathBuf::from(working_dir);
|
||||
let home = dirs_home();
|
||||
|
||||
let mut allowed_dirs = vec![
|
||||
working.clone(),
|
||||
std::env::temp_dir(),
|
||||
];
|
||||
|
||||
if let Some(home) = &home {
|
||||
allowed_dirs.push(home.join(".codeany"));
|
||||
}
|
||||
|
||||
Self {
|
||||
allowed_dirs,
|
||||
read_only_dirs: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_allowed_dir(&mut self, dir: PathBuf) {
|
||||
self.allowed_dirs.push(dir);
|
||||
}
|
||||
|
||||
pub fn add_read_only_dir(&mut self, dir: PathBuf) {
|
||||
self.read_only_dirs.push(dir);
|
||||
}
|
||||
|
||||
/// Check if a path is allowed for the given operation.
|
||||
pub fn validate_path(&self, path: &str, write: bool) -> Result<(), String> {
|
||||
let path = resolve_path(path);
|
||||
|
||||
// Check read-only dirs
|
||||
if write {
|
||||
for ro_dir in &self.read_only_dirs {
|
||||
if path.starts_with(ro_dir) {
|
||||
return Err(format!(
|
||||
"Path {} is in a read-only directory",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check allowed dirs
|
||||
let allowed = self.allowed_dirs.iter().any(|dir| path.starts_with(dir));
|
||||
if !allowed {
|
||||
return Err(format!(
|
||||
"Path {} is outside allowed directories",
|
||||
path.display()
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Check if a path is sensitive (credentials, keys, etc.).
|
||||
pub fn is_sensitive_path(path: &str) -> bool {
|
||||
let path_lower = path.to_lowercase();
|
||||
SENSITIVE_PATTERNS
|
||||
.iter()
|
||||
.any(|pattern| path_lower.ends_with(pattern) || path_lower.contains(pattern))
|
||||
}
|
||||
}
|
||||
|
||||
fn resolve_path(path: &str) -> PathBuf {
|
||||
let p = Path::new(path);
|
||||
if let Ok(canonical) = p.canonicalize() {
|
||||
canonical
|
||||
} else {
|
||||
p.to_path_buf()
|
||||
}
|
||||
}
|
||||
|
||||
fn dirs_home() -> Option<PathBuf> {
|
||||
std::env::var("HOME")
|
||||
.ok()
|
||||
.map(PathBuf::from)
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
mod filesystem;
|
||||
mod permissions;
|
||||
|
||||
pub use filesystem::*;
|
||||
pub use permissions::*;
|
||||
@@ -0,0 +1,119 @@
|
||||
use crate::types::PermissionMode;
|
||||
use serde_json::Value;
|
||||
|
||||
/// A permission rule that matches tool invocations.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Rule {
|
||||
pub tool_name: String,
|
||||
pub pattern: Option<String>,
|
||||
}
|
||||
|
||||
/// Permission configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PermissionConfig {
|
||||
pub mode: PermissionMode,
|
||||
pub allow_rules: Vec<Rule>,
|
||||
pub deny_rules: Vec<Rule>,
|
||||
pub allowed_tools: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl Default for PermissionConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
mode: PermissionMode::BypassPermissions,
|
||||
allow_rules: Vec::new(),
|
||||
deny_rules: Vec::new(),
|
||||
allowed_tools: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a permission check.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PermissionResult {
|
||||
Allow,
|
||||
Deny(String),
|
||||
}
|
||||
|
||||
/// Check if a tool invocation is allowed.
|
||||
pub fn check_permission(
|
||||
config: &PermissionConfig,
|
||||
tool_name: &str,
|
||||
input: &Value,
|
||||
) -> PermissionResult {
|
||||
// Bypass mode allows everything
|
||||
if config.mode == PermissionMode::BypassPermissions {
|
||||
return PermissionResult::Allow;
|
||||
}
|
||||
|
||||
// Check allowed tools whitelist
|
||||
if let Some(allowed) = &config.allowed_tools {
|
||||
if !allowed.iter().any(|a| a == tool_name) {
|
||||
return PermissionResult::Deny(format!("Tool '{}' is not in allowed list", tool_name));
|
||||
}
|
||||
}
|
||||
|
||||
// Check deny rules first
|
||||
for rule in &config.deny_rules {
|
||||
if matches_rule(&rule.tool_name, tool_name) {
|
||||
if let Some(pattern) = &rule.pattern {
|
||||
if matches_input_pattern(pattern, input) {
|
||||
return PermissionResult::Deny(format!(
|
||||
"Tool '{}' denied by rule",
|
||||
tool_name
|
||||
));
|
||||
}
|
||||
} else {
|
||||
return PermissionResult::Deny(format!(
|
||||
"Tool '{}' denied by rule",
|
||||
tool_name
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check allow rules
|
||||
for rule in &config.allow_rules {
|
||||
if matches_rule(&rule.tool_name, tool_name) {
|
||||
return PermissionResult::Allow;
|
||||
}
|
||||
}
|
||||
|
||||
// Default behavior based on mode
|
||||
match config.mode {
|
||||
PermissionMode::DontAsk | PermissionMode::Auto => PermissionResult::Allow,
|
||||
PermissionMode::AcceptEdits => {
|
||||
// Allow read-only tools and file edits
|
||||
let read_only_tools = [
|
||||
"Read", "Glob", "Grep", "WebFetch", "WebSearch", "TaskList", "TaskGet",
|
||||
"ToolSearch",
|
||||
];
|
||||
if read_only_tools.contains(&tool_name) {
|
||||
PermissionResult::Allow
|
||||
} else {
|
||||
PermissionResult::Allow // acceptEdits mode allows edits
|
||||
}
|
||||
}
|
||||
PermissionMode::Plan => PermissionResult::Deny(format!(
|
||||
"Tool '{}' requires approval in plan mode",
|
||||
tool_name
|
||||
)),
|
||||
_ => PermissionResult::Allow,
|
||||
}
|
||||
}
|
||||
|
||||
fn matches_rule(rule_name: &str, tool_name: &str) -> bool {
|
||||
if rule_name == "*" {
|
||||
return true;
|
||||
}
|
||||
if rule_name.ends_with('*') {
|
||||
let prefix = &rule_name[..rule_name.len() - 1];
|
||||
return tool_name.starts_with(prefix);
|
||||
}
|
||||
rule_name == tool_name
|
||||
}
|
||||
|
||||
fn matches_input_pattern(pattern: &str, input: &Value) -> bool {
|
||||
let input_str = input.to_string();
|
||||
input_str.contains(pattern)
|
||||
}
|
||||
@@ -0,0 +1,253 @@
|
||||
//! # Session Storage & Management
|
||||
//!
|
||||
//! Persists conversation transcripts to disk for resumption.
|
||||
//! Manages session lifecycle (create, resume, list, fork, delete, rename, tag).
|
||||
|
||||
use crate::types::Message;
|
||||
use chrono::Utc;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::{Path, PathBuf};
|
||||
use tokio::fs;
|
||||
|
||||
/// Session metadata persisted alongside the transcript.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SessionMetadata {
|
||||
pub id: String,
|
||||
pub cwd: String,
|
||||
pub model: String,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
pub message_count: usize,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub summary: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub title: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub tags: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
/// Session data stored on disk (metadata + messages).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SessionData {
|
||||
pub metadata: SessionMetadata,
|
||||
pub messages: Vec<Message>,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Path helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Return the root sessions directory: `~/.open-agent-sdk/sessions`.
|
||||
fn get_sessions_dir() -> PathBuf {
|
||||
let home = std::env::var("HOME")
|
||||
.or_else(|_| std::env::var("USERPROFILE"))
|
||||
.unwrap_or_else(|_| "/tmp".to_string());
|
||||
Path::new(&home)
|
||||
.join(".open-agent-sdk")
|
||||
.join("sessions")
|
||||
}
|
||||
|
||||
/// Return the directory for a specific session.
|
||||
fn get_session_path(session_id: &str) -> PathBuf {
|
||||
get_sessions_dir().join(session_id)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public API
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Persist a session to `~/.open-agent-sdk/sessions/{id}/transcript.json`.
|
||||
pub async fn save_session(
|
||||
session_id: &str,
|
||||
messages: &[Message],
|
||||
metadata: &SessionMetadata,
|
||||
) -> Result<(), String> {
|
||||
let dir = get_session_path(session_id);
|
||||
fs::create_dir_all(&dir)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to create session directory: {e}"))?;
|
||||
|
||||
let data = SessionData {
|
||||
metadata: SessionMetadata {
|
||||
id: session_id.to_string(),
|
||||
cwd: metadata.cwd.clone(),
|
||||
model: metadata.model.clone(),
|
||||
created_at: metadata.created_at.clone(),
|
||||
updated_at: Utc::now().to_rfc3339(),
|
||||
message_count: messages.len(),
|
||||
summary: metadata.summary.clone(),
|
||||
title: metadata.title.clone(),
|
||||
tags: metadata.tags.clone(),
|
||||
},
|
||||
messages: messages.to_vec(),
|
||||
};
|
||||
|
||||
let json = serde_json::to_string_pretty(&data)
|
||||
.map_err(|e| format!("Failed to serialize session: {e}"))?;
|
||||
|
||||
fs::write(dir.join("transcript.json"), json)
|
||||
.await
|
||||
.map_err(|e| format!("Failed to write session file: {e}"))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load a session from disk. Returns `None` if the session does not exist or
|
||||
/// the transcript cannot be parsed.
|
||||
pub async fn load_session(session_id: &str) -> Result<Option<SessionData>, String> {
|
||||
let file_path = get_session_path(session_id).join("transcript.json");
|
||||
match fs::read_to_string(&file_path).await {
|
||||
Ok(content) => {
|
||||
let data: SessionData = serde_json::from_str(&content)
|
||||
.map_err(|e| format!("Failed to parse session: {e}"))?;
|
||||
Ok(Some(data))
|
||||
}
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
|
||||
Err(e) => Err(format!("Failed to read session file: {e}")),
|
||||
}
|
||||
}
|
||||
|
||||
/// List all sessions sorted by `updatedAt` descending (most recent first).
|
||||
pub async fn list_sessions() -> Result<Vec<SessionMetadata>, String> {
|
||||
let dir = get_sessions_dir();
|
||||
let mut entries = match fs::read_dir(&dir).await {
|
||||
Ok(entries) => entries,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
|
||||
Err(e) => return Err(format!("Failed to read sessions directory: {e}")),
|
||||
};
|
||||
|
||||
let mut sessions: Vec<SessionMetadata> = Vec::new();
|
||||
|
||||
while let Some(entry) = entries
|
||||
.next_entry()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to read directory entry: {e}"))?
|
||||
{
|
||||
let name = entry.file_name().to_string_lossy().to_string();
|
||||
if let Ok(Some(data)) = load_session(&name).await {
|
||||
sessions.push(data.metadata);
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by updatedAt descending
|
||||
sessions.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));
|
||||
|
||||
Ok(sessions)
|
||||
}
|
||||
|
||||
/// Fork a session — create a copy with a new ID. Returns the new session ID.
|
||||
pub async fn fork_session(
|
||||
source_session_id: &str,
|
||||
new_session_id: Option<&str>,
|
||||
) -> Result<Option<String>, String> {
|
||||
let data = match load_session(source_session_id).await? {
|
||||
Some(d) => d,
|
||||
None => return Ok(None),
|
||||
};
|
||||
|
||||
let fork_id = new_session_id
|
||||
.map(|s| s.to_string())
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
|
||||
|
||||
let now = Utc::now().to_rfc3339();
|
||||
let meta = SessionMetadata {
|
||||
id: fork_id.clone(),
|
||||
cwd: data.metadata.cwd,
|
||||
model: data.metadata.model,
|
||||
created_at: now,
|
||||
updated_at: String::new(), // will be overwritten by save_session
|
||||
message_count: data.messages.len(),
|
||||
summary: Some(format!("Forked from session {source_session_id}")),
|
||||
title: data.metadata.title,
|
||||
tags: data.metadata.tags,
|
||||
};
|
||||
|
||||
save_session(&fork_id, &data.messages, &meta).await?;
|
||||
Ok(Some(fork_id))
|
||||
}
|
||||
|
||||
/// Get messages from a session.
|
||||
pub async fn get_session_messages(session_id: &str) -> Result<Vec<Message>, String> {
|
||||
match load_session(session_id).await? {
|
||||
Some(data) => Ok(data.messages),
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Append a message to an existing session.
|
||||
pub async fn append_to_session(session_id: &str, message: Message) -> Result<(), String> {
|
||||
let mut data = match load_session(session_id).await? {
|
||||
Some(d) => d,
|
||||
None => return Err(format!("Session '{session_id}' not found")),
|
||||
};
|
||||
|
||||
data.messages.push(message);
|
||||
data.metadata.message_count = data.messages.len();
|
||||
|
||||
save_session(session_id, &data.messages, &data.metadata).await
|
||||
}
|
||||
|
||||
/// Delete a session directory from disk.
|
||||
pub async fn delete_session(session_id: &str) -> Result<bool, String> {
|
||||
let path = get_session_path(session_id);
|
||||
match fs::remove_dir_all(&path).await {
|
||||
Ok(()) => Ok(true),
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
|
||||
Err(e) => Err(format!("Failed to delete session: {e}")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get metadata about a specific session.
|
||||
pub async fn get_session_info(session_id: &str) -> Result<Option<SessionMetadata>, String> {
|
||||
match load_session(session_id).await? {
|
||||
Some(data) => Ok(Some(data.metadata)),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Rename (update the title of) a session.
|
||||
pub async fn rename_session(session_id: &str, title: &str) -> Result<(), String> {
|
||||
let data = match load_session(session_id).await? {
|
||||
Some(d) => d,
|
||||
None => return Err(format!("Session '{session_id}' not found")),
|
||||
};
|
||||
|
||||
let meta = SessionMetadata {
|
||||
title: Some(title.to_string()),
|
||||
..data.metadata
|
||||
};
|
||||
|
||||
save_session(session_id, &data.messages, &meta).await
|
||||
}
|
||||
|
||||
/// Add or update tags on a session. Pass `None` to clear tags.
|
||||
pub async fn tag_session(session_id: &str, tags: Option<Vec<String>>) -> Result<(), String> {
|
||||
let data = match load_session(session_id).await? {
|
||||
Some(d) => d,
|
||||
None => return Err(format!("Session '{session_id}' not found")),
|
||||
};
|
||||
|
||||
let meta = SessionMetadata {
|
||||
tags,
|
||||
..data.metadata
|
||||
};
|
||||
|
||||
save_session(session_id, &data.messages, &meta).await
|
||||
}
|
||||
|
||||
/// Create a new `SessionMetadata` with sensible defaults.
|
||||
pub fn new_metadata(id: &str, cwd: &str, model: &str) -> SessionMetadata {
|
||||
let now = Utc::now().to_rfc3339();
|
||||
SessionMetadata {
|
||||
id: id.to_string(),
|
||||
cwd: cwd.to_string(),
|
||||
model: model.to_string(),
|
||||
created_at: now.clone(),
|
||||
updated_at: now,
|
||||
message_count: 0,
|
||||
summary: None,
|
||||
title: None,
|
||||
tags: None,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use crate::types::skill::{LoadedSkill, SkillMetadata, SkillSource};
|
||||
|
||||
/// Parse a SKILL.md file into (metadata, markdown_body).
|
||||
/// Frontmatter is delimited by `---` lines at the start of the file.
|
||||
pub fn parse_skill_file(content: &str) -> Option<(SkillMetadata, String)> {
|
||||
let trimmed = content.trim_start();
|
||||
if !trimmed.starts_with("---") {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Find the closing `---`
|
||||
let after_first = &trimmed[3..];
|
||||
let rest = after_first.trim_start_matches(['\r', '\n']);
|
||||
let end_idx = rest.find("\n---")?;
|
||||
let yaml_str = &rest[..end_idx];
|
||||
let after_closing = &rest[end_idx + 4..]; // skip "\n---"
|
||||
let body = after_closing.trim_start_matches(['\r', '\n']).to_string();
|
||||
|
||||
let metadata: SkillMetadata = serde_yaml::from_str(yaml_str).ok()?;
|
||||
Some((metadata, body))
|
||||
}
|
||||
|
||||
/// Load a single skill from a directory containing SKILL.md.
|
||||
pub fn load_skill_from_dir(dir: &Path, source: SkillSource) -> Option<LoadedSkill> {
|
||||
let skill_file = dir.join("SKILL.md");
|
||||
let content = std::fs::read_to_string(&skill_file).ok()?;
|
||||
let (metadata, body) = parse_skill_file(&content)?;
|
||||
|
||||
let name = metadata
|
||||
.name
|
||||
.clone();
|
||||
let name = if name.is_empty() {
|
||||
dir.file_name()?.to_string_lossy().to_string()
|
||||
} else {
|
||||
name
|
||||
};
|
||||
|
||||
Some(LoadedSkill {
|
||||
name,
|
||||
path: skill_file,
|
||||
metadata,
|
||||
content: body,
|
||||
source,
|
||||
group: None,
|
||||
})
|
||||
}
|
||||
|
||||
/// Scan all subdirectories of `parent_dir` for SKILL.md files.
|
||||
pub fn load_from_dir(parent_dir: &Path, source: SkillSource) -> Vec<LoadedSkill> {
|
||||
let entries = match std::fs::read_dir(parent_dir) {
|
||||
Ok(e) => e,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
|
||||
let mut skills = Vec::new();
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.is_dir() {
|
||||
if let Some(skill) = load_skill_from_dir(&path, source.clone()) {
|
||||
skills.push(skill);
|
||||
} else {
|
||||
// No SKILL.md — treat as a group container and recurse
|
||||
let group_name = match path.file_name() {
|
||||
Some(n) => n.to_string_lossy().to_string(),
|
||||
None => continue,
|
||||
};
|
||||
let mut group_skills = Vec::new();
|
||||
collect_skills_recursive(&path, source.clone(), &mut group_skills, 0, 3);
|
||||
for mut skill in group_skills {
|
||||
skill.group = Some(group_name.clone());
|
||||
skills.push(skill);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
skills
|
||||
}
|
||||
|
||||
/// Recursively collect skills from nested directories up to `max_depth`.
|
||||
fn collect_skills_recursive(
|
||||
dir: &Path,
|
||||
source: SkillSource,
|
||||
skills: &mut Vec<LoadedSkill>,
|
||||
depth: u32,
|
||||
max_depth: u32,
|
||||
) {
|
||||
if depth >= max_depth {
|
||||
return;
|
||||
}
|
||||
let entries = match std::fs::read_dir(dir) {
|
||||
Ok(e) => e,
|
||||
Err(_) => return,
|
||||
};
|
||||
for entry in entries.flatten() {
|
||||
let path = entry.path();
|
||||
if path.is_dir() {
|
||||
if let Some(skill) = load_skill_from_dir(&path, source.clone()) {
|
||||
skills.push(skill);
|
||||
} else {
|
||||
collect_skills_recursive(&path, source.clone(), skills, depth + 1, max_depth);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Load skills from all global directories in priority order.
|
||||
/// Deduplicates by name — first found wins.
|
||||
pub fn load_all_global(home_dir: &Path) -> Vec<LoadedSkill> {
|
||||
let dirs: Vec<(PathBuf, SkillSource)> = vec![
|
||||
(home_dir.join(".boosapi").join("skills"), SkillSource::BoosApi),
|
||||
(home_dir.join(".claude").join("skills"), SkillSource::Claude),
|
||||
(home_dir.join(".agents").join("skills"), SkillSource::Agents),
|
||||
];
|
||||
|
||||
let mut result = Vec::new();
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
|
||||
for (dir, source) in dirs {
|
||||
for skill in load_from_dir(&dir, source) {
|
||||
if seen.insert(skill.name.clone()) {
|
||||
result.push(skill);
|
||||
}
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
||||
/// Walk up from `cwd` looking for `.git/` or `.boosapi/` to detect the project root.
|
||||
pub fn detect_project_root(cwd: &Path) -> Option<PathBuf> {
|
||||
let mut current = cwd.to_path_buf();
|
||||
loop {
|
||||
if current.join(".git").is_dir() || current.join(".boosapi").is_dir() {
|
||||
return Some(current);
|
||||
}
|
||||
if !current.pop() {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Load skills from the project's `.boosapi/skills/` directory.
|
||||
pub fn load_project_skills(cwd: &Path) -> Vec<LoadedSkill> {
|
||||
let root = match detect_project_root(cwd) {
|
||||
Some(r) => r,
|
||||
None => return Vec::new(),
|
||||
};
|
||||
load_from_dir(&root.join(".boosapi").join("skills"), SkillSource::Project)
|
||||
}
|
||||
|
||||
/// Substitute template variables in skill content.
|
||||
/// Replaces `$ARGUMENTS`, `${SKILL_DIR}`, and `${CLAUDE_SKILL_DIR}` (compat).
|
||||
pub fn substitute_variables(content: &str, arguments: &str, skill_dir: &str) -> String {
|
||||
content
|
||||
.replace("$ARGUMENTS", arguments)
|
||||
.replace("${SKILL_DIR}", skill_dir)
|
||||
.replace("${CLAUDE_SKILL_DIR}", skill_dir)
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod loader;
|
||||
pub mod registry;
|
||||
|
||||
pub use loader::*;
|
||||
pub use registry::*;
|
||||
@@ -0,0 +1,102 @@
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
use crate::types::skill::LoadedSkill;
|
||||
|
||||
pub struct SkillRegistry {
|
||||
skills: HashMap<String, LoadedSkill>,
|
||||
disabled: HashSet<String>,
|
||||
}
|
||||
|
||||
impl SkillRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
skills: HashMap::new(),
|
||||
disabled: HashSet::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Register a skill. If a skill with the same name exists, ignore (first wins).
|
||||
pub fn register(&mut self, skill: LoadedSkill) {
|
||||
self.skills.entry(skill.name.clone()).or_insert(skill);
|
||||
}
|
||||
|
||||
/// Set the disabled skills set (from database).
|
||||
pub fn set_disabled(&mut self, disabled: HashSet<String>) {
|
||||
self.disabled = disabled;
|
||||
}
|
||||
|
||||
/// Get all registered skills (including disabled).
|
||||
pub fn all(&self) -> Vec<&LoadedSkill> {
|
||||
self.skills.values().collect()
|
||||
}
|
||||
|
||||
/// Get all enabled skills.
|
||||
pub fn all_enabled(&self) -> Vec<&LoadedSkill> {
|
||||
self.skills
|
||||
.values()
|
||||
.filter(|s| !self.disabled.contains(&s.name))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get a skill by name (regardless of enabled/disabled).
|
||||
pub fn get(&self, name: &str) -> Option<&LoadedSkill> {
|
||||
self.skills.get(name)
|
||||
}
|
||||
|
||||
/// Get skills that the model can auto-invoke (enabled + not disable_model_invocation).
|
||||
pub fn auto_invocable(&self) -> Vec<&LoadedSkill> {
|
||||
self.skills
|
||||
.values()
|
||||
.filter(|s| !self.disabled.contains(&s.name) && !s.metadata.disable_model_invocation)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get skills that the user can invoke (enabled + user_invocable).
|
||||
pub fn user_invocable(&self) -> Vec<&LoadedSkill> {
|
||||
self.skills
|
||||
.values()
|
||||
.filter(|s| !self.disabled.contains(&s.name) && s.metadata.user_invocable)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Generate a summary for system prompt injection (only enabled skills).
|
||||
pub fn generate_context_summary(&self) -> String {
|
||||
let skills = self.all_enabled();
|
||||
if skills.is_empty() {
|
||||
return String::new();
|
||||
}
|
||||
let mut summary = String::from(
|
||||
"# Available Skills\n\nYou have access to the following skills via the Skill tool:\n\n",
|
||||
);
|
||||
for skill in &skills {
|
||||
summary.push_str(&format!("- **{}**", skill.name));
|
||||
if let Some(desc) = &skill.metadata.description {
|
||||
let desc = if desc.len() > 250 {
|
||||
format!("{}...", &desc[..250])
|
||||
} else {
|
||||
desc.clone()
|
||||
};
|
||||
summary.push_str(&format!(": {}", desc));
|
||||
}
|
||||
if let Some(when) = &skill.metadata.when_to_use {
|
||||
summary.push_str(&format!(" TRIGGER when: {}", when));
|
||||
}
|
||||
if let Some(hint) = &skill.metadata.argument_hint {
|
||||
summary.push_str(&format!(" (usage: /{})", hint));
|
||||
}
|
||||
summary.push('\n');
|
||||
}
|
||||
summary.push_str(
|
||||
"\nWhen a skill matches the user's request, invoke it using the Skill tool. \
|
||||
When the Skill tool returns instructions, treat them as mandatory operational \
|
||||
procedures, not as optional reference material. Follow the instructions step by step.\n",
|
||||
);
|
||||
summary
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for SkillRegistry {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// Structured request passed to the ask callback.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AskUserRequest {
|
||||
pub question: String,
|
||||
pub options: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
/// Callback function for asking the user a question.
|
||||
pub type AskUserFn = Arc<
|
||||
dyn Fn(AskUserRequest) -> futures::future::BoxFuture<'static, Result<String, String>>
|
||||
+ Send
|
||||
+ Sync,
|
||||
>;
|
||||
|
||||
pub struct AskUserTool {
|
||||
ask_fn: Option<AskUserFn>,
|
||||
}
|
||||
|
||||
impl Default for AskUserTool {
|
||||
fn default() -> Self {
|
||||
Self { ask_fn: None }
|
||||
}
|
||||
}
|
||||
|
||||
impl AskUserTool {
|
||||
pub fn new(ask_fn: AskUserFn) -> Self {
|
||||
Self {
|
||||
ask_fn: Some(ask_fn),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for AskUserTool {
|
||||
fn name(&self) -> &str {
|
||||
"AskUserQuestion"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Ask the user a question and wait for their response. Use when you need clarification or input. Optionally provide a list of options for the user to choose from."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"question".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The question to ask the user"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"options".to_string(),
|
||||
json!({
|
||||
"type": "array",
|
||||
"items": { "type": "string" },
|
||||
"description": "Optional list of choices for the user to select from (multi-select). The user can also provide additional free-form input."
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["question".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let question = input
|
||||
.get("question")
|
||||
.and_then(|q| q.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'question' field".to_string()))?;
|
||||
|
||||
let options: Option<Vec<String>> = input.get("options").and_then(|o| {
|
||||
o.as_array().map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str().map(String::from))
|
||||
.collect()
|
||||
})
|
||||
});
|
||||
|
||||
let request = AskUserRequest {
|
||||
question: question.to_string(),
|
||||
options,
|
||||
};
|
||||
|
||||
match &self.ask_fn {
|
||||
Some(ask_fn) => {
|
||||
let answer = (ask_fn)(request)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionError(e))?;
|
||||
Ok(ToolResult::text(answer))
|
||||
}
|
||||
None => Ok(ToolResult::error(
|
||||
"User interaction is not configured. Provide an ask_fn when creating the agent."
|
||||
.to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,368 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::process::Stdio;
|
||||
use tokio::process::Command;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
const DEFAULT_TIMEOUT_MS: u64 = 120_000;
|
||||
const MAX_TIMEOUT_MS: u64 = 600_000;
|
||||
const MAX_OUTPUT_SIZE: usize = 100_000;
|
||||
|
||||
/// Destructive command patterns that should be flagged.
|
||||
const DESTRUCTIVE_PATTERNS: &[&str] = &[
|
||||
"rm -rf /",
|
||||
"rm -rf ~",
|
||||
"rm -rf .",
|
||||
"git push --force",
|
||||
"git push -f",
|
||||
"git reset --hard",
|
||||
"chmod 777",
|
||||
"chmod -R 777",
|
||||
"> /dev/sda",
|
||||
"mkfs.",
|
||||
"dd if=",
|
||||
":(){ :|:& };:",
|
||||
];
|
||||
|
||||
pub struct BashTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for BashTool {
|
||||
fn name(&self) -> &str {
|
||||
"Bash"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Executes a given bash command and returns its output. Use for system commands and terminal operations that require shell execution."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"command".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The command to execute"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"timeout".to_string(),
|
||||
json!({
|
||||
"type": "number",
|
||||
"description": "Optional timeout in milliseconds (max 600000)"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"description".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "Clear description of what this command does"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["command".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, input: &Value) -> bool {
|
||||
let command = input.get("command").and_then(|c| c.as_str()).unwrap_or("");
|
||||
|
||||
// Check if command starts with a read-only command
|
||||
let cmd_trimmed = command.trim();
|
||||
let first_cmd = cmd_trimmed.split_whitespace().next().unwrap_or("");
|
||||
let single_word_reads = [
|
||||
"ls", "cat", "head", "tail", "find", "grep", "rg", "wc", "pwd", "echo", "which",
|
||||
"type", "file", "stat", "du", "df",
|
||||
];
|
||||
let prefix_reads = [
|
||||
"git status",
|
||||
"git log",
|
||||
"git diff",
|
||||
"git show",
|
||||
"git branch",
|
||||
"cargo check",
|
||||
"cargo test --no-run",
|
||||
"rustc --version",
|
||||
];
|
||||
|
||||
single_word_reads.contains(&first_cmd)
|
||||
|| prefix_reads.iter().any(|p| cmd_trimmed.starts_with(p))
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let command = input
|
||||
.get("command")
|
||||
.and_then(|c| c.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'command' field".to_string()))?;
|
||||
|
||||
// Security check
|
||||
if let Some(warning) = check_destructive(command) {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"Potentially destructive command detected: {}. Proceed with caution.",
|
||||
warning
|
||||
)));
|
||||
}
|
||||
|
||||
let timeout_ms = input
|
||||
.get("timeout")
|
||||
.and_then(|t| t.as_u64())
|
||||
.unwrap_or(DEFAULT_TIMEOUT_MS)
|
||||
.min(MAX_TIMEOUT_MS);
|
||||
|
||||
let output = tokio::time::timeout(
|
||||
std::time::Duration::from_millis(timeout_ms),
|
||||
run_command(command, &context.working_dir),
|
||||
)
|
||||
.await;
|
||||
|
||||
match output {
|
||||
Ok(Ok((stdout, stderr, exit_code))) => {
|
||||
let mut result = String::new();
|
||||
|
||||
if !stdout.is_empty() {
|
||||
result.push_str(&stdout);
|
||||
}
|
||||
if !stderr.is_empty() {
|
||||
if !result.is_empty() {
|
||||
result.push('\n');
|
||||
}
|
||||
result.push_str("STDERR:\n");
|
||||
result.push_str(&stderr);
|
||||
}
|
||||
|
||||
if result.len() > MAX_OUTPUT_SIZE {
|
||||
result.truncate(MAX_OUTPUT_SIZE);
|
||||
result.push_str("\n... (output truncated)");
|
||||
}
|
||||
|
||||
if exit_code != 0 {
|
||||
result.push_str(&format!("\n\nExit code: {}", exit_code));
|
||||
}
|
||||
|
||||
if result.is_empty() {
|
||||
result = "(no output)".to_string();
|
||||
}
|
||||
|
||||
Ok(if exit_code != 0 {
|
||||
ToolResult::error(result)
|
||||
} else {
|
||||
ToolResult::text(result)
|
||||
})
|
||||
}
|
||||
Ok(Err(e)) => Ok(ToolResult::error(format!("Command failed: {}", e))),
|
||||
Err(_) => Ok(ToolResult::error(format!(
|
||||
"Command timed out after {}ms",
|
||||
timeout_ms
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_command(
|
||||
command: &str,
|
||||
working_dir: &str,
|
||||
) -> Result<(String, String, i32), std::io::Error> {
|
||||
let shell_path = crate::mcp::shell_path::get_shell_path();
|
||||
let shell = build_shell_runner(command);
|
||||
let mut cmd = Command::new(&shell.program);
|
||||
cmd.args(&shell.args)
|
||||
.current_dir(working_dir)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
if !shell_path.is_empty() {
|
||||
cmd.env("PATH", shell_path);
|
||||
}
|
||||
let output = cmd.output().await?;
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout).to_string();
|
||||
let stderr = String::from_utf8_lossy(&output.stderr).to_string();
|
||||
let exit_code = output.status.code().unwrap_or(-1);
|
||||
|
||||
Ok((stdout, stderr, exit_code))
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct ShellRunner {
|
||||
program: String,
|
||||
args: Vec<String>,
|
||||
}
|
||||
|
||||
fn build_shell_runner(command: &str) -> ShellRunner {
|
||||
select_shell_runner(
|
||||
command,
|
||||
cfg!(windows),
|
||||
std::env::var("ComSpec").ok(),
|
||||
program_exists,
|
||||
)
|
||||
}
|
||||
|
||||
fn select_shell_runner<F>(
|
||||
command: &str,
|
||||
is_windows: bool,
|
||||
comspec: Option<String>,
|
||||
mut has_program: F,
|
||||
) -> ShellRunner
|
||||
where
|
||||
F: FnMut(&str) -> bool,
|
||||
{
|
||||
if is_windows {
|
||||
for candidate in ["bash.exe", "bash"] {
|
||||
if has_program(candidate) {
|
||||
return ShellRunner {
|
||||
program: candidate.to_string(),
|
||||
args: vec!["-c".to_string(), command.to_string()],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
for candidate in ["pwsh.exe", "pwsh", "powershell.exe", "powershell"] {
|
||||
if has_program(candidate) {
|
||||
return ShellRunner {
|
||||
program: candidate.to_string(),
|
||||
args: vec![
|
||||
"-NoLogo".to_string(),
|
||||
"-NoProfile".to_string(),
|
||||
"-NonInteractive".to_string(),
|
||||
"-Command".to_string(),
|
||||
command.to_string(),
|
||||
],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
return ShellRunner {
|
||||
program: comspec
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_else(|| "cmd.exe".to_string()),
|
||||
args: vec![
|
||||
"/d".to_string(),
|
||||
"/s".to_string(),
|
||||
"/c".to_string(),
|
||||
command.to_string(),
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
for candidate in ["/bin/bash", "bash", "/bin/sh", "sh"] {
|
||||
if has_program(candidate) {
|
||||
return ShellRunner {
|
||||
program: candidate.to_string(),
|
||||
args: vec!["-c".to_string(), command.to_string()],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
ShellRunner {
|
||||
program: "sh".to_string(),
|
||||
args: vec!["-c".to_string(), command.to_string()],
|
||||
}
|
||||
}
|
||||
|
||||
fn program_exists(candidate: &str) -> bool {
|
||||
if candidate.contains(std::path::MAIN_SEPARATOR)
|
||||
|| candidate.contains('/')
|
||||
|| candidate.contains('\\')
|
||||
{
|
||||
return Path::new(candidate).is_file();
|
||||
}
|
||||
|
||||
std::env::var_os("PATH")
|
||||
.map(|paths| {
|
||||
std::env::split_paths(&paths).any(|dir| {
|
||||
let full_path = dir.join(candidate);
|
||||
if full_path.is_file() {
|
||||
return true;
|
||||
}
|
||||
|
||||
#[cfg(windows)]
|
||||
{
|
||||
let exe_path = dir.join(format!("{candidate}.exe"));
|
||||
return exe_path.is_file();
|
||||
}
|
||||
|
||||
#[cfg(not(windows))]
|
||||
{
|
||||
false
|
||||
}
|
||||
})
|
||||
})
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn check_destructive(command: &str) -> Option<&'static str> {
|
||||
for pattern in DESTRUCTIVE_PATTERNS {
|
||||
if command.contains(pattern) {
|
||||
return Some(pattern);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{select_shell_runner, ShellRunner};
|
||||
|
||||
#[test]
|
||||
fn select_shell_runner_falls_back_to_powershell_on_windows() {
|
||||
let runner = select_shell_runner("node -v", true, None, |candidate| {
|
||||
matches!(candidate, "pwsh.exe")
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
runner,
|
||||
ShellRunner {
|
||||
program: "pwsh.exe".to_string(),
|
||||
args: vec![
|
||||
"-NoLogo".to_string(),
|
||||
"-NoProfile".to_string(),
|
||||
"-NonInteractive".to_string(),
|
||||
"-Command".to_string(),
|
||||
"node -v".to_string(),
|
||||
],
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_shell_runner_falls_back_to_cmd_on_windows() {
|
||||
let runner = select_shell_runner(
|
||||
"node -v",
|
||||
true,
|
||||
Some("C:\\Windows\\System32\\cmd.exe".to_string()),
|
||||
|_| false,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
runner,
|
||||
ShellRunner {
|
||||
program: "C:\\Windows\\System32\\cmd.exe".to_string(),
|
||||
args: vec![
|
||||
"/d".to_string(),
|
||||
"/s".to_string(),
|
||||
"/c".to_string(),
|
||||
"node -v".to_string(),
|
||||
],
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_shell_runner_prefers_bash_on_unix() {
|
||||
let runner =
|
||||
select_shell_runner("node -v", false, None, |candidate| candidate == "/bin/bash");
|
||||
|
||||
assert_eq!(
|
||||
runner,
|
||||
ShellRunner {
|
||||
program: "/bin/bash".to_string(),
|
||||
args: vec!["-c".to_string(), "node -v".to_string()],
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,136 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// In-memory config store for session configuration.
|
||||
#[derive(Clone)]
|
||||
pub struct ConfigStore {
|
||||
data: Arc<RwLock<HashMap<String, Value>>>,
|
||||
}
|
||||
|
||||
impl ConfigStore {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
data: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// ConfigTool for get/set/list session settings.
|
||||
pub struct ConfigTool {
|
||||
store: ConfigStore,
|
||||
}
|
||||
|
||||
impl ConfigTool {
|
||||
pub fn new(store: ConfigStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ConfigTool {
|
||||
fn name(&self) -> &str {
|
||||
"Config"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Get or set configuration values. Supports session-scoped settings."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"action".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["get", "set", "list"],
|
||||
"description": "Operation to perform"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"key".to_string(),
|
||||
json!({ "type": "string", "description": "Config key" }),
|
||||
),
|
||||
(
|
||||
"value".to_string(),
|
||||
json!({ "description": "Config value (for set)" }),
|
||||
),
|
||||
]),
|
||||
required: vec!["action".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, input: &Value) -> bool {
|
||||
input
|
||||
.get("action")
|
||||
.and_then(|a| a.as_str())
|
||||
.map(|a| a == "get" || a == "list")
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let action = input
|
||||
.get("action")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'action'".to_string()))?;
|
||||
|
||||
match action {
|
||||
"get" => {
|
||||
let key = input
|
||||
.get("key")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("key required for get".to_string()))?;
|
||||
|
||||
let data = self.store.data.read().await;
|
||||
match data.get(key) {
|
||||
Some(value) => Ok(ToolResult::text(serde_json::to_string(value).unwrap_or_default())),
|
||||
None => Ok(ToolResult::text(format!("Config key \"{}\" not found", key))),
|
||||
}
|
||||
}
|
||||
"set" => {
|
||||
let key = input
|
||||
.get("key")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("key required for set".to_string()))?;
|
||||
|
||||
let value = input.get("value").cloned().unwrap_or(Value::Null);
|
||||
|
||||
let mut data = self.store.data.write().await;
|
||||
data.insert(key.to_string(), value.clone());
|
||||
|
||||
Ok(ToolResult::text(format!(
|
||||
"Config set: {} = {}",
|
||||
key,
|
||||
serde_json::to_string(&value).unwrap_or_default()
|
||||
)))
|
||||
}
|
||||
"list" => {
|
||||
let data = self.store.data.read().await;
|
||||
|
||||
if data.is_empty() {
|
||||
return Ok(ToolResult::text("No config values set."));
|
||||
}
|
||||
|
||||
let mut entries: Vec<(&String, &Value)> = data.iter().collect();
|
||||
entries.sort_by_key(|(k, _)| (*k).clone());
|
||||
|
||||
let lines: Vec<String> = entries
|
||||
.iter()
|
||||
.map(|(k, v)| {
|
||||
format!("{} = {}", k, serde_json::to_string(v).unwrap_or_default())
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(ToolResult::text(lines.join("\n")))
|
||||
}
|
||||
_ => Ok(ToolResult::error(format!("Unknown action: {}", action))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// A cron job definition.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CronJob {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub schedule: String,
|
||||
pub command: String,
|
||||
pub enabled: bool,
|
||||
pub created_at: String,
|
||||
#[serde(default)]
|
||||
pub last_run_at: Option<String>,
|
||||
#[serde(default)]
|
||||
pub next_run_at: Option<String>,
|
||||
}
|
||||
|
||||
/// In-memory cron job store shared across cron tools.
|
||||
#[derive(Clone)]
|
||||
pub struct CronStore {
|
||||
jobs: Arc<RwLock<HashMap<String, CronJob>>>,
|
||||
counter: Arc<RwLock<u64>>,
|
||||
}
|
||||
|
||||
impl CronStore {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
jobs: Arc::new(RwLock::new(HashMap::new())),
|
||||
counter: Arc::new(RwLock::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_id(&self) -> String {
|
||||
let mut counter = self.counter.write().await;
|
||||
*counter += 1;
|
||||
format!("cron_{}", counter)
|
||||
}
|
||||
}
|
||||
|
||||
// --- CronCreateTool ---
|
||||
|
||||
pub struct CronCreateTool {
|
||||
store: CronStore,
|
||||
}
|
||||
|
||||
impl CronCreateTool {
|
||||
pub fn new(store: CronStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for CronCreateTool {
|
||||
fn name(&self) -> &str {
|
||||
"CronCreate"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Create a scheduled recurring task (cron job). Supports cron expressions for scheduling."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"name".to_string(),
|
||||
json!({ "type": "string", "description": "Job name" }),
|
||||
),
|
||||
(
|
||||
"schedule".to_string(),
|
||||
json!({ "type": "string", "description": "Cron expression (e.g., \"*/5 * * * *\" for every 5 minutes)" }),
|
||||
),
|
||||
(
|
||||
"command".to_string(),
|
||||
json!({ "type": "string", "description": "Command or prompt to execute" }),
|
||||
),
|
||||
]),
|
||||
required: vec![
|
||||
"name".to_string(),
|
||||
"schedule".to_string(),
|
||||
"command".to_string(),
|
||||
],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let name = input
|
||||
.get("name")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'name'".to_string()))?;
|
||||
let schedule = input
|
||||
.get("schedule")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'schedule'".to_string()))?;
|
||||
let command = input
|
||||
.get("command")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'command'".to_string()))?;
|
||||
|
||||
let id = self.store.next_id().await;
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
|
||||
let job = CronJob {
|
||||
id: id.clone(),
|
||||
name: name.to_string(),
|
||||
schedule: schedule.to_string(),
|
||||
command: command.to_string(),
|
||||
enabled: true,
|
||||
created_at: now,
|
||||
last_run_at: None,
|
||||
next_run_at: None,
|
||||
};
|
||||
|
||||
let mut jobs = self.store.jobs.write().await;
|
||||
jobs.insert(id.clone(), job.clone());
|
||||
|
||||
Ok(ToolResult::text(format!(
|
||||
"Cron job created: {} \"{}\" schedule=\"{}\"",
|
||||
id, job.name, job.schedule
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
// --- CronDeleteTool ---
|
||||
|
||||
pub struct CronDeleteTool {
|
||||
store: CronStore,
|
||||
}
|
||||
|
||||
impl CronDeleteTool {
|
||||
pub fn new(store: CronStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for CronDeleteTool {
|
||||
fn name(&self) -> &str {
|
||||
"CronDelete"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Delete a scheduled cron job."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([(
|
||||
"id".to_string(),
|
||||
json!({ "type": "string", "description": "Cron job ID to delete" }),
|
||||
)]),
|
||||
required: vec!["id".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let id = input
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'id'".to_string()))?;
|
||||
|
||||
let mut jobs = self.store.jobs.write().await;
|
||||
if jobs.remove(id).is_some() {
|
||||
Ok(ToolResult::text(format!("Cron job deleted: {}", id)))
|
||||
} else {
|
||||
Ok(ToolResult::error(format!("Cron job not found: {}", id)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- CronListTool ---
|
||||
|
||||
pub struct CronListTool {
|
||||
store: CronStore,
|
||||
}
|
||||
|
||||
impl CronListTool {
|
||||
pub fn new(store: CronStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for CronListTool {
|
||||
fn name(&self) -> &str {
|
||||
"CronList"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"List all scheduled cron jobs."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema::default()
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(
|
||||
&self,
|
||||
_input: Value,
|
||||
_context: &ToolUseContext,
|
||||
) -> Result<ToolResult, ToolError> {
|
||||
let jobs = self.store.jobs.read().await;
|
||||
|
||||
if jobs.is_empty() {
|
||||
return Ok(ToolResult::text("No cron jobs scheduled."));
|
||||
}
|
||||
|
||||
let mut job_list: Vec<&CronJob> = jobs.values().collect();
|
||||
job_list.sort_by(|a, b| a.created_at.cmp(&b.created_at));
|
||||
|
||||
let lines: Vec<String> = job_list
|
||||
.iter()
|
||||
.map(|j| {
|
||||
let status = if j.enabled { "enabled" } else { "disabled" };
|
||||
format!(
|
||||
"[{}] {} \"{}\" schedule=\"{}\" command=\"{}\"",
|
||||
j.id,
|
||||
status,
|
||||
j.name,
|
||||
j.schedule,
|
||||
if j.command.len() > 50 {
|
||||
format!("{}...", &j.command[..50])
|
||||
} else {
|
||||
j.command.clone()
|
||||
}
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(ToolResult::text(lines.join("\n")))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
/// Generate a unified diff between two strings.
|
||||
pub fn unified_diff(old: &str, new: &str, filename: &str) -> String {
|
||||
let old_lines: Vec<&str> = old.lines().collect();
|
||||
let new_lines: Vec<&str> = new.lines().collect();
|
||||
|
||||
let mut result = String::new();
|
||||
result.push_str(&format!("--- a/{}\n", filename));
|
||||
result.push_str(&format!("+++ b/{}\n", filename));
|
||||
|
||||
let hunks = compute_hunks(&old_lines, &new_lines);
|
||||
|
||||
for hunk in hunks {
|
||||
result.push_str(&format!(
|
||||
"@@ -{},{} +{},{} @@\n",
|
||||
hunk.old_start + 1,
|
||||
hunk.old_count,
|
||||
hunk.new_start + 1,
|
||||
hunk.new_count
|
||||
));
|
||||
|
||||
for line in &hunk.lines {
|
||||
match line {
|
||||
DiffLine::Context(s) => {
|
||||
result.push(' ');
|
||||
result.push_str(s);
|
||||
result.push('\n');
|
||||
}
|
||||
DiffLine::Added(s) => {
|
||||
result.push('+');
|
||||
result.push_str(s);
|
||||
result.push('\n');
|
||||
}
|
||||
DiffLine::Removed(s) => {
|
||||
result.push('-');
|
||||
result.push_str(s);
|
||||
result.push('\n');
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
/// Count lines added and removed.
|
||||
pub fn count_changes(old: &str, new: &str) -> (usize, usize) {
|
||||
let old_lines: Vec<&str> = old.lines().collect();
|
||||
let new_lines: Vec<&str> = new.lines().collect();
|
||||
let hunks = compute_hunks(&old_lines, &new_lines);
|
||||
|
||||
let mut added = 0;
|
||||
let mut removed = 0;
|
||||
for hunk in hunks {
|
||||
for line in &hunk.lines {
|
||||
match line {
|
||||
DiffLine::Added(_) => added += 1,
|
||||
DiffLine::Removed(_) => removed += 1,
|
||||
DiffLine::Context(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
(added, removed)
|
||||
}
|
||||
|
||||
enum DiffLine<'a> {
|
||||
Context(&'a str),
|
||||
Added(&'a str),
|
||||
Removed(&'a str),
|
||||
}
|
||||
|
||||
struct Hunk<'a> {
|
||||
old_start: usize,
|
||||
old_count: usize,
|
||||
new_start: usize,
|
||||
new_count: usize,
|
||||
lines: Vec<DiffLine<'a>>,
|
||||
}
|
||||
|
||||
fn compute_hunks<'a>(old: &[&'a str], new: &[&'a str]) -> Vec<Hunk<'a>> {
|
||||
// Simple LCS-based diff
|
||||
let n = old.len();
|
||||
let m = new.len();
|
||||
|
||||
// Compute edit script using Myers-like approach (simplified)
|
||||
let mut edits: Vec<DiffLine<'a>> = Vec::new();
|
||||
let mut old_idx = 0;
|
||||
let mut new_idx = 0;
|
||||
|
||||
// Simple line-by-line comparison with longest common subsequence
|
||||
let lcs = compute_lcs(old, new);
|
||||
let mut lcs_idx = 0;
|
||||
|
||||
while old_idx < n || new_idx < m {
|
||||
if lcs_idx < lcs.len() && old_idx == lcs[lcs_idx].0 && new_idx == lcs[lcs_idx].1 {
|
||||
edits.push(DiffLine::Context(old[old_idx]));
|
||||
old_idx += 1;
|
||||
new_idx += 1;
|
||||
lcs_idx += 1;
|
||||
} else if old_idx < n
|
||||
&& (lcs_idx >= lcs.len() || old_idx < lcs[lcs_idx].0)
|
||||
{
|
||||
edits.push(DiffLine::Removed(old[old_idx]));
|
||||
old_idx += 1;
|
||||
} else if new_idx < m {
|
||||
edits.push(DiffLine::Added(new[new_idx]));
|
||||
new_idx += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Group edits into hunks with 3 lines of context
|
||||
let context_lines = 3;
|
||||
let mut hunks = Vec::new();
|
||||
let mut i = 0;
|
||||
|
||||
while i < edits.len() {
|
||||
// Find next change
|
||||
while i < edits.len() {
|
||||
if !matches!(edits[i], DiffLine::Context(_)) {
|
||||
break;
|
||||
}
|
||||
i += 1;
|
||||
}
|
||||
if i >= edits.len() {
|
||||
break;
|
||||
}
|
||||
|
||||
// Start hunk with context before
|
||||
let context_start = if i >= context_lines {
|
||||
i - context_lines
|
||||
} else {
|
||||
0
|
||||
};
|
||||
|
||||
let mut hunk_lines = Vec::new();
|
||||
|
||||
// Calculate starting positions
|
||||
let mut oi = 0;
|
||||
let mut ni = 0;
|
||||
for j in 0..context_start {
|
||||
match &edits[j] {
|
||||
DiffLine::Context(_) => {
|
||||
oi += 1;
|
||||
ni += 1;
|
||||
}
|
||||
DiffLine::Removed(_) => oi += 1,
|
||||
DiffLine::Added(_) => ni += 1,
|
||||
}
|
||||
}
|
||||
let old_start = oi;
|
||||
let new_start = ni;
|
||||
|
||||
// Add context before
|
||||
let mut j = context_start;
|
||||
while j < i {
|
||||
hunk_lines.push(DiffLine::Context(match &edits[j] {
|
||||
DiffLine::Context(s) => s,
|
||||
_ => unreachable!(),
|
||||
}));
|
||||
j += 1;
|
||||
}
|
||||
|
||||
// Add changes and merge close hunks
|
||||
let mut consecutive_context = 0;
|
||||
while j < edits.len() {
|
||||
match &edits[j] {
|
||||
DiffLine::Context(s) => {
|
||||
consecutive_context += 1;
|
||||
if consecutive_context > context_lines * 2 {
|
||||
// End hunk, remove trailing context beyond limit
|
||||
let to_remove = consecutive_context - context_lines;
|
||||
for _ in 0..to_remove {
|
||||
hunk_lines.pop();
|
||||
}
|
||||
break;
|
||||
}
|
||||
hunk_lines.push(DiffLine::Context(s));
|
||||
}
|
||||
DiffLine::Added(s) => {
|
||||
consecutive_context = 0;
|
||||
hunk_lines.push(DiffLine::Added(s));
|
||||
}
|
||||
DiffLine::Removed(s) => {
|
||||
consecutive_context = 0;
|
||||
hunk_lines.push(DiffLine::Removed(s));
|
||||
}
|
||||
}
|
||||
j += 1;
|
||||
}
|
||||
|
||||
// Remove trailing context beyond limit
|
||||
while hunk_lines.len() > 0 {
|
||||
if matches!(hunk_lines.last(), Some(DiffLine::Context(_))) {
|
||||
let trailing_context = hunk_lines
|
||||
.iter()
|
||||
.rev()
|
||||
.take_while(|l| matches!(l, DiffLine::Context(_)))
|
||||
.count();
|
||||
if trailing_context > context_lines {
|
||||
hunk_lines.pop();
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let mut old_count = 0;
|
||||
let mut new_count = 0;
|
||||
for line in &hunk_lines {
|
||||
match line {
|
||||
DiffLine::Context(_) => {
|
||||
old_count += 1;
|
||||
new_count += 1;
|
||||
}
|
||||
DiffLine::Removed(_) => old_count += 1,
|
||||
DiffLine::Added(_) => new_count += 1,
|
||||
}
|
||||
}
|
||||
|
||||
if old_count > 0 || new_count > 0 {
|
||||
hunks.push(Hunk {
|
||||
old_start,
|
||||
old_count,
|
||||
new_start,
|
||||
new_count,
|
||||
lines: hunk_lines,
|
||||
});
|
||||
}
|
||||
|
||||
i = j;
|
||||
}
|
||||
|
||||
hunks
|
||||
}
|
||||
|
||||
fn compute_lcs<'a>(old: &[&'a str], new: &[&'a str]) -> Vec<(usize, usize)> {
|
||||
let n = old.len();
|
||||
let m = new.len();
|
||||
|
||||
if n == 0 || m == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// DP table (optimized for memory with two rows)
|
||||
let mut dp = vec![vec![0u32; m + 1]; n + 1];
|
||||
|
||||
for i in 1..=n {
|
||||
for j in 1..=m {
|
||||
if old[i - 1] == new[j - 1] {
|
||||
dp[i][j] = dp[i - 1][j - 1] + 1;
|
||||
} else {
|
||||
dp[i][j] = dp[i - 1][j].max(dp[i][j - 1]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Backtrack to find LCS
|
||||
let mut result = Vec::new();
|
||||
let mut i = n;
|
||||
let mut j = m;
|
||||
|
||||
while i > 0 && j > 0 {
|
||||
if old[i - 1] == new[j - 1] {
|
||||
result.push((i - 1, j - 1));
|
||||
i -= 1;
|
||||
j -= 1;
|
||||
} else if dp[i - 1][j] > dp[i][j - 1] {
|
||||
i -= 1;
|
||||
} else {
|
||||
j -= 1;
|
||||
}
|
||||
}
|
||||
|
||||
result.reverse();
|
||||
result
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
use crate::types::{
|
||||
CanUseToolFn, ContentBlock, Message, MessageRole, PermissionDecision, SDKMessage, Tool,
|
||||
ToolError, ToolResult, ToolResultContentBlock, ToolUseContext,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use super::registry::ToolRegistry;
|
||||
|
||||
/// Execute a set of tool calls from an assistant message.
|
||||
/// Concurrent-safe tools run in parallel; others run sequentially.
|
||||
pub async fn execute_tools(
|
||||
message: &Message,
|
||||
registry: &ToolRegistry,
|
||||
context: &ToolUseContext,
|
||||
permission_fn: Option<&CanUseToolFn>,
|
||||
event_sender: Option<mpsc::Sender<SDKMessage>>,
|
||||
) -> Vec<(String, String, ToolResult)> {
|
||||
let tool_uses: Vec<(String, String, Value)> = message
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::ToolUse { id, name, input } => {
|
||||
Some((id.clone(), name.clone(), input.clone()))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
if tool_uses.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// Partition into concurrent and sequential
|
||||
let mut concurrent_calls = Vec::new();
|
||||
let mut sequential_calls = Vec::new();
|
||||
|
||||
for (id, name, input) in &tool_uses {
|
||||
if let Some(tool) = registry.get(name) {
|
||||
if tool.is_concurrency_safe(input) {
|
||||
concurrent_calls.push((id.clone(), name.clone(), input.clone(), tool));
|
||||
} else {
|
||||
sequential_calls.push((id.clone(), name.clone(), input.clone(), tool));
|
||||
}
|
||||
} else {
|
||||
// Unknown tool
|
||||
sequential_calls.push((
|
||||
id.clone(),
|
||||
name.clone(),
|
||||
input.clone(),
|
||||
Arc::new(UnknownTool(name.clone())) as Arc<dyn Tool>,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let mut results = Vec::new();
|
||||
|
||||
// Run concurrent tools in parallel
|
||||
if !concurrent_calls.is_empty() {
|
||||
let mut handles = Vec::new();
|
||||
for (id, name, input, tool) in concurrent_calls {
|
||||
let ctx = context.clone();
|
||||
let perm_fn = permission_fn.cloned();
|
||||
let tool = tool.clone();
|
||||
let sender = event_sender.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
if let Some(sender) = &sender {
|
||||
let _ = sender
|
||||
.send(SDKMessage::ToolStart {
|
||||
tool_use_id: id.clone(),
|
||||
tool_name: name.clone(),
|
||||
input: input.clone(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
let input =
|
||||
check_permission(&name, input, perm_fn.as_ref(), &ctx.abort_signal).await;
|
||||
match input {
|
||||
Ok(input) => {
|
||||
let result = call_tool_with_cancel(tool.as_ref(), input, &ctx).await;
|
||||
let tool_result = match result {
|
||||
Ok(r) => r,
|
||||
Err(e) => ToolResult::error(e.to_string()),
|
||||
};
|
||||
(id, name, tool_result)
|
||||
}
|
||||
Err(msg) => (id, name, ToolResult::error(msg)),
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
if let Ok(result) = handle.await {
|
||||
results.push(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Run sequential tools one at a time
|
||||
for (id, name, input, tool) in sequential_calls {
|
||||
if let Some(sender) = &event_sender {
|
||||
let _ = sender
|
||||
.send(SDKMessage::ToolStart {
|
||||
tool_use_id: id.clone(),
|
||||
tool_name: name.clone(),
|
||||
input: input.clone(),
|
||||
})
|
||||
.await;
|
||||
}
|
||||
let input = check_permission(&name, input, permission_fn, &context.abort_signal).await;
|
||||
match input {
|
||||
Ok(input) => {
|
||||
let result = call_tool_with_cancel(tool.as_ref(), input, context).await;
|
||||
let tool_result = match result {
|
||||
Ok(r) => r,
|
||||
Err(e) => ToolResult::error(e.to_string()),
|
||||
};
|
||||
results.push((id, name, tool_result));
|
||||
}
|
||||
Err(msg) => {
|
||||
results.push((id, name, ToolResult::error(msg)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
async fn check_permission(
|
||||
tool_name: &str,
|
||||
input: Value,
|
||||
permission_fn: Option<&CanUseToolFn>,
|
||||
abort_signal: &tokio_util::sync::CancellationToken,
|
||||
) -> Result<Value, String> {
|
||||
if abort_signal.is_cancelled() {
|
||||
return Err("Tool aborted".to_string());
|
||||
}
|
||||
|
||||
if let Some(perm_fn) = permission_fn {
|
||||
let decision = tokio::select! {
|
||||
decision = perm_fn(tool_name, &input) => decision,
|
||||
_ = abort_signal.cancelled() => return Err("Tool aborted".to_string()),
|
||||
};
|
||||
match decision {
|
||||
PermissionDecision::Allow => Ok(input),
|
||||
PermissionDecision::Deny(msg) => Err(msg),
|
||||
PermissionDecision::AllowWithModifiedInput(new_input) => Ok(new_input),
|
||||
}
|
||||
} else {
|
||||
Ok(input)
|
||||
}
|
||||
}
|
||||
|
||||
async fn call_tool_with_cancel(
|
||||
tool: &dyn Tool,
|
||||
input: Value,
|
||||
context: &ToolUseContext,
|
||||
) -> Result<ToolResult, ToolError> {
|
||||
tokio::select! {
|
||||
result = tool.call(input, context) => result,
|
||||
_ = context.abort_signal.cancelled() => Ok(ToolResult::error("Tool aborted")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a user message containing tool results.
|
||||
pub fn build_tool_results_message(results: &[(String, String, ToolResult)]) -> Message {
|
||||
let content: Vec<ContentBlock> = results
|
||||
.iter()
|
||||
.map(|(id, _name, result)| {
|
||||
let content_blocks: Vec<ToolResultContentBlock> = result
|
||||
.content
|
||||
.iter()
|
||||
.map(|c| match c {
|
||||
crate::types::ToolResultContent::Text { text } => {
|
||||
ToolResultContentBlock::Text { text: text.clone() }
|
||||
}
|
||||
crate::types::ToolResultContent::Image { source } => {
|
||||
ToolResultContentBlock::Image {
|
||||
source: crate::types::ImageContentSource {
|
||||
source_type: source.source_type.clone(),
|
||||
media_type: source.media_type.clone(),
|
||||
data: source.data.clone(),
|
||||
},
|
||||
}
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
ContentBlock::ToolResult {
|
||||
tool_use_id: id.clone(),
|
||||
content: content_blocks,
|
||||
is_error: result.is_error,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
Message {
|
||||
role: MessageRole::User,
|
||||
content,
|
||||
}
|
||||
}
|
||||
|
||||
/// Placeholder tool for unknown tool names.
|
||||
struct UnknownTool(String);
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl Tool for UnknownTool {
|
||||
fn name(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Unknown tool"
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> crate::types::ToolInputSchema {
|
||||
crate::types::ToolInputSchema::default()
|
||||
}
|
||||
|
||||
async fn call(
|
||||
&self,
|
||||
_input: Value,
|
||||
_context: &ToolUseContext,
|
||||
) -> Result<ToolResult, ToolError> {
|
||||
Ok(ToolResult::error(format!(
|
||||
"Unknown tool: {}. Use ToolSearch to discover available tools.",
|
||||
self.0
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
pub struct FileEditTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for FileEditTool {
|
||||
fn name(&self) -> &str {
|
||||
"Edit"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Performs exact string replacements in files. The old_string must be unique in the file unless replace_all is true."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"file_path".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The absolute path to the file to modify"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"old_string".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The text to replace"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"new_string".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The replacement text"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"replace_all".to_string(),
|
||||
json!({
|
||||
"type": "boolean",
|
||||
"description": "Replace all occurrences (default false)",
|
||||
"default": false
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec![
|
||||
"file_path".to_string(),
|
||||
"old_string".to_string(),
|
||||
"new_string".to_string(),
|
||||
],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let file_path = input
|
||||
.get("file_path")
|
||||
.and_then(|p| p.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'file_path' field".to_string()))?;
|
||||
|
||||
let old_string = input
|
||||
.get("old_string")
|
||||
.and_then(|s| s.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'old_string' field".to_string()))?;
|
||||
|
||||
let new_string = input
|
||||
.get("new_string")
|
||||
.and_then(|s| s.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'new_string' field".to_string()))?;
|
||||
|
||||
let replace_all = input
|
||||
.get("replace_all")
|
||||
.and_then(|r| r.as_bool())
|
||||
.unwrap_or(false);
|
||||
|
||||
if old_string == new_string {
|
||||
return Ok(ToolResult::error(
|
||||
"old_string and new_string must be different".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
let path = if Path::new(file_path).is_absolute() {
|
||||
std::path::PathBuf::from(file_path)
|
||||
} else {
|
||||
Path::new(&context.working_dir).join(file_path)
|
||||
};
|
||||
|
||||
if !path.exists() {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"File not found: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
|
||||
let content = tokio::fs::read_to_string(&path)
|
||||
.await
|
||||
.map_err(ToolError::IoError)?;
|
||||
|
||||
// Check staleness
|
||||
{
|
||||
let current_hash = simple_hash(&content);
|
||||
let state = context.read_file_state.read().await;
|
||||
if let Some(saved_hash) = state.get(&path.to_string_lossy().to_string()) {
|
||||
if *saved_hash != current_hash {
|
||||
return Ok(ToolResult::error(
|
||||
"File has been modified since last read. Please Read the file again before editing."
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Count occurrences
|
||||
let count = content.matches(old_string).count();
|
||||
|
||||
if count == 0 {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"old_string not found in {}. Make sure the string matches exactly, including whitespace and indentation.",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
|
||||
if count > 1 && !replace_all {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"old_string found {} times in {}. Either provide more context to make it unique, or set replace_all to true.",
|
||||
count,
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
|
||||
let new_content = if replace_all {
|
||||
content.replace(old_string, new_string)
|
||||
} else {
|
||||
content.replacen(old_string, new_string, 1)
|
||||
};
|
||||
|
||||
let diff = super::diff::unified_diff(&content, &new_content, file_path);
|
||||
|
||||
tokio::fs::write(&path, &new_content)
|
||||
.await
|
||||
.map_err(ToolError::IoError)?;
|
||||
|
||||
// Update file state
|
||||
let hash = simple_hash(&new_content);
|
||||
let mut state = context.read_file_state.write().await;
|
||||
state.insert(path.to_string_lossy().to_string(), hash);
|
||||
|
||||
let replacements = if replace_all {
|
||||
format!(" ({} replacements)", count)
|
||||
} else {
|
||||
String::new()
|
||||
};
|
||||
|
||||
Ok(ToolResult::text(format!(
|
||||
"Edited file: {}{}\n\n{}",
|
||||
path.display(),
|
||||
replacements,
|
||||
diff
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
fn simple_hash(content: &str) -> String {
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = DefaultHasher::new();
|
||||
content.hash(&mut hasher);
|
||||
format!("{:x}", hasher.finish())
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
use async_trait::async_trait;
|
||||
use base64::Engine;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
const DEFAULT_LINE_LIMIT: usize = 2000;
|
||||
const MAX_CONTENT_SIZE: usize = 100_000;
|
||||
|
||||
pub struct FileReadTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for FileReadTool {
|
||||
fn name(&self) -> &str {
|
||||
"Read"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Reads a file from the local filesystem. Returns content with line numbers."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"file_path".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The absolute path to the file to read"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"offset".to_string(),
|
||||
json!({
|
||||
"type": "integer",
|
||||
"description": "The line number to start reading from (0-based)"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"limit".to_string(),
|
||||
json!({
|
||||
"type": "integer",
|
||||
"description": "The number of lines to read"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["file_path".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let file_path = input
|
||||
.get("file_path")
|
||||
.and_then(|p| p.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'file_path' field".to_string()))?;
|
||||
|
||||
let path = resolve_path(file_path, &context.working_dir);
|
||||
|
||||
// Check if file is a device file
|
||||
if file_path.starts_with("/dev/") || file_path.starts_with("/proc/") {
|
||||
return Ok(ToolResult::error(
|
||||
"Cannot read device or proc files".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
if !path.exists() {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"File not found: {}",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
|
||||
if path.is_dir() {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"{} is a directory, not a file. Use Bash with 'ls' to list directory contents.",
|
||||
path.display()
|
||||
)));
|
||||
}
|
||||
|
||||
// Check for image files
|
||||
let ext = path
|
||||
.extension()
|
||||
.and_then(|e| e.to_str())
|
||||
.unwrap_or("")
|
||||
.to_lowercase();
|
||||
|
||||
if matches!(ext.as_str(), "png" | "jpg" | "jpeg" | "gif" | "webp" | "bmp" | "svg") {
|
||||
return read_image(&path, &ext).await;
|
||||
}
|
||||
|
||||
// Read text file
|
||||
let content = tokio::fs::read_to_string(&path)
|
||||
.await
|
||||
.map_err(|e| ToolError::IoError(e))?;
|
||||
|
||||
// Track file state for staleness detection
|
||||
{
|
||||
let hash = simple_hash(&content);
|
||||
let mut state = context.read_file_state.write().await;
|
||||
state.insert(path.to_string_lossy().to_string(), hash);
|
||||
}
|
||||
|
||||
let offset = input.get("offset").and_then(|o| o.as_u64()).unwrap_or(0) as usize;
|
||||
let limit = input
|
||||
.get("limit")
|
||||
.and_then(|l| l.as_u64())
|
||||
.unwrap_or(DEFAULT_LINE_LIMIT as u64) as usize;
|
||||
|
||||
let lines: Vec<&str> = content.lines().collect();
|
||||
let total_lines = lines.len();
|
||||
|
||||
let start = offset.min(total_lines);
|
||||
let end = (start + limit).min(total_lines);
|
||||
let selected_lines = &lines[start..end];
|
||||
|
||||
let mut result = String::new();
|
||||
for (i, line) in selected_lines.iter().enumerate() {
|
||||
let line_num = start + i + 1;
|
||||
result.push_str(&format!("{}\t{}\n", line_num, line));
|
||||
}
|
||||
|
||||
if result.len() > MAX_CONTENT_SIZE {
|
||||
result.truncate(MAX_CONTENT_SIZE);
|
||||
result.push_str("\n... (content truncated)");
|
||||
}
|
||||
|
||||
if end < total_lines {
|
||||
result.push_str(&format!(
|
||||
"\n(showing lines {}-{} of {})",
|
||||
start + 1,
|
||||
end,
|
||||
total_lines
|
||||
));
|
||||
}
|
||||
|
||||
Ok(ToolResult::text(result))
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_image(path: &Path, ext: &str) -> Result<ToolResult, ToolError> {
|
||||
let data = tokio::fs::read(path).await.map_err(ToolError::IoError)?;
|
||||
let encoded = base64::engine::general_purpose::STANDARD.encode(&data);
|
||||
let media_type = match ext {
|
||||
"png" => "image/png",
|
||||
"jpg" | "jpeg" => "image/jpeg",
|
||||
"gif" => "image/gif",
|
||||
"webp" => "image/webp",
|
||||
"bmp" => "image/bmp",
|
||||
"svg" => "image/svg+xml",
|
||||
_ => "application/octet-stream",
|
||||
};
|
||||
Ok(ToolResult::image(encoded, media_type.to_string()))
|
||||
}
|
||||
|
||||
fn resolve_path(file_path: &str, working_dir: &str) -> PathBuf {
|
||||
let path = Path::new(file_path);
|
||||
if path.is_absolute() {
|
||||
path.to_path_buf()
|
||||
} else {
|
||||
Path::new(working_dir).join(file_path)
|
||||
}
|
||||
}
|
||||
|
||||
fn simple_hash(content: &str) -> String {
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = DefaultHasher::new();
|
||||
content.hash(&mut hasher);
|
||||
format!("{:x}", hasher.finish())
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
pub struct FileWriteTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for FileWriteTool {
|
||||
fn name(&self) -> &str {
|
||||
"Write"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Writes a file to the local filesystem. Will overwrite existing files. Use Read tool first for existing files."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"file_path".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The absolute path to the file to write"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"content".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The content to write to the file"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["file_path".to_string(), "content".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let file_path = input
|
||||
.get("file_path")
|
||||
.and_then(|p| p.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'file_path' field".to_string()))?;
|
||||
|
||||
let content = input
|
||||
.get("content")
|
||||
.and_then(|c| c.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'content' field".to_string()))?;
|
||||
|
||||
let path = if Path::new(file_path).is_absolute() {
|
||||
std::path::PathBuf::from(file_path)
|
||||
} else {
|
||||
Path::new(&context.working_dir).join(file_path)
|
||||
};
|
||||
|
||||
let is_new = !path.exists();
|
||||
|
||||
// Create parent directories if needed
|
||||
if let Some(parent) = path.parent() {
|
||||
tokio::fs::create_dir_all(parent)
|
||||
.await
|
||||
.map_err(ToolError::IoError)?;
|
||||
}
|
||||
|
||||
// Check staleness for existing files
|
||||
if !is_new {
|
||||
let current_content = tokio::fs::read_to_string(&path).await.ok();
|
||||
if let Some(current) = ¤t_content {
|
||||
let current_hash = simple_hash(current);
|
||||
let state = context.read_file_state.read().await;
|
||||
if let Some(saved_hash) = state.get(&path.to_string_lossy().to_string()) {
|
||||
if *saved_hash != current_hash {
|
||||
return Ok(ToolResult::error(
|
||||
"File has been modified since last read. Please Read the file again before writing."
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Generate diff for existing files
|
||||
if let Some(old_content) = current_content {
|
||||
let diff = super::diff::unified_diff(&old_content, content, file_path);
|
||||
tokio::fs::write(&path, content)
|
||||
.await
|
||||
.map_err(ToolError::IoError)?;
|
||||
|
||||
// Update file state
|
||||
let hash = simple_hash(content);
|
||||
let mut state = context.read_file_state.write().await;
|
||||
state.insert(path.to_string_lossy().to_string(), hash);
|
||||
|
||||
return Ok(ToolResult::text(format!(
|
||||
"Updated file: {}\n\n{}",
|
||||
path.display(),
|
||||
diff
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
tokio::fs::write(&path, content)
|
||||
.await
|
||||
.map_err(ToolError::IoError)?;
|
||||
|
||||
// Update file state
|
||||
let hash = simple_hash(content);
|
||||
let mut state = context.read_file_state.write().await;
|
||||
state.insert(path.to_string_lossy().to_string(), hash);
|
||||
|
||||
let action = if is_new { "Created" } else { "Updated" };
|
||||
Ok(ToolResult::text(format!(
|
||||
"{} file: {} ({} bytes)",
|
||||
action,
|
||||
path.display(),
|
||||
content.len()
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
fn simple_hash(content: &str) -> String {
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
let mut hasher = DefaultHasher::new();
|
||||
content.hash(&mut hasher);
|
||||
format!("{:x}", hasher.finish())
|
||||
}
|
||||
@@ -0,0 +1,274 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
const MAX_RESULTS: usize = 100;
|
||||
const MAX_VISITED_ENTRIES: usize = 20_000;
|
||||
const MAX_SEARCH_MS: u64 = 3_000;
|
||||
|
||||
/// Directories to skip during glob matching.
|
||||
const SKIP_DIRS: &[&str] = &[
|
||||
".git",
|
||||
"node_modules",
|
||||
".next",
|
||||
"__pycache__",
|
||||
".mypy_cache",
|
||||
"target",
|
||||
"dist",
|
||||
"build",
|
||||
];
|
||||
|
||||
pub struct GlobTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for GlobTool {
|
||||
fn name(&self) -> &str {
|
||||
"Glob"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Fast file pattern matching tool. Supports glob patterns like \"**/*.rs\" or \"src/**/*.ts\". Returns matching file paths sorted by modification time."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"pattern".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The glob pattern to match files against"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"path".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The directory to search in (defaults to working directory)"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["pattern".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let pattern = input
|
||||
.get("pattern")
|
||||
.and_then(|p| p.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'pattern' field".to_string()))?;
|
||||
|
||||
let search_path_value = input
|
||||
.get("path")
|
||||
.and_then(|p| p.as_str())
|
||||
.unwrap_or(&context.working_dir);
|
||||
let search_path = resolve_search_path(search_path_value, &context.working_dir);
|
||||
let pattern = normalize_pattern(pattern)?;
|
||||
let matcher = glob::Pattern::new(&pattern)
|
||||
.map_err(|e| ToolError::ExecutionError(format!("Invalid glob pattern: {}", e)))?;
|
||||
let (walk_root, max_depth) = walk_root_and_depth(&search_path, &pattern);
|
||||
|
||||
let scan = scan_files(
|
||||
&walk_root,
|
||||
&search_path,
|
||||
&matcher,
|
||||
pattern.starts_with('/'),
|
||||
max_depth,
|
||||
context,
|
||||
)?;
|
||||
let mut files = scan.files;
|
||||
|
||||
// Sort by modification time (newest first)
|
||||
files.sort_by(|a, b| b.1.cmp(&a.1));
|
||||
|
||||
let total = files.len();
|
||||
let truncated = scan.truncated || total > MAX_RESULTS;
|
||||
let files: Vec<String> = files
|
||||
.into_iter()
|
||||
.take(MAX_RESULTS)
|
||||
.map(|(p, _)| p)
|
||||
.collect();
|
||||
|
||||
if files.is_empty() {
|
||||
return Ok(ToolResult::text(format!(
|
||||
"No files found matching pattern: {}",
|
||||
pattern
|
||||
)));
|
||||
}
|
||||
|
||||
let mut result = files.join("\n");
|
||||
if truncated {
|
||||
result.push_str(&format!(
|
||||
"\n\n(showing first {} matches; search was bounded)",
|
||||
files.len().min(MAX_RESULTS)
|
||||
));
|
||||
}
|
||||
|
||||
Ok(ToolResult::text(result))
|
||||
}
|
||||
}
|
||||
|
||||
struct GlobScan {
|
||||
files: Vec<(String, std::time::SystemTime)>,
|
||||
truncated: bool,
|
||||
}
|
||||
|
||||
fn normalize_pattern(pattern: &str) -> Result<String, ToolError> {
|
||||
let trimmed = pattern.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Err(ToolError::InvalidInput("Empty glob pattern".to_string()));
|
||||
}
|
||||
Ok(match trimmed {
|
||||
"." | "./" => "*".to_string(),
|
||||
other => other.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
fn resolve_search_path(input_path: &str, working_dir: &str) -> PathBuf {
|
||||
let path = Path::new(input_path);
|
||||
if path.is_absolute() {
|
||||
path.to_path_buf()
|
||||
} else {
|
||||
Path::new(working_dir).join(path)
|
||||
}
|
||||
}
|
||||
|
||||
fn walk_root_and_depth(search_path: &Path, pattern: &str) -> (PathBuf, Option<usize>) {
|
||||
let base = extract_glob_base(pattern);
|
||||
let root = if pattern.starts_with('/') {
|
||||
PathBuf::from(if base.is_empty() { "/" } else { &base })
|
||||
} else if base.is_empty() {
|
||||
search_path.to_path_buf()
|
||||
} else {
|
||||
search_path.join(base)
|
||||
};
|
||||
let remaining = remaining_segments(pattern, &extract_glob_base(pattern));
|
||||
let max_depth = if remaining.iter().any(|part| *part == "**") {
|
||||
None
|
||||
} else {
|
||||
Some(remaining.len().max(1))
|
||||
};
|
||||
(root, max_depth)
|
||||
}
|
||||
|
||||
fn extract_glob_base(pattern: &str) -> String {
|
||||
let mut parts = Vec::new();
|
||||
for segment in pattern.split('/') {
|
||||
if segment.contains('*')
|
||||
|| segment.contains('?')
|
||||
|| segment.contains('[')
|
||||
|| segment.contains('{')
|
||||
{
|
||||
break;
|
||||
}
|
||||
parts.push(segment);
|
||||
}
|
||||
parts.join("/")
|
||||
}
|
||||
|
||||
fn remaining_segments<'a>(pattern: &'a str, base: &str) -> Vec<&'a str> {
|
||||
let base_count = base.split('/').filter(|part| !part.is_empty()).count();
|
||||
pattern
|
||||
.split('/')
|
||||
.filter(|part| !part.is_empty())
|
||||
.skip(base_count)
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn scan_files(
|
||||
root: &Path,
|
||||
search_path: &Path,
|
||||
matcher: &glob::Pattern,
|
||||
match_absolute: bool,
|
||||
max_depth: Option<usize>,
|
||||
context: &ToolUseContext,
|
||||
) -> Result<GlobScan, ToolError> {
|
||||
let deadline = Duration::from_millis(MAX_SEARCH_MS);
|
||||
let started = Instant::now();
|
||||
let mut stack = vec![(root.to_path_buf(), 0usize)];
|
||||
let mut files = Vec::new();
|
||||
let mut visited = 0usize;
|
||||
let mut truncated = false;
|
||||
|
||||
if root.is_file() && path_matches(root, search_path, matcher, match_absolute) {
|
||||
let modified = root
|
||||
.metadata()
|
||||
.and_then(|m| m.modified())
|
||||
.unwrap_or(std::time::UNIX_EPOCH);
|
||||
files.push((root.to_string_lossy().to_string(), modified));
|
||||
return Ok(GlobScan { files, truncated });
|
||||
}
|
||||
|
||||
while let Some((dir, depth)) = stack.pop() {
|
||||
if context.abort_signal.is_cancelled() {
|
||||
return Ok(GlobScan {
|
||||
files,
|
||||
truncated: true,
|
||||
});
|
||||
}
|
||||
if visited >= MAX_VISITED_ENTRIES || started.elapsed() > deadline {
|
||||
truncated = true;
|
||||
break;
|
||||
}
|
||||
let entries = match std::fs::read_dir(&dir) {
|
||||
Ok(entries) => entries,
|
||||
Err(_) => continue,
|
||||
};
|
||||
for entry in entries.flatten() {
|
||||
visited += 1;
|
||||
let path = entry.path();
|
||||
if path.is_dir() {
|
||||
if should_skip_dir(&path) {
|
||||
continue;
|
||||
}
|
||||
if max_depth.map_or(true, |limit| depth + 1 < limit) {
|
||||
stack.push((path, depth + 1));
|
||||
}
|
||||
} else if path.is_file() && path_matches(&path, search_path, matcher, match_absolute) {
|
||||
let modified = entry
|
||||
.metadata()
|
||||
.and_then(|m| m.modified())
|
||||
.unwrap_or(std::time::UNIX_EPOCH);
|
||||
files.push((path.to_string_lossy().to_string(), modified));
|
||||
if files.len() >= MAX_RESULTS {
|
||||
truncated = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
if truncated {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(GlobScan { files, truncated })
|
||||
}
|
||||
|
||||
fn path_matches(path: &Path, search_path: &Path, matcher: &glob::Pattern, absolute: bool) -> bool {
|
||||
if absolute {
|
||||
return matcher.matches_path(path);
|
||||
}
|
||||
path.strip_prefix(search_path)
|
||||
.ok()
|
||||
.is_some_and(|relative| matcher.matches_path(relative))
|
||||
}
|
||||
|
||||
fn should_skip_dir(path: &Path) -> bool {
|
||||
path.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.is_some_and(|name| SKIP_DIRS.contains(&name))
|
||||
}
|
||||
@@ -0,0 +1,290 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::process::Stdio;
|
||||
use tokio::process::Command;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
const DEFAULT_HEAD_LIMIT: usize = 250;
|
||||
|
||||
pub struct GrepTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for GrepTool {
|
||||
fn name(&self) -> &str {
|
||||
"Grep"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"A powerful search tool built on ripgrep. Supports full regex syntax, file type filtering, and multiple output modes."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"pattern".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The regex pattern to search for"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"path".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "File or directory to search in"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"output_mode".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["content", "files_with_matches", "count"],
|
||||
"description": "Output mode (default: files_with_matches)"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"glob".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "Glob pattern to filter files (e.g. \"*.rs\")"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"type".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "File type to search (e.g. rust, js, py)"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"-i".to_string(),
|
||||
json!({
|
||||
"type": "boolean",
|
||||
"description": "Case insensitive search"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"-n".to_string(),
|
||||
json!({
|
||||
"type": "boolean",
|
||||
"description": "Show line numbers"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"-A".to_string(),
|
||||
json!({
|
||||
"type": "number",
|
||||
"description": "Lines to show after each match"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"-B".to_string(),
|
||||
json!({
|
||||
"type": "number",
|
||||
"description": "Lines to show before each match"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"-C".to_string(),
|
||||
json!({
|
||||
"type": "number",
|
||||
"description": "Lines of context around each match"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"head_limit".to_string(),
|
||||
json!({
|
||||
"type": "number",
|
||||
"description": "Limit output to first N entries (default 250)"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"multiline".to_string(),
|
||||
json!({
|
||||
"type": "boolean",
|
||||
"description": "Enable multiline matching"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["pattern".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let pattern = input
|
||||
.get("pattern")
|
||||
.and_then(|p| p.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'pattern' field".to_string()))?;
|
||||
|
||||
let search_path = input
|
||||
.get("path")
|
||||
.and_then(|p| p.as_str())
|
||||
.unwrap_or(&context.working_dir);
|
||||
|
||||
let output_mode = input
|
||||
.get("output_mode")
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or("files_with_matches");
|
||||
|
||||
let head_limit = input
|
||||
.get("head_limit")
|
||||
.and_then(|h| h.as_u64())
|
||||
.unwrap_or(DEFAULT_HEAD_LIMIT as u64) as usize;
|
||||
|
||||
// Try rg first, fall back to grep
|
||||
let result = run_ripgrep(pattern, search_path, output_mode, &input).await;
|
||||
|
||||
match result {
|
||||
Ok(output) => {
|
||||
if output.is_empty() {
|
||||
return Ok(ToolResult::text("No matches found.".to_string()));
|
||||
}
|
||||
|
||||
let lines: Vec<&str> = output.lines().collect();
|
||||
let total = lines.len();
|
||||
let truncated = total > head_limit;
|
||||
let lines: Vec<&str> = lines.into_iter().take(head_limit).collect();
|
||||
let mut result = lines.join("\n");
|
||||
|
||||
if truncated {
|
||||
result.push_str(&format!(
|
||||
"\n\n(showing {} of {} results)",
|
||||
head_limit, total
|
||||
));
|
||||
}
|
||||
|
||||
Ok(ToolResult::text(result))
|
||||
}
|
||||
Err(e) => {
|
||||
// Fall back to grep
|
||||
match run_grep(pattern, search_path, output_mode, &input).await {
|
||||
Ok(output) => {
|
||||
if output.is_empty() {
|
||||
return Ok(ToolResult::text("No matches found.".to_string()));
|
||||
}
|
||||
Ok(ToolResult::text(output))
|
||||
}
|
||||
Err(_) => Ok(ToolResult::error(format!("Search failed: {}", e))),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_ripgrep(
|
||||
pattern: &str,
|
||||
path: &str,
|
||||
output_mode: &str,
|
||||
input: &Value,
|
||||
) -> Result<String, String> {
|
||||
let mut args = vec!["--no-heading".to_string()];
|
||||
|
||||
match output_mode {
|
||||
"files_with_matches" => args.push("-l".to_string()),
|
||||
"count" => args.push("-c".to_string()),
|
||||
"content" => {
|
||||
if input.get("-n").and_then(|n| n.as_bool()).unwrap_or(true) {
|
||||
args.push("-n".to_string());
|
||||
}
|
||||
}
|
||||
_ => args.push("-l".to_string()),
|
||||
}
|
||||
|
||||
if input.get("-i").and_then(|i| i.as_bool()).unwrap_or(false) {
|
||||
args.push("-i".to_string());
|
||||
}
|
||||
|
||||
if let Some(after) = input.get("-A").and_then(|a| a.as_u64()) {
|
||||
args.push(format!("-A{}", after));
|
||||
}
|
||||
if let Some(before) = input.get("-B").and_then(|b| b.as_u64()) {
|
||||
args.push(format!("-B{}", before));
|
||||
}
|
||||
if let Some(ctx) = input.get("-C").and_then(|c| c.as_u64()) {
|
||||
args.push(format!("-C{}", ctx));
|
||||
}
|
||||
|
||||
if let Some(glob_pattern) = input.get("glob").and_then(|g| g.as_str()) {
|
||||
args.push("--glob".to_string());
|
||||
args.push(glob_pattern.to_string());
|
||||
}
|
||||
|
||||
if let Some(file_type) = input.get("type").and_then(|t| t.as_str()) {
|
||||
args.push("--type".to_string());
|
||||
args.push(file_type.to_string());
|
||||
}
|
||||
|
||||
if input
|
||||
.get("multiline")
|
||||
.and_then(|m| m.as_bool())
|
||||
.unwrap_or(false)
|
||||
{
|
||||
args.push("-U".to_string());
|
||||
args.push("--multiline-dotall".to_string());
|
||||
}
|
||||
|
||||
args.push(pattern.to_string());
|
||||
args.push(path.to_string());
|
||||
|
||||
let output = Command::new("rg")
|
||||
.args(&args)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
if output.status.success() || output.status.code() == Some(1) {
|
||||
Ok(String::from_utf8_lossy(&output.stdout).to_string())
|
||||
} else {
|
||||
Err(String::from_utf8_lossy(&output.stderr).to_string())
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_grep(
|
||||
pattern: &str,
|
||||
path: &str,
|
||||
output_mode: &str,
|
||||
input: &Value,
|
||||
) -> Result<String, String> {
|
||||
let mut args = vec!["-r".to_string()];
|
||||
|
||||
match output_mode {
|
||||
"files_with_matches" => args.push("-l".to_string()),
|
||||
"count" => args.push("-c".to_string()),
|
||||
"content" => {
|
||||
args.push("-n".to_string());
|
||||
}
|
||||
_ => args.push("-l".to_string()),
|
||||
}
|
||||
|
||||
if input.get("-i").and_then(|i| i.as_bool()).unwrap_or(false) {
|
||||
args.push("-i".to_string());
|
||||
}
|
||||
|
||||
args.push(pattern.to_string());
|
||||
args.push(path.to_string());
|
||||
|
||||
let output = Command::new("grep")
|
||||
.args(&args)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.output()
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
|
||||
Ok(String::from_utf8_lossy(&output.stdout).to_string())
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
use std::process::Command;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// LSPTool - Language Server Protocol integration.
|
||||
///
|
||||
/// Provides code intelligence: go-to-definition, find-references,
|
||||
/// hover, document symbols. Uses ripgrep/grep as fallback for symbol lookup.
|
||||
pub struct LSPTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for LSPTool {
|
||||
fn name(&self) -> &str {
|
||||
"LSP"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Language Server Protocol operations for code intelligence. Supports go-to-definition, find-references, hover, and symbol lookup."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"operation".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["goToDefinition", "findReferences", "hover", "documentSymbol"],
|
||||
"description": "LSP operation to perform"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"file_path".to_string(),
|
||||
json!({ "type": "string", "description": "File path for the operation" }),
|
||||
),
|
||||
(
|
||||
"line".to_string(),
|
||||
json!({ "type": "number", "description": "Line number (0-based)" }),
|
||||
),
|
||||
(
|
||||
"character".to_string(),
|
||||
json!({ "type": "number", "description": "Character position (0-based)" }),
|
||||
),
|
||||
(
|
||||
"query".to_string(),
|
||||
json!({ "type": "string", "description": "Symbol name (for workspace symbol search)" }),
|
||||
),
|
||||
]),
|
||||
required: vec!["operation".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let operation = input
|
||||
.get("operation")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'operation'".to_string()))?;
|
||||
|
||||
let file_path = input.get("file_path").and_then(|v| v.as_str());
|
||||
let line = input.get("line").and_then(|v| v.as_u64());
|
||||
let character = input.get("character").and_then(|v| v.as_u64()).unwrap_or(0);
|
||||
|
||||
let cwd = &context.working_dir;
|
||||
|
||||
match operation {
|
||||
"goToDefinition" => {
|
||||
let file_path = file_path
|
||||
.ok_or_else(|| ToolError::InvalidInput("file_path required".to_string()))?;
|
||||
let line = line
|
||||
.ok_or_else(|| ToolError::InvalidInput("line required".to_string()))?;
|
||||
|
||||
let symbol = get_symbol_at_position(file_path, line as usize, character as usize, cwd);
|
||||
match symbol {
|
||||
Some(sym) => {
|
||||
let pattern = format!(
|
||||
r"(?:fn |struct |enum |trait |type |const |let |pub fn |pub struct |pub enum |pub trait |pub type |pub const |impl |mod |use )\s*{}",
|
||||
regex::escape(&sym)
|
||||
);
|
||||
let result = run_rg_or_grep(&pattern, cwd);
|
||||
if result.is_empty() {
|
||||
Ok(ToolResult::text(format!("No definition found for \"{}\"", sym)))
|
||||
} else {
|
||||
Ok(ToolResult::text(result))
|
||||
}
|
||||
}
|
||||
None => Ok(ToolResult::text("Could not identify symbol at position")),
|
||||
}
|
||||
}
|
||||
|
||||
"findReferences" => {
|
||||
let file_path = file_path
|
||||
.ok_or_else(|| ToolError::InvalidInput("file_path required".to_string()))?;
|
||||
let line = line
|
||||
.ok_or_else(|| ToolError::InvalidInput("line required".to_string()))?;
|
||||
|
||||
let symbol = get_symbol_at_position(file_path, line as usize, character as usize, cwd);
|
||||
match symbol {
|
||||
Some(sym) => {
|
||||
let result = run_rg_or_grep(®ex::escape(&sym), cwd);
|
||||
if result.is_empty() {
|
||||
Ok(ToolResult::text(format!("No references found for \"{}\"", sym)))
|
||||
} else {
|
||||
Ok(ToolResult::text(result))
|
||||
}
|
||||
}
|
||||
None => Ok(ToolResult::text("Could not identify symbol at position")),
|
||||
}
|
||||
}
|
||||
|
||||
"hover" => Ok(ToolResult::text(
|
||||
"Hover information requires a running language server. Use FileRead tool to examine the file content.",
|
||||
)),
|
||||
|
||||
"documentSymbol" => {
|
||||
let file_path = file_path
|
||||
.ok_or_else(|| ToolError::InvalidInput("file_path required".to_string()))?;
|
||||
|
||||
let pattern = r"^\s*(pub\s+)?(fn |struct |enum |trait |type |const |let |impl |mod |use )";
|
||||
let result = run_rg_on_file(pattern, file_path, cwd);
|
||||
if result.is_empty() {
|
||||
Ok(ToolResult::text("No symbols found"))
|
||||
} else {
|
||||
Ok(ToolResult::text(result))
|
||||
}
|
||||
}
|
||||
|
||||
_ => Ok(ToolResult::text(format!(
|
||||
"LSP operation \"{}\" requires a running language server.",
|
||||
operation
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract the symbol (word) at a given line and character position in a file.
|
||||
fn get_symbol_at_position(
|
||||
file_path: &str,
|
||||
line: usize,
|
||||
character: usize,
|
||||
cwd: &str,
|
||||
) -> Option<String> {
|
||||
let full_path = Path::new(cwd).join(file_path);
|
||||
let content = std::fs::read_to_string(full_path).ok()?;
|
||||
let lines: Vec<&str> = content.lines().collect();
|
||||
|
||||
if line >= lines.len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let line_text = lines[line];
|
||||
if character >= line_text.len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
// Find the word boundary around the character position
|
||||
let bytes = line_text.as_bytes();
|
||||
let mut start = character;
|
||||
let mut end = character;
|
||||
|
||||
while start > 0 && is_word_char(bytes[start - 1]) {
|
||||
start -= 1;
|
||||
}
|
||||
while end < bytes.len() && is_word_char(bytes[end]) {
|
||||
end += 1;
|
||||
}
|
||||
|
||||
if start == end {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(line_text[start..end].to_string())
|
||||
}
|
||||
|
||||
fn is_word_char(b: u8) -> bool {
|
||||
b.is_ascii_alphanumeric() || b == b'_'
|
||||
}
|
||||
|
||||
/// Run ripgrep on the working directory, falling back to grep.
|
||||
fn run_rg_or_grep(pattern: &str, cwd: &str) -> String {
|
||||
// Try ripgrep first
|
||||
if let Ok(output) = Command::new("rg")
|
||||
.args(["-n", pattern, "--type-add", "src:*.{ts,tsx,js,jsx,py,go,rs,java}", "-t", "src"])
|
||||
.arg(cwd)
|
||||
.output()
|
||||
{
|
||||
if output.status.success() {
|
||||
let text = String::from_utf8_lossy(&output.stdout);
|
||||
// Limit output to 50 lines
|
||||
let lines: Vec<&str> = text.lines().take(50).collect();
|
||||
return lines.join("\n");
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to grep
|
||||
if let Ok(output) = Command::new("grep")
|
||||
.args(["-rn", pattern, cwd, "--include=*.rs", "--include=*.ts", "--include=*.py", "--include=*.go", "--include=*.java"])
|
||||
.output()
|
||||
{
|
||||
if output.status.success() {
|
||||
let text = String::from_utf8_lossy(&output.stdout);
|
||||
let lines: Vec<&str> = text.lines().take(50).collect();
|
||||
return lines.join("\n");
|
||||
}
|
||||
}
|
||||
|
||||
String::new()
|
||||
}
|
||||
|
||||
/// Run ripgrep on a single file, falling back to grep.
|
||||
fn run_rg_on_file(pattern: &str, file_path: &str, cwd: &str) -> String {
|
||||
let full_path = Path::new(cwd).join(file_path);
|
||||
let path_str = full_path.to_string_lossy();
|
||||
|
||||
if let Ok(output) = Command::new("rg")
|
||||
.args(["-n", pattern])
|
||||
.arg(path_str.as_ref())
|
||||
.output()
|
||||
{
|
||||
if output.status.success() {
|
||||
return String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(output) = Command::new("grep")
|
||||
.args(["-n", pattern])
|
||||
.arg(path_str.as_ref())
|
||||
.output()
|
||||
{
|
||||
if output.status.success() {
|
||||
return String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
}
|
||||
}
|
||||
|
||||
String::new()
|
||||
}
|
||||
@@ -0,0 +1,174 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::mcp::McpClient;
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// ListMcpResourcesTool - List available resources from connected MCP servers.
|
||||
pub struct ListMcpResourcesTool {
|
||||
mcp_client: Arc<RwLock<Option<Arc<McpClient>>>>,
|
||||
}
|
||||
|
||||
impl ListMcpResourcesTool {
|
||||
pub fn new(mcp_client: Arc<RwLock<Option<Arc<McpClient>>>>) -> Self {
|
||||
Self { mcp_client }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ListMcpResourcesTool {
|
||||
fn name(&self) -> &str {
|
||||
"ListMcpResources"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"List available resources from connected MCP servers. Resources can include files, databases, and other data sources."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([(
|
||||
"server".to_string(),
|
||||
json!({ "type": "string", "description": "Filter by MCP server name" }),
|
||||
)]),
|
||||
required: vec![],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let client_guard = self.mcp_client.read().await;
|
||||
let client = match client_guard.as_ref() {
|
||||
Some(c) => c.clone(),
|
||||
None => return Ok(ToolResult::text("No MCP servers connected.")),
|
||||
};
|
||||
|
||||
let server_filter = input.get("server").and_then(|v| v.as_str());
|
||||
|
||||
let connections = client.get_connections().await;
|
||||
|
||||
let filtered: Vec<_> = if let Some(server) = server_filter {
|
||||
connections.into_iter().filter(|c| c.name == server).collect()
|
||||
} else {
|
||||
connections
|
||||
};
|
||||
|
||||
if filtered.is_empty() {
|
||||
return Ok(ToolResult::text("No MCP servers connected."));
|
||||
}
|
||||
|
||||
let mut results = Vec::new();
|
||||
for conn in &filtered {
|
||||
if !conn.connected {
|
||||
continue;
|
||||
}
|
||||
results.push(format!(
|
||||
"Server: {} ({} tools available)",
|
||||
conn.name,
|
||||
conn.tools.len()
|
||||
));
|
||||
// Note: Full resource listing requires MCP resources/list support.
|
||||
// Currently showing available tool count as a proxy.
|
||||
}
|
||||
|
||||
if results.is_empty() {
|
||||
Ok(ToolResult::text("No connected MCP servers found."))
|
||||
} else {
|
||||
Ok(ToolResult::text(results.join("\n")))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// ReadMcpResourceTool - Read a specific resource from an MCP server.
|
||||
pub struct ReadMcpResourceTool {
|
||||
mcp_client: Arc<RwLock<Option<Arc<McpClient>>>>,
|
||||
}
|
||||
|
||||
impl ReadMcpResourceTool {
|
||||
pub fn new(mcp_client: Arc<RwLock<Option<Arc<McpClient>>>>) -> Self {
|
||||
Self { mcp_client }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ReadMcpResourceTool {
|
||||
fn name(&self) -> &str {
|
||||
"ReadMcpResource"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Read a specific resource from an MCP server."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"server".to_string(),
|
||||
json!({ "type": "string", "description": "MCP server name" }),
|
||||
),
|
||||
(
|
||||
"uri".to_string(),
|
||||
json!({ "type": "string", "description": "Resource URI to read" }),
|
||||
),
|
||||
]),
|
||||
required: vec!["server".to_string(), "uri".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let server = input
|
||||
.get("server")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'server'".to_string()))?;
|
||||
let uri = input
|
||||
.get("uri")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'uri'".to_string()))?;
|
||||
|
||||
let client_guard = self.mcp_client.read().await;
|
||||
let client = match client_guard.as_ref() {
|
||||
Some(c) => c.clone(),
|
||||
None => {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"MCP server not found: {}",
|
||||
server
|
||||
)))
|
||||
}
|
||||
};
|
||||
|
||||
let connections = client.get_connections().await;
|
||||
let conn = connections.iter().find(|c| c.name == server);
|
||||
|
||||
if conn.is_none() {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"MCP server not found: {}",
|
||||
server
|
||||
)));
|
||||
}
|
||||
|
||||
// MCP resource reading is done via the resources/read method.
|
||||
// Since the current McpClient doesn't expose a read_resource method,
|
||||
// we provide a stub that indicates the operation is recognized but
|
||||
// requires the MCP server to support the resources/read capability.
|
||||
Ok(ToolResult::text(format!(
|
||||
"Resource read request: server=\"{}\" uri=\"{}\". \
|
||||
Full resource reading requires MCP server support for the resources/read method.",
|
||||
server, uri
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
pub mod askuser;
|
||||
pub mod bash;
|
||||
pub mod config_tool;
|
||||
pub mod cron_tools;
|
||||
pub mod diff;
|
||||
pub mod executor;
|
||||
pub mod fileedit;
|
||||
pub mod fileread;
|
||||
pub mod filewrite;
|
||||
pub mod glob_tool;
|
||||
pub mod grep;
|
||||
pub mod lsp_tool;
|
||||
pub mod mcp_resource_tools;
|
||||
pub mod notebook_edit;
|
||||
pub mod plan_tools;
|
||||
pub mod registry;
|
||||
pub mod send_message;
|
||||
pub mod skill_tool;
|
||||
pub mod tasks;
|
||||
pub mod team_tools;
|
||||
pub mod todo_tool;
|
||||
pub mod toolsearch;
|
||||
pub mod webfetch;
|
||||
pub mod websearch;
|
||||
pub mod worktree_tools;
|
||||
|
||||
pub use executor::*;
|
||||
pub use registry::*;
|
||||
@@ -0,0 +1,180 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::path::Path;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// NotebookEditTool - Edit Jupyter notebook (.ipynb) cells.
|
||||
/// Supports insert, replace, and delete operations on notebook cells.
|
||||
pub struct NotebookEditTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for NotebookEditTool {
|
||||
fn name(&self) -> &str {
|
||||
"NotebookEdit"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Edit Jupyter notebook (.ipynb) cells. Can insert, replace, or delete cells."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"file_path".to_string(),
|
||||
json!({ "type": "string", "description": "Path to the .ipynb file" }),
|
||||
),
|
||||
(
|
||||
"command".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["insert", "replace", "delete"],
|
||||
"description": "The edit operation to perform"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"cell_number".to_string(),
|
||||
json!({
|
||||
"type": "number",
|
||||
"description": "Cell index (0-based) to operate on"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"cell_type".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["code", "markdown"],
|
||||
"description": "Type of cell (for insert/replace)"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"source".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "Cell content (for insert/replace)"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec![
|
||||
"file_path".to_string(),
|
||||
"command".to_string(),
|
||||
"cell_number".to_string(),
|
||||
],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let file_path = input
|
||||
.get("file_path")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'file_path'".to_string()))?;
|
||||
let command = input
|
||||
.get("command")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'command'".to_string()))?;
|
||||
let cell_number = input
|
||||
.get("cell_number")
|
||||
.and_then(|v| v.as_u64())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'cell_number'".to_string()))?
|
||||
as usize;
|
||||
|
||||
let resolved = Path::new(&context.working_dir).join(file_path);
|
||||
let resolved_str = resolved.to_string_lossy().to_string();
|
||||
|
||||
let content = match tokio::fs::read_to_string(&resolved).await {
|
||||
Ok(c) => c,
|
||||
Err(e) => return Ok(ToolResult::error(format!("Error: {}", e))),
|
||||
};
|
||||
|
||||
let mut notebook: Value = match serde_json::from_str(&content) {
|
||||
Ok(v) => v,
|
||||
Err(e) => return Ok(ToolResult::error(format!("Error: Invalid JSON: {}", e))),
|
||||
};
|
||||
|
||||
let cells = match notebook.get_mut("cells").and_then(|c| c.as_array_mut()) {
|
||||
Some(c) => c,
|
||||
None => return Ok(ToolResult::error("Error: Invalid notebook format".to_string())),
|
||||
};
|
||||
|
||||
let cell_type = input
|
||||
.get("cell_type")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("code");
|
||||
let source = input
|
||||
.get("source")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
// Split source into lines matching Jupyter format (each line ends with \n except last)
|
||||
let source_lines: Vec<Value> = source
|
||||
.split('\n')
|
||||
.enumerate()
|
||||
.map(|(i, line)| {
|
||||
let total = source.split('\n').count();
|
||||
if i < total - 1 {
|
||||
json!(format!("{}\n", line))
|
||||
} else {
|
||||
json!(line)
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
match command {
|
||||
"insert" => {
|
||||
let mut new_cell = json!({
|
||||
"cell_type": cell_type,
|
||||
"source": source_lines,
|
||||
"metadata": {}
|
||||
});
|
||||
if cell_type != "markdown" {
|
||||
new_cell["outputs"] = json!([]);
|
||||
new_cell["execution_count"] = Value::Null;
|
||||
}
|
||||
cells.insert(cell_number, new_cell);
|
||||
}
|
||||
"replace" => {
|
||||
if cell_number >= cells.len() {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"Error: Cell {} does not exist",
|
||||
cell_number
|
||||
)));
|
||||
}
|
||||
cells[cell_number]["source"] = json!(source_lines);
|
||||
if !cell_type.is_empty() {
|
||||
cells[cell_number]["cell_type"] = json!(cell_type);
|
||||
}
|
||||
}
|
||||
"delete" => {
|
||||
if cell_number >= cells.len() {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"Error: Cell {} does not exist",
|
||||
cell_number
|
||||
)));
|
||||
}
|
||||
cells.remove(cell_number);
|
||||
}
|
||||
_ => {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"Error: Unknown command '{}'",
|
||||
command
|
||||
)));
|
||||
}
|
||||
}
|
||||
|
||||
let output = serde_json::to_string_pretty(¬ebook)
|
||||
.map_err(|e| ToolError::ExecutionError(format!("Failed to serialize: {}", e)))?;
|
||||
|
||||
if let Err(e) = tokio::fs::write(&resolved, output).await {
|
||||
return Ok(ToolResult::error(format!("Error writing file: {}", e)));
|
||||
}
|
||||
|
||||
Ok(ToolResult::text(format!(
|
||||
"Notebook {}: cell {} in {}",
|
||||
command, cell_number, resolved_str
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// Shared plan mode state.
|
||||
#[derive(Clone)]
|
||||
pub struct PlanState {
|
||||
active: Arc<RwLock<bool>>,
|
||||
current_plan: Arc<RwLock<Option<String>>>,
|
||||
}
|
||||
|
||||
impl PlanState {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
active: Arc::new(RwLock::new(false)),
|
||||
current_plan: Arc::new(RwLock::new(None)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn is_active(&self) -> bool {
|
||||
*self.active.read().await
|
||||
}
|
||||
|
||||
pub async fn get_plan(&self) -> Option<String> {
|
||||
self.current_plan.read().await.clone()
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// EnterPlanModeTool
|
||||
// ============================================================================
|
||||
|
||||
pub struct EnterPlanModeTool {
|
||||
state: PlanState,
|
||||
}
|
||||
|
||||
impl EnterPlanModeTool {
|
||||
pub fn new(state: PlanState) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for EnterPlanModeTool {
|
||||
fn name(&self) -> &str {
|
||||
"EnterPlanMode"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Enter plan/design mode for complex tasks. In plan mode, the agent focuses on designing the approach before executing."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::new(),
|
||||
required: Vec::new(),
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, _input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let mut active = self.state.active.write().await;
|
||||
if *active {
|
||||
return Ok(ToolResult::text("Already in plan mode."));
|
||||
}
|
||||
|
||||
*active = true;
|
||||
let mut plan = self.state.current_plan.write().await;
|
||||
*plan = None;
|
||||
|
||||
Ok(ToolResult::text(
|
||||
"Entered plan mode. Design your approach before executing. Use ExitPlanMode when the plan is ready.",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// ExitPlanModeTool
|
||||
// ============================================================================
|
||||
|
||||
pub struct ExitPlanModeTool {
|
||||
state: PlanState,
|
||||
}
|
||||
|
||||
impl ExitPlanModeTool {
|
||||
pub fn new(state: PlanState) -> Self {
|
||||
Self { state }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ExitPlanModeTool {
|
||||
fn name(&self) -> &str {
|
||||
"ExitPlanMode"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Exit plan mode with a completed plan. The plan will be recorded and execution can proceed."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"plan".to_string(),
|
||||
json!({ "type": "string", "description": "The completed plan" }),
|
||||
),
|
||||
(
|
||||
"approved".to_string(),
|
||||
json!({ "type": "boolean", "description": "Whether the plan is approved for execution" }),
|
||||
),
|
||||
]),
|
||||
required: Vec::new(),
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let mut active = self.state.active.write().await;
|
||||
if !*active {
|
||||
return Ok(ToolResult::error("Not in plan mode."));
|
||||
}
|
||||
|
||||
*active = false;
|
||||
|
||||
let plan_text = input.get("plan").and_then(|v| v.as_str()).map(String::from);
|
||||
let approved = input
|
||||
.get("approved")
|
||||
.and_then(|v| v.as_bool())
|
||||
.unwrap_or(true);
|
||||
|
||||
let mut current_plan = self.state.current_plan.write().await;
|
||||
*current_plan = plan_text.clone();
|
||||
|
||||
let status = if approved {
|
||||
"approved"
|
||||
} else {
|
||||
"pending approval"
|
||||
};
|
||||
|
||||
let mut msg = format!("Plan mode exited. Plan status: {}.", status);
|
||||
if let Some(ref plan) = plan_text {
|
||||
msg.push_str(&format!("\n\nPlan:\n{}", plan));
|
||||
}
|
||||
|
||||
Ok(ToolResult::text(msg))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
use crate::types::Tool;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// Registry of available tools.
|
||||
pub struct ToolRegistry {
|
||||
tools: HashMap<String, Arc<dyn Tool>>,
|
||||
}
|
||||
|
||||
impl ToolRegistry {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tools: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a registry with all default built-in tools.
|
||||
pub fn default_registry() -> Self {
|
||||
let mut registry = Self::new();
|
||||
|
||||
// File I/O tools
|
||||
registry.register(Arc::new(super::bash::BashTool));
|
||||
registry.register(Arc::new(super::fileread::FileReadTool));
|
||||
registry.register(Arc::new(super::filewrite::FileWriteTool));
|
||||
registry.register(Arc::new(super::fileedit::FileEditTool));
|
||||
registry.register(Arc::new(super::glob_tool::GlobTool));
|
||||
registry.register(Arc::new(super::grep::GrepTool));
|
||||
|
||||
// Web tools
|
||||
registry.register(Arc::new(super::webfetch::WebFetchTool));
|
||||
registry.register(Arc::new(super::websearch::WebSearchTool::default()));
|
||||
|
||||
// User interaction
|
||||
registry.register(Arc::new(super::askuser::AskUserTool::default()));
|
||||
|
||||
// Task tools
|
||||
let task_store = super::tasks::TaskStore::new();
|
||||
registry.register(Arc::new(super::tasks::TaskCreateTool::new(
|
||||
task_store.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::tasks::TaskGetTool::new(
|
||||
task_store.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::tasks::TaskListTool::new(
|
||||
task_store.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::tasks::TaskUpdateTool::new(
|
||||
task_store.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::tasks::TaskStopTool::new(
|
||||
task_store.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::tasks::TaskOutputTool::new(task_store)));
|
||||
|
||||
// Tool search
|
||||
registry.register(Arc::new(super::toolsearch::ToolSearchTool::default()));
|
||||
|
||||
// Notebook edit
|
||||
registry.register(Arc::new(super::notebook_edit::NotebookEditTool));
|
||||
|
||||
// Inter-agent messaging
|
||||
let mailbox = super::send_message::new_mailbox();
|
||||
registry.register(Arc::new(super::send_message::SendMessageTool::new(mailbox)));
|
||||
|
||||
// Team tools
|
||||
let team_store = super::team_tools::TeamStore::new();
|
||||
registry.register(Arc::new(super::team_tools::TeamCreateTool::new(
|
||||
team_store.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::team_tools::TeamDeleteTool::new(team_store)));
|
||||
|
||||
// Plan mode tools
|
||||
let plan_state = super::plan_tools::PlanState::new();
|
||||
registry.register(Arc::new(super::plan_tools::EnterPlanModeTool::new(
|
||||
plan_state.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::plan_tools::ExitPlanModeTool::new(plan_state)));
|
||||
|
||||
// Worktree tools
|
||||
let worktree_store = super::worktree_tools::new_worktree_store();
|
||||
registry.register(Arc::new(super::worktree_tools::EnterWorktreeTool::new(
|
||||
worktree_store.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::worktree_tools::ExitWorktreeTool::new(
|
||||
worktree_store,
|
||||
)));
|
||||
|
||||
// Todo tool
|
||||
let todo_store = super::todo_tool::TodoStore::new();
|
||||
registry.register(Arc::new(super::todo_tool::TodoWriteTool::new(todo_store)));
|
||||
|
||||
// Cron tools
|
||||
let cron_store = super::cron_tools::CronStore::new();
|
||||
registry.register(Arc::new(super::cron_tools::CronCreateTool::new(
|
||||
cron_store.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::cron_tools::CronDeleteTool::new(
|
||||
cron_store.clone(),
|
||||
)));
|
||||
registry.register(Arc::new(super::cron_tools::CronListTool::new(cron_store)));
|
||||
|
||||
// Config tool
|
||||
let config_store = super::config_tool::ConfigStore::new();
|
||||
registry.register(Arc::new(super::config_tool::ConfigTool::new(config_store)));
|
||||
|
||||
// LSP tool
|
||||
registry.register(Arc::new(super::lsp_tool::LSPTool));
|
||||
|
||||
// MCP resource tools
|
||||
let mcp_client: Arc<RwLock<Option<Arc<crate::mcp::McpClient>>>> =
|
||||
Arc::new(RwLock::new(None));
|
||||
registry.register(Arc::new(
|
||||
super::mcp_resource_tools::ListMcpResourcesTool::new(mcp_client.clone()),
|
||||
));
|
||||
registry.register(Arc::new(
|
||||
super::mcp_resource_tools::ReadMcpResourceTool::new(mcp_client),
|
||||
));
|
||||
|
||||
registry
|
||||
}
|
||||
|
||||
pub fn register(&mut self, tool: Arc<dyn Tool>) {
|
||||
self.tools.insert(tool.name().to_string(), tool);
|
||||
}
|
||||
|
||||
/// Register a tool, replacing any existing tool with the same name.
|
||||
pub fn register_replace(&mut self, tool: Arc<dyn Tool>) {
|
||||
self.tools.insert(tool.name().to_string(), tool);
|
||||
}
|
||||
|
||||
pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
|
||||
self.tools.get(name).cloned()
|
||||
}
|
||||
|
||||
pub fn all(&self) -> Vec<Arc<dyn Tool>> {
|
||||
self.tools.values().cloned().collect()
|
||||
}
|
||||
|
||||
pub fn names(&self) -> Vec<String> {
|
||||
self.tools.keys().cloned().collect()
|
||||
}
|
||||
|
||||
pub fn filter<F>(&self, predicate: F) -> Vec<Arc<dyn Tool>>
|
||||
where
|
||||
F: Fn(&dyn Tool) -> bool,
|
||||
{
|
||||
self.tools
|
||||
.values()
|
||||
.filter(|t| predicate(t.as_ref()))
|
||||
.cloned()
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn len(&self) -> usize {
|
||||
self.tools.len()
|
||||
}
|
||||
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.tools.is_empty()
|
||||
}
|
||||
|
||||
/// Remove tools by name.
|
||||
pub fn remove(&mut self, names: &[&str]) {
|
||||
for name in names {
|
||||
self.tools.remove(*name);
|
||||
}
|
||||
}
|
||||
|
||||
/// Keep only the specified tools.
|
||||
pub fn retain(&mut self, names: &[&str]) {
|
||||
let name_set: std::collections::HashSet<&str> = names.iter().copied().collect();
|
||||
self.tools.retain(|k, _| name_set.contains(k.as_str()));
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ToolRegistry {
|
||||
fn default() -> Self {
|
||||
Self::default_registry()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// Shared mailbox for inter-agent communication.
|
||||
pub type Mailbox = Arc<RwLock<HashMap<String, Vec<AgentMessage>>>>;
|
||||
|
||||
/// A message between agents.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AgentMessage {
|
||||
pub from: String,
|
||||
pub to: String,
|
||||
pub content: String,
|
||||
pub timestamp: String,
|
||||
pub message_type: String, // "text", "shutdown_request", "shutdown_response", "plan_approval_response"
|
||||
}
|
||||
|
||||
/// Create a new empty mailbox.
|
||||
pub fn new_mailbox() -> Mailbox {
|
||||
Arc::new(RwLock::new(HashMap::new()))
|
||||
}
|
||||
|
||||
/// Read and drain messages for an agent.
|
||||
pub async fn read_mailbox(mailbox: &Mailbox, agent_name: &str) -> Vec<AgentMessage> {
|
||||
let mut mb = mailbox.write().await;
|
||||
mb.remove(agent_name).unwrap_or_default()
|
||||
}
|
||||
|
||||
/// Write a message to an agent's mailbox.
|
||||
pub async fn write_to_mailbox(mailbox: &Mailbox, agent_name: &str, message: AgentMessage) {
|
||||
let mut mb = mailbox.write().await;
|
||||
mb.entry(agent_name.to_string())
|
||||
.or_insert_with(Vec::new)
|
||||
.push(message);
|
||||
}
|
||||
|
||||
/// SendMessageTool - Send a message to another agent or teammate.
|
||||
pub struct SendMessageTool {
|
||||
mailbox: Mailbox,
|
||||
}
|
||||
|
||||
impl SendMessageTool {
|
||||
pub fn new(mailbox: Mailbox) -> Self {
|
||||
Self { mailbox }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SendMessageTool {
|
||||
fn name(&self) -> &str {
|
||||
"SendMessage"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Send a message to another agent or teammate. Supports plain text and structured protocol messages."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"to".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "Recipient agent name or ID. Use \"*\" for broadcast."
|
||||
}),
|
||||
),
|
||||
(
|
||||
"content".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "Message content"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"type".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["text", "shutdown_request", "shutdown_response", "plan_approval_response"],
|
||||
"description": "Message type (default: text)"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["to".to_string(), "content".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let to = input
|
||||
.get("to")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'to'".to_string()))?;
|
||||
let content = input
|
||||
.get("content")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'content'".to_string()))?;
|
||||
let msg_type = input
|
||||
.get("type")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("text");
|
||||
|
||||
let timestamp = chrono::Utc::now().to_rfc3339();
|
||||
|
||||
let message = AgentMessage {
|
||||
from: "self".to_string(),
|
||||
to: to.to_string(),
|
||||
content: content.to_string(),
|
||||
timestamp,
|
||||
message_type: msg_type.to_string(),
|
||||
};
|
||||
|
||||
if to == "*" {
|
||||
// Broadcast to all known mailboxes
|
||||
let mb = self.mailbox.read().await;
|
||||
let names: Vec<String> = mb.keys().cloned().collect();
|
||||
drop(mb);
|
||||
|
||||
for name in &names {
|
||||
let mut broadcast_msg = message.clone();
|
||||
broadcast_msg.to = name.clone();
|
||||
write_to_mailbox(&self.mailbox, name, broadcast_msg).await;
|
||||
}
|
||||
|
||||
Ok(ToolResult::text("Message broadcast to all agents"))
|
||||
} else {
|
||||
write_to_mailbox(&self.mailbox, to, message).await;
|
||||
Ok(ToolResult::text(format!("Message sent to {}", to)))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::skills::SkillRegistry;
|
||||
use crate::skills::loader::substitute_variables;
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
pub struct SkillTool {
|
||||
registry: Arc<RwLock<SkillRegistry>>,
|
||||
}
|
||||
|
||||
impl SkillTool {
|
||||
pub fn new(registry: Arc<RwLock<SkillRegistry>>) -> Self {
|
||||
Self { registry }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for SkillTool {
|
||||
fn name(&self) -> &str {
|
||||
"Skill"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Invoke a skill by name. Skills provide specialized instructions and workflows for specific tasks. Use the skill name from the Available Skills list."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"skill_name".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The name of the skill to invoke"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"arguments".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "Arguments to pass to the skill (replaces $ARGUMENTS in skill content)"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["skill_name".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(
|
||||
&self,
|
||||
input: Value,
|
||||
_context: &ToolUseContext,
|
||||
) -> Result<ToolResult, ToolError> {
|
||||
let skill_name = input
|
||||
.get("skill_name")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'skill_name' field".to_string()))?;
|
||||
|
||||
let arguments = input
|
||||
.get("arguments")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("");
|
||||
|
||||
let registry = self.registry.read().await;
|
||||
|
||||
let skill = registry.get(skill_name).ok_or_else(|| {
|
||||
ToolError::ExecutionError(format!(
|
||||
"Skill '{}' not found. Available skills: {}",
|
||||
skill_name,
|
||||
registry
|
||||
.user_invocable()
|
||||
.iter()
|
||||
.map(|s| s.name.as_str())
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
))
|
||||
})?;
|
||||
|
||||
if !registry.all_enabled().iter().any(|s| s.name == skill_name) {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"Skill '{}' is currently disabled.",
|
||||
skill_name
|
||||
)));
|
||||
}
|
||||
|
||||
let skill_dir = skill
|
||||
.path
|
||||
.parent()
|
||||
.map(|p| p.to_string_lossy().to_string())
|
||||
.unwrap_or_default();
|
||||
|
||||
let content = substitute_variables(&skill.content, arguments, &skill_dir);
|
||||
|
||||
let result = format!(
|
||||
"[SKILL INSTRUCTIONS - You MUST follow these instructions precisely]\n\n{}\n\n[END SKILL INSTRUCTIONS]",
|
||||
content
|
||||
);
|
||||
|
||||
Ok(ToolResult::text(result))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,452 @@
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// Task status values.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum TaskStatus {
|
||||
Pending,
|
||||
InProgress,
|
||||
Completed,
|
||||
Failed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
/// A task in the task store.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Task {
|
||||
pub id: String,
|
||||
pub subject: String,
|
||||
pub status: TaskStatus,
|
||||
#[serde(default)]
|
||||
pub owner: Option<String>,
|
||||
#[serde(default)]
|
||||
pub description: Option<String>,
|
||||
#[serde(default)]
|
||||
pub metadata: HashMap<String, Value>,
|
||||
#[serde(default)]
|
||||
pub output: Option<String>,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
/// In-memory task store shared across tools.
|
||||
#[derive(Clone)]
|
||||
pub struct TaskStore {
|
||||
tasks: Arc<RwLock<HashMap<String, Task>>>,
|
||||
counter: Arc<RwLock<u64>>,
|
||||
}
|
||||
|
||||
impl TaskStore {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tasks: Arc::new(RwLock::new(HashMap::new())),
|
||||
counter: Arc::new(RwLock::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_id(&self) -> String {
|
||||
let mut counter = self.counter.write().await;
|
||||
*counter += 1;
|
||||
format!("task_{}", counter)
|
||||
}
|
||||
}
|
||||
|
||||
// --- TaskCreateTool ---
|
||||
|
||||
pub struct TaskCreateTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
impl TaskCreateTool {
|
||||
pub fn new(store: TaskStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskCreateTool {
|
||||
fn name(&self) -> &str {
|
||||
"TaskCreate"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Create a new task for tracking work."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"subject".to_string(),
|
||||
json!({ "type": "string", "description": "Task subject/title" }),
|
||||
),
|
||||
(
|
||||
"description".to_string(),
|
||||
json!({ "type": "string", "description": "Optional task description" }),
|
||||
),
|
||||
(
|
||||
"owner".to_string(),
|
||||
json!({ "type": "string", "description": "Optional task owner" }),
|
||||
),
|
||||
]),
|
||||
required: vec!["subject".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let subject = input
|
||||
.get("subject")
|
||||
.and_then(|s| s.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'subject'".to_string()))?;
|
||||
|
||||
let id = self.store.next_id().await;
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
|
||||
let task = Task {
|
||||
id: id.clone(),
|
||||
subject: subject.to_string(),
|
||||
status: TaskStatus::Pending,
|
||||
owner: input.get("owner").and_then(|o| o.as_str()).map(String::from),
|
||||
description: input
|
||||
.get("description")
|
||||
.and_then(|d| d.as_str())
|
||||
.map(String::from),
|
||||
output: None,
|
||||
metadata: HashMap::new(),
|
||||
created_at: now.clone(),
|
||||
updated_at: now,
|
||||
};
|
||||
|
||||
let mut tasks = self.store.tasks.write().await;
|
||||
tasks.insert(id.clone(), task);
|
||||
|
||||
Ok(ToolResult::text(format!("Created task: {}", id)))
|
||||
}
|
||||
}
|
||||
|
||||
// --- TaskGetTool ---
|
||||
|
||||
pub struct TaskGetTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
impl TaskGetTool {
|
||||
pub fn new(store: TaskStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskGetTool {
|
||||
fn name(&self) -> &str {
|
||||
"TaskGet"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Get a task by ID."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([(
|
||||
"id".to_string(),
|
||||
json!({ "type": "string", "description": "Task ID" }),
|
||||
)]),
|
||||
required: vec!["id".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let id = input
|
||||
.get("id")
|
||||
.and_then(|i| i.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'id'".to_string()))?;
|
||||
|
||||
let tasks = self.store.tasks.read().await;
|
||||
match tasks.get(id) {
|
||||
Some(task) => Ok(ToolResult::text(
|
||||
serde_json::to_string_pretty(task).unwrap_or_default(),
|
||||
)),
|
||||
None => Ok(ToolResult::error(format!("Task not found: {}", id))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- TaskListTool ---
|
||||
|
||||
pub struct TaskListTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
impl TaskListTool {
|
||||
pub fn new(store: TaskStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskListTool {
|
||||
fn name(&self) -> &str {
|
||||
"TaskList"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"List all tasks."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema::default()
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, _input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let tasks = self.store.tasks.read().await;
|
||||
|
||||
if tasks.is_empty() {
|
||||
return Ok(ToolResult::text("No tasks.".to_string()));
|
||||
}
|
||||
|
||||
let mut task_list: Vec<&Task> = tasks.values().collect();
|
||||
task_list.sort_by(|a, b| a.created_at.cmp(&b.created_at));
|
||||
|
||||
let formatted: Vec<String> = task_list
|
||||
.iter()
|
||||
.map(|t| {
|
||||
format!(
|
||||
"- [{}] {} ({:?}){}",
|
||||
t.id,
|
||||
t.subject,
|
||||
t.status,
|
||||
t.owner
|
||||
.as_ref()
|
||||
.map(|o| format!(" @{}", o))
|
||||
.unwrap_or_default()
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(ToolResult::text(formatted.join("\n")))
|
||||
}
|
||||
}
|
||||
|
||||
// --- TaskUpdateTool ---
|
||||
|
||||
pub struct TaskUpdateTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
impl TaskUpdateTool {
|
||||
pub fn new(store: TaskStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskUpdateTool {
|
||||
fn name(&self) -> &str {
|
||||
"TaskUpdate"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Update a task's status or other properties."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"id".to_string(),
|
||||
json!({ "type": "string", "description": "Task ID" }),
|
||||
),
|
||||
(
|
||||
"status".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["pending", "in_progress", "completed", "failed", "cancelled"],
|
||||
"description": "New task status"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"owner".to_string(),
|
||||
json!({ "type": "string", "description": "New task owner" }),
|
||||
),
|
||||
(
|
||||
"output".to_string(),
|
||||
json!({ "type": "string", "description": "Task output/result" }),
|
||||
),
|
||||
]),
|
||||
required: vec!["id".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let id = input
|
||||
.get("id")
|
||||
.and_then(|i| i.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'id'".to_string()))?;
|
||||
|
||||
let mut tasks = self.store.tasks.write().await;
|
||||
match tasks.get_mut(id) {
|
||||
Some(task) => {
|
||||
if let Some(status) = input.get("status").and_then(|s| s.as_str()) {
|
||||
task.status = match status {
|
||||
"pending" => TaskStatus::Pending,
|
||||
"in_progress" => TaskStatus::InProgress,
|
||||
"completed" => TaskStatus::Completed,
|
||||
"failed" => TaskStatus::Failed,
|
||||
"cancelled" => TaskStatus::Cancelled,
|
||||
_ => {
|
||||
return Ok(ToolResult::error(format!("Invalid status: {}", status)));
|
||||
}
|
||||
};
|
||||
}
|
||||
if let Some(owner) = input.get("owner").and_then(|o| o.as_str()) {
|
||||
task.owner = Some(owner.to_string());
|
||||
}
|
||||
if let Some(output) = input.get("output").and_then(|o| o.as_str()) {
|
||||
task.output = Some(output.to_string());
|
||||
}
|
||||
task.updated_at = chrono::Utc::now().to_rfc3339();
|
||||
|
||||
Ok(ToolResult::text(format!("Updated task: {}", id)))
|
||||
}
|
||||
None => Ok(ToolResult::error(format!("Task not found: {}", id))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- TaskStopTool ---
|
||||
|
||||
pub struct TaskStopTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
impl TaskStopTool {
|
||||
pub fn new(store: TaskStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskStopTool {
|
||||
fn name(&self) -> &str {
|
||||
"TaskStop"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Stop/cancel a running task."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"id".to_string(),
|
||||
json!({ "type": "string", "description": "Task ID to stop" }),
|
||||
),
|
||||
(
|
||||
"reason".to_string(),
|
||||
json!({ "type": "string", "description": "Reason for stopping" }),
|
||||
),
|
||||
]),
|
||||
required: vec!["id".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let id = input
|
||||
.get("id")
|
||||
.and_then(|i| i.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'id'".to_string()))?;
|
||||
|
||||
let mut tasks = self.store.tasks.write().await;
|
||||
match tasks.get_mut(id) {
|
||||
Some(task) => {
|
||||
task.status = TaskStatus::Cancelled;
|
||||
task.updated_at = chrono::Utc::now().to_rfc3339();
|
||||
if let Some(reason) = input.get("reason").and_then(|r| r.as_str()) {
|
||||
task.output = Some(format!("Stopped: {}", reason));
|
||||
}
|
||||
Ok(ToolResult::text(format!("Task stopped: {}", id)))
|
||||
}
|
||||
None => Ok(ToolResult::error(format!("Task not found: {}", id))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- TaskOutputTool ---
|
||||
|
||||
pub struct TaskOutputTool {
|
||||
store: TaskStore,
|
||||
}
|
||||
|
||||
impl TaskOutputTool {
|
||||
pub fn new(store: TaskStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TaskOutputTool {
|
||||
fn name(&self) -> &str {
|
||||
"TaskOutput"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Get the output/result of a task."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([(
|
||||
"id".to_string(),
|
||||
json!({ "type": "string", "description": "Task ID" }),
|
||||
)]),
|
||||
required: vec!["id".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let id = input
|
||||
.get("id")
|
||||
.and_then(|i| i.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'id'".to_string()))?;
|
||||
|
||||
let tasks = self.store.tasks.read().await;
|
||||
match tasks.get(id) {
|
||||
Some(task) => {
|
||||
let output = task.output.as_deref().unwrap_or("(no output yet)");
|
||||
Ok(ToolResult::text(output))
|
||||
}
|
||||
None => Ok(ToolResult::error(format!("Task not found: {}", id))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// A team for multi-agent coordination.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Team {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub members: Vec<String>,
|
||||
pub leader: String,
|
||||
pub tasks: Vec<String>,
|
||||
pub created_at: String,
|
||||
pub status: String, // "active" or "disbanded"
|
||||
}
|
||||
|
||||
/// In-memory team store shared across tools.
|
||||
#[derive(Clone)]
|
||||
pub struct TeamStore {
|
||||
teams: Arc<RwLock<HashMap<String, Team>>>,
|
||||
counter: Arc<RwLock<u64>>,
|
||||
}
|
||||
|
||||
impl TeamStore {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
teams: Arc::new(RwLock::new(HashMap::new())),
|
||||
counter: Arc::new(RwLock::new(0)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_id(&self) -> String {
|
||||
let mut counter = self.counter.write().await;
|
||||
*counter += 1;
|
||||
format!("team_{}", counter)
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// TeamCreateTool
|
||||
// ============================================================================
|
||||
|
||||
pub struct TeamCreateTool {
|
||||
store: TeamStore,
|
||||
}
|
||||
|
||||
impl TeamCreateTool {
|
||||
pub fn new(store: TeamStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TeamCreateTool {
|
||||
fn name(&self) -> &str {
|
||||
"TeamCreate"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Create a multi-agent team for coordinated work. Assigns a lead and manages member composition."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"name".to_string(),
|
||||
json!({ "type": "string", "description": "Team name" }),
|
||||
),
|
||||
(
|
||||
"members".to_string(),
|
||||
json!({
|
||||
"type": "array",
|
||||
"items": { "type": "string" },
|
||||
"description": "List of agent/teammate names"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"task_description".to_string(),
|
||||
json!({ "type": "string", "description": "Description of the team's mission" }),
|
||||
),
|
||||
]),
|
||||
required: vec!["name".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let name = input
|
||||
.get("name")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'name'".to_string()))?;
|
||||
|
||||
let members: Vec<String> = input
|
||||
.get("members")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str().map(String::from))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let tasks: Vec<String> = input
|
||||
.get("task_description")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|t| vec![t.to_string()])
|
||||
.unwrap_or_default();
|
||||
|
||||
let id = self.store.next_id().await;
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
|
||||
let member_count = members.len();
|
||||
let team = Team {
|
||||
id: id.clone(),
|
||||
name: name.to_string(),
|
||||
members,
|
||||
leader: "self".to_string(),
|
||||
tasks,
|
||||
created_at: now,
|
||||
status: "active".to_string(),
|
||||
};
|
||||
|
||||
let mut teams = self.store.teams.write().await;
|
||||
teams.insert(id.clone(), team);
|
||||
|
||||
Ok(ToolResult::text(format!(
|
||||
"Team created: {} \"{}\" with {} members",
|
||||
id, name, member_count
|
||||
)))
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// TeamDeleteTool
|
||||
// ============================================================================
|
||||
|
||||
pub struct TeamDeleteTool {
|
||||
store: TeamStore,
|
||||
}
|
||||
|
||||
impl TeamDeleteTool {
|
||||
pub fn new(store: TeamStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TeamDeleteTool {
|
||||
fn name(&self) -> &str {
|
||||
"TeamDelete"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Disband a team and clean up resources."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([(
|
||||
"id".to_string(),
|
||||
json!({ "type": "string", "description": "Team ID to disband" }),
|
||||
)]),
|
||||
required: vec!["id".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let id = input
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'id'".to_string()))?;
|
||||
|
||||
let mut teams = self.store.teams.write().await;
|
||||
match teams.remove(id) {
|
||||
Some(team) => Ok(ToolResult::text(format!("Team disbanded: {}", team.name))),
|
||||
None => Ok(ToolResult::error(format!("Team not found: {}", id))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// A todo item.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TodoItem {
|
||||
pub id: u64,
|
||||
pub text: String,
|
||||
pub done: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub priority: Option<String>, // "high", "medium", "low"
|
||||
}
|
||||
|
||||
/// In-memory todo store.
|
||||
#[derive(Clone)]
|
||||
pub struct TodoStore {
|
||||
items: Arc<RwLock<Vec<TodoItem>>>,
|
||||
counter: Arc<RwLock<u64>>,
|
||||
}
|
||||
|
||||
impl TodoStore {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
items: Arc::new(RwLock::new(Vec::new())),
|
||||
counter: Arc::new(RwLock::new(0)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// TodoWriteTool - Manage a session todo/checklist.
|
||||
pub struct TodoWriteTool {
|
||||
store: TodoStore,
|
||||
}
|
||||
|
||||
impl TodoWriteTool {
|
||||
pub fn new(store: TodoStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TodoWriteTool {
|
||||
fn name(&self) -> &str {
|
||||
"TodoWrite"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Manage a session todo/checklist. Supports add, toggle, remove, list, and clear operations."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"action".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["add", "toggle", "remove", "list", "clear"],
|
||||
"description": "Operation to perform"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"text".to_string(),
|
||||
json!({ "type": "string", "description": "Todo item text (for add)" }),
|
||||
),
|
||||
(
|
||||
"id".to_string(),
|
||||
json!({ "type": "number", "description": "Todo item ID (for toggle/remove)" }),
|
||||
),
|
||||
(
|
||||
"priority".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["high", "medium", "low"],
|
||||
"description": "Priority level (for add)"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["action".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let action = input
|
||||
.get("action")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'action'".to_string()))?;
|
||||
|
||||
match action {
|
||||
"add" => {
|
||||
let text = match input.get("text").and_then(|v| v.as_str()) {
|
||||
Some(t) => t,
|
||||
None => return Ok(ToolResult::error("text required")),
|
||||
};
|
||||
let priority = input
|
||||
.get("priority")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
|
||||
let mut counter = self.store.counter.write().await;
|
||||
*counter += 1;
|
||||
let id = *counter;
|
||||
|
||||
let item = TodoItem {
|
||||
id,
|
||||
text: text.to_string(),
|
||||
done: false,
|
||||
priority,
|
||||
};
|
||||
|
||||
let mut items = self.store.items.write().await;
|
||||
items.push(item);
|
||||
|
||||
Ok(ToolResult::text(format!(
|
||||
"Todo added: #{} \"{}\"",
|
||||
id, text
|
||||
)))
|
||||
}
|
||||
|
||||
"toggle" => {
|
||||
let id = input
|
||||
.get("id")
|
||||
.and_then(|v| v.as_u64())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'id'".to_string()))?;
|
||||
|
||||
let mut items = self.store.items.write().await;
|
||||
match items.iter_mut().find(|t| t.id == id) {
|
||||
Some(item) => {
|
||||
item.done = !item.done;
|
||||
let status = if item.done { "completed" } else { "reopened" };
|
||||
Ok(ToolResult::text(format!("Todo #{} {}", id, status)))
|
||||
}
|
||||
None => Ok(ToolResult::error(format!("Todo #{} not found", id))),
|
||||
}
|
||||
}
|
||||
|
||||
"remove" => {
|
||||
let id = input
|
||||
.get("id")
|
||||
.and_then(|v| v.as_u64())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'id'".to_string()))?;
|
||||
|
||||
let mut items = self.store.items.write().await;
|
||||
let len_before = items.len();
|
||||
items.retain(|t| t.id != id);
|
||||
|
||||
if items.len() < len_before {
|
||||
Ok(ToolResult::text(format!("Todo #{} removed", id)))
|
||||
} else {
|
||||
Ok(ToolResult::error(format!("Todo #{} not found", id)))
|
||||
}
|
||||
}
|
||||
|
||||
"list" => {
|
||||
let items = self.store.items.read().await;
|
||||
if items.is_empty() {
|
||||
return Ok(ToolResult::text("No todos."));
|
||||
}
|
||||
|
||||
let lines: Vec<String> = items
|
||||
.iter()
|
||||
.map(|t| {
|
||||
let checkbox = if t.done { "[x]" } else { "[ ]" };
|
||||
let priority = t
|
||||
.priority
|
||||
.as_ref()
|
||||
.map(|p| format!(" ({})", p))
|
||||
.unwrap_or_default();
|
||||
format!("{} #{} {}{}", checkbox, t.id, t.text, priority)
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(ToolResult::text(lines.join("\n")))
|
||||
}
|
||||
|
||||
"clear" => {
|
||||
let mut items = self.store.items.write().await;
|
||||
items.clear();
|
||||
Ok(ToolResult::text("All todos cleared."))
|
||||
}
|
||||
|
||||
_ => Ok(ToolResult::error(format!("Unknown action: {}", action))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// ToolSearch allows lazy discovery of tools by name or keyword.
|
||||
pub struct ToolSearchTool {
|
||||
/// All available tools for searching.
|
||||
available_tools: Arc<RwLock<Vec<ToolInfo>>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ToolInfo {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub input_schema: Value,
|
||||
}
|
||||
|
||||
impl Default for ToolSearchTool {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
available_tools: Arc::new(RwLock::new(Vec::new())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl ToolSearchTool {
|
||||
pub fn new(tools: Vec<ToolInfo>) -> Self {
|
||||
Self {
|
||||
available_tools: Arc::new(RwLock::new(tools)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn set_tools(&self, tools: Vec<ToolInfo>) {
|
||||
let mut available = self.available_tools.write().await;
|
||||
*available = tools;
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ToolSearchTool {
|
||||
fn name(&self) -> &str {
|
||||
"ToolSearch"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Search for available tools by name or keyword. Use 'select:Read,Edit' for exact names, or keywords for fuzzy search."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"query".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "Search query. Use 'select:Tool1,Tool2' for exact matches, or keywords"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"max_results".to_string(),
|
||||
json!({
|
||||
"type": "number",
|
||||
"description": "Maximum results to return (default 5)"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["query".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let query = input
|
||||
.get("query")
|
||||
.and_then(|q| q.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'query'".to_string()))?;
|
||||
|
||||
let max_results = input
|
||||
.get("max_results")
|
||||
.and_then(|m| m.as_u64())
|
||||
.unwrap_or(5) as usize;
|
||||
|
||||
let tools = self.available_tools.read().await;
|
||||
|
||||
let matches: Vec<&ToolInfo> = if query.starts_with("select:") {
|
||||
let names: Vec<&str> = query[7..].split(',').map(|s| s.trim()).collect();
|
||||
tools
|
||||
.iter()
|
||||
.filter(|t| names.iter().any(|n| t.name.eq_ignore_ascii_case(n)))
|
||||
.collect()
|
||||
} else {
|
||||
let query_lower = query.to_lowercase();
|
||||
let keywords: Vec<&str> = query_lower.split_whitespace().collect();
|
||||
let mut scored: Vec<(&ToolInfo, usize)> = tools
|
||||
.iter()
|
||||
.filter_map(|t| {
|
||||
let name_lower = t.name.to_lowercase();
|
||||
let desc_lower = t.description.to_lowercase();
|
||||
let score: usize = keywords
|
||||
.iter()
|
||||
.filter(|kw| name_lower.contains(*kw) || desc_lower.contains(*kw))
|
||||
.count();
|
||||
if score > 0 {
|
||||
Some((t, score))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
scored.sort_by(|a, b| b.1.cmp(&a.1));
|
||||
scored.into_iter().map(|(t, _)| t).collect()
|
||||
};
|
||||
|
||||
let matches: Vec<&ToolInfo> = matches.into_iter().take(max_results).collect();
|
||||
|
||||
if matches.is_empty() {
|
||||
return Ok(ToolResult::text(format!(
|
||||
"No tools found matching: {}",
|
||||
query
|
||||
)));
|
||||
}
|
||||
|
||||
let mut result = String::new();
|
||||
for tool in matches {
|
||||
result.push_str(&format!(
|
||||
"## {}\n{}\n\nInput Schema:\n```json\n{}\n```\n\n",
|
||||
tool.name,
|
||||
tool.description,
|
||||
serde_json::to_string_pretty(&tool.input_schema).unwrap_or_default()
|
||||
));
|
||||
}
|
||||
|
||||
Ok(ToolResult::text(result))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
const MAX_RESPONSE_SIZE: usize = 512 * 1024; // 512KB
|
||||
const TIMEOUT_SECS: u64 = 30;
|
||||
|
||||
pub struct WebFetchTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for WebFetchTool {
|
||||
fn name(&self) -> &str {
|
||||
"WebFetch"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Fetches content from a URL. Returns the response body with a 512KB size limit."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"url".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The URL to fetch"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"headers".to_string(),
|
||||
json!({
|
||||
"type": "object",
|
||||
"description": "Optional HTTP headers"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["url".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let url = input
|
||||
.get("url")
|
||||
.and_then(|u| u.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'url' field".to_string()))?;
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(TIMEOUT_SECS))
|
||||
.build()
|
||||
.map_err(|e| ToolError::ExecutionError(e.to_string()))?;
|
||||
|
||||
let mut req = client.get(url);
|
||||
|
||||
if let Some(headers) = input.get("headers").and_then(|h| h.as_object()) {
|
||||
for (key, value) in headers {
|
||||
if let Some(val) = value.as_str() {
|
||||
req = req.header(key, val);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let response = req
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionError(format!("Fetch failed: {}", e)))?;
|
||||
|
||||
let status = response.status().as_u16();
|
||||
if status >= 400 {
|
||||
return Ok(ToolResult::error(format!(
|
||||
"HTTP {} error fetching {}",
|
||||
status, url
|
||||
)));
|
||||
}
|
||||
|
||||
let body = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionError(format!("Failed to read response: {}", e)))?;
|
||||
|
||||
let mut text = body;
|
||||
if text.len() > MAX_RESPONSE_SIZE {
|
||||
text.truncate(MAX_RESPONSE_SIZE);
|
||||
text.push_str("\n... (response truncated at 512KB)");
|
||||
}
|
||||
|
||||
Ok(ToolResult::text(text))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,123 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// Search result from a web search provider.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SearchResult {
|
||||
pub title: String,
|
||||
pub url: String,
|
||||
pub snippet: String,
|
||||
}
|
||||
|
||||
/// Type alias for a pluggable search function.
|
||||
pub type SearchFn =
|
||||
Arc<dyn Fn(&str, usize) -> futures::future::BoxFuture<'static, Result<Vec<SearchResult>, String>> + Send + Sync>;
|
||||
|
||||
pub struct WebSearchTool {
|
||||
search_fn: Option<SearchFn>,
|
||||
}
|
||||
|
||||
impl Default for WebSearchTool {
|
||||
fn default() -> Self {
|
||||
Self { search_fn: None }
|
||||
}
|
||||
}
|
||||
|
||||
impl WebSearchTool {
|
||||
pub fn new(search_fn: SearchFn) -> Self {
|
||||
Self {
|
||||
search_fn: Some(search_fn),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for WebSearchTool {
|
||||
fn name(&self) -> &str {
|
||||
"WebSearch"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Searches the web for information. Requires a configured search provider."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"query".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"description": "The search query"
|
||||
}),
|
||||
),
|
||||
(
|
||||
"max_results".to_string(),
|
||||
json!({
|
||||
"type": "number",
|
||||
"description": "Maximum number of results (default 5)"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["query".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, _input: &Value) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let query = input
|
||||
.get("query")
|
||||
.and_then(|q| q.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'query' field".to_string()))?;
|
||||
|
||||
let max_results = input
|
||||
.get("max_results")
|
||||
.and_then(|m| m.as_u64())
|
||||
.unwrap_or(5) as usize;
|
||||
|
||||
match &self.search_fn {
|
||||
Some(search_fn) => {
|
||||
let results = (search_fn)(query, max_results)
|
||||
.await
|
||||
.map_err(|e| ToolError::ExecutionError(e))?;
|
||||
|
||||
if results.is_empty() {
|
||||
return Ok(ToolResult::text("No results found.".to_string()));
|
||||
}
|
||||
|
||||
let formatted: Vec<String> = results
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, r)| {
|
||||
format!(
|
||||
"{}. **{}**\n URL: {}\n {}",
|
||||
i + 1,
|
||||
r.title,
|
||||
r.url,
|
||||
r.snippet
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(ToolResult::text(formatted.join("\n\n")))
|
||||
}
|
||||
None => Ok(ToolResult::error(
|
||||
"Web search is not configured. Provide a search function when creating the agent."
|
||||
.to_string(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::HashMap;
|
||||
use std::process::Stdio;
|
||||
use std::sync::Arc;
|
||||
use tokio::process::Command;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::types::{Tool, ToolError, ToolInputSchema, ToolResult, ToolUseContext};
|
||||
|
||||
/// Info about an active worktree.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WorktreeInfo {
|
||||
pub path: String,
|
||||
pub branch: String,
|
||||
pub original_cwd: String,
|
||||
}
|
||||
|
||||
/// Shared worktree tracker.
|
||||
pub type WorktreeStore = Arc<RwLock<HashMap<String, WorktreeInfo>>>;
|
||||
|
||||
pub fn new_worktree_store() -> WorktreeStore {
|
||||
Arc::new(RwLock::new(HashMap::new()))
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// EnterWorktreeTool
|
||||
// ============================================================================
|
||||
|
||||
pub struct EnterWorktreeTool {
|
||||
store: WorktreeStore,
|
||||
}
|
||||
|
||||
impl EnterWorktreeTool {
|
||||
pub fn new(store: WorktreeStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for EnterWorktreeTool {
|
||||
fn name(&self) -> &str {
|
||||
"EnterWorktree"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Create an isolated git worktree for parallel work. The agent will work in the worktree without affecting the main working tree."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"branch".to_string(),
|
||||
json!({ "type": "string", "description": "Branch name for the worktree (auto-generated if not provided)" }),
|
||||
),
|
||||
(
|
||||
"path".to_string(),
|
||||
json!({ "type": "string", "description": "Path for the worktree (auto-generated if not provided)" }),
|
||||
),
|
||||
]),
|
||||
required: Vec::new(),
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
// Check if we're in a git repo
|
||||
let check = Command::new("git")
|
||||
.args(["rev-parse", "--git-dir"])
|
||||
.current_dir(&context.working_dir)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.output()
|
||||
.await;
|
||||
|
||||
if check.is_err() || !check.unwrap().status.success() {
|
||||
return Ok(ToolResult::error("Not in a git repository"));
|
||||
}
|
||||
|
||||
let now = chrono::Utc::now().timestamp_millis();
|
||||
let branch = input
|
||||
.get("branch")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
.unwrap_or_else(|| format!("worktree-{}", now));
|
||||
|
||||
let worktree_path = input
|
||||
.get("path")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from)
|
||||
.unwrap_or_else(|| {
|
||||
let parent = std::path::Path::new(&context.working_dir)
|
||||
.parent()
|
||||
.unwrap_or_else(|| std::path::Path::new("/tmp"));
|
||||
parent
|
||||
.join(format!(".worktree-{}", branch))
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
});
|
||||
|
||||
// Create branch if it doesn't exist (ignore errors if it already exists)
|
||||
let _ = Command::new("git")
|
||||
.args(["branch", &branch])
|
||||
.current_dir(&context.working_dir)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.output()
|
||||
.await;
|
||||
|
||||
// Create worktree
|
||||
let result = Command::new("git")
|
||||
.args(["worktree", "add", &worktree_path, &branch])
|
||||
.current_dir(&context.working_dir)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.output()
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(output) if output.status.success() => {
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
let info = WorktreeInfo {
|
||||
path: worktree_path.clone(),
|
||||
branch: branch.clone(),
|
||||
original_cwd: context.working_dir.clone(),
|
||||
};
|
||||
|
||||
let mut store = self.store.write().await;
|
||||
store.insert(id.clone(), info);
|
||||
|
||||
Ok(ToolResult::text(format!(
|
||||
"Worktree created:\n ID: {}\n Path: {}\n Branch: {}\n\nYou are now working in the isolated worktree.",
|
||||
id, worktree_path, branch
|
||||
)))
|
||||
}
|
||||
Ok(output) => {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
Ok(ToolResult::error(format!(
|
||||
"Error creating worktree: {}",
|
||||
stderr
|
||||
)))
|
||||
}
|
||||
Err(e) => Ok(ToolResult::error(format!(
|
||||
"Error creating worktree: {}",
|
||||
e
|
||||
))),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// ExitWorktreeTool
|
||||
// ============================================================================
|
||||
|
||||
pub struct ExitWorktreeTool {
|
||||
store: WorktreeStore,
|
||||
}
|
||||
|
||||
impl ExitWorktreeTool {
|
||||
pub fn new(store: WorktreeStore) -> Self {
|
||||
Self { store }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for ExitWorktreeTool {
|
||||
fn name(&self) -> &str {
|
||||
"ExitWorktree"
|
||||
}
|
||||
|
||||
fn description(&self) -> &str {
|
||||
"Exit and optionally remove a git worktree. Use \"keep\" to preserve changes or \"remove\" to clean up."
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::from([
|
||||
(
|
||||
"id".to_string(),
|
||||
json!({ "type": "string", "description": "Worktree ID" }),
|
||||
),
|
||||
(
|
||||
"action".to_string(),
|
||||
json!({
|
||||
"type": "string",
|
||||
"enum": ["keep", "remove"],
|
||||
"description": "Whether to keep or remove the worktree (default: remove)"
|
||||
}),
|
||||
),
|
||||
]),
|
||||
required: vec!["id".to_string()],
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn call(&self, input: Value, _context: &ToolUseContext) -> Result<ToolResult, ToolError> {
|
||||
let id = input
|
||||
.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.ok_or_else(|| ToolError::InvalidInput("Missing 'id'".to_string()))?;
|
||||
|
||||
let action = input
|
||||
.get("action")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("remove");
|
||||
|
||||
let mut store = self.store.write().await;
|
||||
let worktree = match store.get(id) {
|
||||
Some(w) => w.clone(),
|
||||
None => return Ok(ToolResult::error(format!("Worktree not found: {}", id))),
|
||||
};
|
||||
|
||||
if action == "remove" {
|
||||
// Remove worktree
|
||||
let result = Command::new("git")
|
||||
.args(["worktree", "remove", &worktree.path, "--force"])
|
||||
.current_dir(&worktree.original_cwd)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.output()
|
||||
.await;
|
||||
|
||||
if let Ok(output) = result {
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
return Ok(ToolResult::error(format!("Error: {}", stderr)));
|
||||
}
|
||||
}
|
||||
|
||||
// Try to delete the branch (ignore errors)
|
||||
let _ = Command::new("git")
|
||||
.args(["branch", "-D", &worktree.branch])
|
||||
.current_dir(&worktree.original_cwd)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.output()
|
||||
.await;
|
||||
}
|
||||
|
||||
store.remove(id);
|
||||
|
||||
Ok(ToolResult::text(format!(
|
||||
"Worktree {}: {}",
|
||||
if action == "remove" {
|
||||
"removed"
|
||||
} else {
|
||||
"kept"
|
||||
},
|
||||
worktree.path
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// MCP server configuration.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "transport")]
|
||||
pub enum McpServerConfig {
|
||||
#[serde(rename = "stdio")]
|
||||
Stdio {
|
||||
command: String,
|
||||
#[serde(default)]
|
||||
args: Vec<String>,
|
||||
#[serde(default)]
|
||||
env: HashMap<String, String>,
|
||||
},
|
||||
#[serde(rename = "sse")]
|
||||
Sse {
|
||||
url: String,
|
||||
#[serde(default)]
|
||||
headers: HashMap<String, String>,
|
||||
},
|
||||
#[serde(rename = "http")]
|
||||
Http {
|
||||
url: String,
|
||||
#[serde(default)]
|
||||
headers: HashMap<String, String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// MCP tool definition received from server.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpToolDefinition {
|
||||
pub name: String,
|
||||
#[serde(default)]
|
||||
pub description: Option<String>,
|
||||
#[serde(default)]
|
||||
pub input_schema: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
/// Result from calling an MCP tool.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct McpToolCallResult {
|
||||
#[serde(default)]
|
||||
pub content: Vec<McpContentBlock>,
|
||||
#[serde(default)]
|
||||
pub is_error: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum McpContentBlock {
|
||||
#[serde(rename = "text")]
|
||||
Text { text: String },
|
||||
#[serde(rename = "image")]
|
||||
Image {
|
||||
data: String,
|
||||
#[serde(rename = "mimeType")]
|
||||
mime_type: String,
|
||||
},
|
||||
}
|
||||
|
||||
/// MCP server connection state.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct McpServerConnection {
|
||||
pub name: String,
|
||||
pub config: McpServerConfig,
|
||||
pub tools: Vec<McpToolDefinition>,
|
||||
pub connected: bool,
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
/// Message role in conversation.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum MessageRole {
|
||||
User,
|
||||
Assistant,
|
||||
}
|
||||
|
||||
/// A message in the conversation history.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Message {
|
||||
pub role: MessageRole,
|
||||
pub content: Vec<ContentBlock>,
|
||||
}
|
||||
|
||||
/// Content block types used in messages.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum ContentBlock {
|
||||
#[serde(rename = "text")]
|
||||
Text { text: String },
|
||||
|
||||
#[serde(rename = "tool_use")]
|
||||
ToolUse {
|
||||
id: String,
|
||||
name: String,
|
||||
input: Value,
|
||||
},
|
||||
|
||||
#[serde(rename = "tool_result")]
|
||||
ToolResult {
|
||||
tool_use_id: String,
|
||||
#[serde(default)]
|
||||
content: Vec<ToolResultContentBlock>,
|
||||
#[serde(default)]
|
||||
is_error: bool,
|
||||
},
|
||||
|
||||
#[serde(rename = "thinking")]
|
||||
Thinking {
|
||||
thinking: String,
|
||||
#[serde(default)]
|
||||
signature: Option<String>,
|
||||
},
|
||||
|
||||
#[serde(rename = "image")]
|
||||
Image { source: ImageContentSource },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ImageContentSource {
|
||||
#[serde(rename = "type")]
|
||||
pub source_type: String,
|
||||
pub media_type: String,
|
||||
pub data: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum ToolResultContentBlock {
|
||||
#[serde(rename = "text")]
|
||||
Text { text: String },
|
||||
#[serde(rename = "image")]
|
||||
Image { source: ImageContentSource },
|
||||
}
|
||||
|
||||
/// System prompt block with optional cache control.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SystemBlock {
|
||||
#[serde(rename = "type")]
|
||||
pub block_type: String,
|
||||
pub text: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub cache_control: Option<CacheControl>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct CacheControl {
|
||||
#[serde(rename = "type")]
|
||||
pub control_type: String,
|
||||
}
|
||||
|
||||
impl CacheControl {
|
||||
pub fn ephemeral() -> Self {
|
||||
Self {
|
||||
control_type: "ephemeral".to_string(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Token usage information from API response.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct Usage {
|
||||
#[serde(default)]
|
||||
pub input_tokens: u64,
|
||||
#[serde(default)]
|
||||
pub output_tokens: u64,
|
||||
#[serde(default)]
|
||||
pub cache_creation_input_tokens: u64,
|
||||
#[serde(default)]
|
||||
pub cache_read_input_tokens: u64,
|
||||
}
|
||||
|
||||
/// SDK message types emitted during the agent loop.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum SDKMessage {
|
||||
#[serde(rename = "system")]
|
||||
System { message: String },
|
||||
|
||||
#[serde(rename = "assistant")]
|
||||
Assistant {
|
||||
message: Message,
|
||||
#[serde(default)]
|
||||
usage: Option<Usage>,
|
||||
},
|
||||
|
||||
#[serde(rename = "tool_result")]
|
||||
ToolResult {
|
||||
tool_use_id: String,
|
||||
tool_name: String,
|
||||
content: String,
|
||||
#[serde(default)]
|
||||
is_error: bool,
|
||||
},
|
||||
|
||||
#[serde(rename = "result")]
|
||||
Result {
|
||||
#[serde(default)]
|
||||
text: String,
|
||||
#[serde(default)]
|
||||
usage: Usage,
|
||||
#[serde(default)]
|
||||
num_turns: u32,
|
||||
#[serde(default)]
|
||||
cost_usd: f64,
|
||||
#[serde(default)]
|
||||
duration_ms: u64,
|
||||
#[serde(default)]
|
||||
messages: Vec<Message>,
|
||||
},
|
||||
|
||||
#[serde(rename = "partial_message")]
|
||||
PartialMessage { text: String },
|
||||
|
||||
#[serde(rename = "compact_boundary")]
|
||||
CompactBoundary {
|
||||
#[serde(default)]
|
||||
summary: String,
|
||||
},
|
||||
|
||||
#[serde(rename = "status")]
|
||||
Status { message: String },
|
||||
|
||||
#[serde(rename = "task_notification")]
|
||||
TaskNotification {
|
||||
task_id: String,
|
||||
status: String,
|
||||
#[serde(default)]
|
||||
message: Option<String>,
|
||||
},
|
||||
|
||||
#[serde(rename = "rate_limit")]
|
||||
RateLimit {
|
||||
#[serde(default)]
|
||||
retry_after_ms: u64,
|
||||
#[serde(default)]
|
||||
message: String,
|
||||
},
|
||||
|
||||
#[serde(rename = "progress")]
|
||||
Progress { message: String },
|
||||
|
||||
#[serde(rename = "tool_start")]
|
||||
ToolStart {
|
||||
tool_use_id: String,
|
||||
tool_name: String,
|
||||
input: Value,
|
||||
},
|
||||
|
||||
#[serde(rename = "permission_request")]
|
||||
PermissionRequest {
|
||||
tool_use_id: String,
|
||||
tool_name: String,
|
||||
input: Value,
|
||||
risk_level: String,
|
||||
},
|
||||
|
||||
#[serde(rename = "error")]
|
||||
Error { message: String },
|
||||
|
||||
/// Real-time thinking delta from the API stream.
|
||||
#[serde(rename = "thinking_delta")]
|
||||
ThinkingDelta { thinking: String },
|
||||
|
||||
/// Real-time text delta from the API stream.
|
||||
#[serde(rename = "text_delta")]
|
||||
TextDelta { text: String },
|
||||
}
|
||||
|
||||
/// Query result returned by the blocking `prompt()` method.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct QueryResult {
|
||||
pub text: String,
|
||||
pub usage: Usage,
|
||||
pub num_turns: u32,
|
||||
pub cost_usd: f64,
|
||||
pub duration_ms: u64,
|
||||
pub messages: Vec<Message>,
|
||||
}
|
||||
|
||||
/// Tool definition in API format (for sending to Claude).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ApiToolParam {
|
||||
pub name: String,
|
||||
pub description: String,
|
||||
pub input_schema: Value,
|
||||
}
|
||||
|
||||
/// Extract text content from a message.
|
||||
pub fn extract_text(message: &Message) -> String {
|
||||
message
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::Text { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("")
|
||||
}
|
||||
|
||||
/// Extract tool use blocks from a message.
|
||||
pub fn extract_tool_uses(message: &Message) -> Vec<(String, String, Value)> {
|
||||
message
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::ToolUse { id, name, input } => {
|
||||
Some((id.clone(), name.clone(), input.clone()))
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Thinking configuration for extended thinking.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ThinkingConfig {
|
||||
#[serde(rename = "type")]
|
||||
pub thinking_type: String,
|
||||
#[serde(default)]
|
||||
pub budget_tokens: Option<u64>,
|
||||
}
|
||||
|
||||
impl ThinkingConfig {
|
||||
pub fn enabled(budget_tokens: u64) -> Self {
|
||||
Self {
|
||||
thinking_type: "enabled".to_string(),
|
||||
budget_tokens: Some(budget_tokens),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn disabled() -> Self {
|
||||
Self {
|
||||
thinking_type: "disabled".to_string(),
|
||||
budget_tokens: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
pub mod message;
|
||||
pub mod mcp;
|
||||
pub mod skill;
|
||||
pub mod tool;
|
||||
|
||||
pub use message::*;
|
||||
pub use mcp::*;
|
||||
pub use skill::*;
|
||||
pub use tool::*;
|
||||
@@ -0,0 +1,95 @@
|
||||
use serde::{Deserialize, Deserializer, Serialize};
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// SKILL.md YAML frontmatter metadata.
|
||||
/// Uses kebab-case in YAML, mapped to snake_case in Rust via serde rename_all.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub struct SkillMetadata {
|
||||
pub name: String,
|
||||
#[serde(default)]
|
||||
pub description: Option<String>,
|
||||
#[serde(default)]
|
||||
pub license: Option<String>,
|
||||
#[serde(default)]
|
||||
pub metadata: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub argument_hint: Option<String>,
|
||||
#[serde(default)]
|
||||
pub when_to_use: Option<String>,
|
||||
#[serde(default)]
|
||||
pub disable_model_invocation: bool,
|
||||
#[serde(default = "default_true")]
|
||||
pub user_invocable: bool,
|
||||
#[serde(default, deserialize_with = "deserialize_allowed_tools")]
|
||||
pub allowed_tools: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub context: Option<String>,
|
||||
#[serde(default)]
|
||||
pub agent: Option<String>,
|
||||
}
|
||||
|
||||
fn default_true() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Custom deserializer for allowed-tools field.
|
||||
/// Supports YAML array or space-separated string (Claude Code compat).
|
||||
fn deserialize_allowed_tools<'de, D>(deserializer: D) -> Result<Option<Vec<String>>, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
{
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum AllowedToolsFormat {
|
||||
Array(Vec<String>),
|
||||
SpaceSeparated(String),
|
||||
}
|
||||
|
||||
let opt = Option::<AllowedToolsFormat>::deserialize(deserializer)?;
|
||||
match opt {
|
||||
None => Ok(None),
|
||||
Some(AllowedToolsFormat::Array(v)) => Ok(Some(v)),
|
||||
Some(AllowedToolsFormat::SpaceSeparated(s)) => {
|
||||
let trimmed = s.trim();
|
||||
if trimmed.is_empty() {
|
||||
Ok(None)
|
||||
} else {
|
||||
Ok(Some(trimmed.split_whitespace().map(String::from).collect()))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Where a skill was loaded from (priority order, highest first).
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
pub enum SkillSource {
|
||||
BoosApi,
|
||||
Claude,
|
||||
Agents,
|
||||
Project,
|
||||
}
|
||||
|
||||
impl SkillSource {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
SkillSource::BoosApi => "boosapi",
|
||||
SkillSource::Claude => "claude",
|
||||
SkillSource::Agents => "agents",
|
||||
SkillSource::Project => "project",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A fully loaded skill with parsed metadata and content.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LoadedSkill {
|
||||
pub name: String,
|
||||
pub path: PathBuf,
|
||||
pub metadata: SkillMetadata,
|
||||
/// Markdown body (after frontmatter)
|
||||
pub content: String,
|
||||
pub source: SkillSource,
|
||||
/// If this skill is part of a group/collection, the group name
|
||||
pub group: Option<String>,
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
use async_trait::async_trait;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::collections::HashMap;
|
||||
use std::fmt;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// Tool input schema following JSON Schema format.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolInputSchema {
|
||||
#[serde(rename = "type")]
|
||||
pub schema_type: String,
|
||||
#[serde(default)]
|
||||
pub properties: HashMap<String, Value>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub required: Vec<String>,
|
||||
#[serde(default, rename = "additionalProperties")]
|
||||
pub additional_properties: Option<bool>,
|
||||
}
|
||||
|
||||
impl Default for ToolInputSchema {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
schema_type: "object".to_string(),
|
||||
properties: HashMap::new(),
|
||||
required: Vec::new(),
|
||||
additional_properties: Some(false),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Result returned by a tool execution.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ToolResult {
|
||||
#[serde(default)]
|
||||
pub content: Vec<ToolResultContent>,
|
||||
#[serde(default)]
|
||||
pub is_error: bool,
|
||||
}
|
||||
|
||||
impl ToolResult {
|
||||
pub fn text(text: impl Into<String>) -> Self {
|
||||
Self {
|
||||
content: vec![ToolResultContent::Text { text: text.into() }],
|
||||
is_error: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn error(text: impl Into<String>) -> Self {
|
||||
Self {
|
||||
content: vec![ToolResultContent::Text { text: text.into() }],
|
||||
is_error: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn image(data: String, media_type: String) -> Self {
|
||||
Self {
|
||||
content: vec![ToolResultContent::Image {
|
||||
source: ImageSource {
|
||||
source_type: "base64".to_string(),
|
||||
media_type,
|
||||
data,
|
||||
},
|
||||
}],
|
||||
is_error: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn get_text(&self) -> String {
|
||||
self.content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
ToolResultContent::Text { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n")
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type")]
|
||||
pub enum ToolResultContent {
|
||||
#[serde(rename = "text")]
|
||||
Text { text: String },
|
||||
#[serde(rename = "image")]
|
||||
Image { source: ImageSource },
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct ImageSource {
|
||||
#[serde(rename = "type")]
|
||||
pub source_type: String,
|
||||
pub media_type: String,
|
||||
pub data: String,
|
||||
}
|
||||
|
||||
/// Context passed to tools during execution.
|
||||
#[derive(Clone)]
|
||||
pub struct ToolUseContext {
|
||||
pub working_dir: String,
|
||||
pub abort_signal: tokio_util::sync::CancellationToken,
|
||||
pub read_file_state: Arc<RwLock<HashMap<String, String>>>,
|
||||
}
|
||||
|
||||
impl ToolUseContext {
|
||||
pub fn new(working_dir: String) -> Self {
|
||||
Self::with_abort(working_dir, tokio_util::sync::CancellationToken::new())
|
||||
}
|
||||
|
||||
pub fn with_abort(
|
||||
working_dir: String,
|
||||
abort_signal: tokio_util::sync::CancellationToken,
|
||||
) -> Self {
|
||||
Self {
|
||||
working_dir,
|
||||
abort_signal,
|
||||
read_file_state: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Debug for ToolUseContext {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
f.debug_struct("ToolUseContext")
|
||||
.field("working_dir", &self.working_dir)
|
||||
.finish()
|
||||
}
|
||||
}
|
||||
|
||||
/// Permission decision for tool execution.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub enum PermissionDecision {
|
||||
Allow,
|
||||
Deny(String),
|
||||
AllowWithModifiedInput(Value),
|
||||
}
|
||||
|
||||
/// Callback function for custom permission logic (async).
|
||||
pub type CanUseToolFn = Arc<
|
||||
dyn Fn(
|
||||
&str,
|
||||
&Value,
|
||||
)
|
||||
-> std::pin::Pin<Box<dyn std::future::Future<Output = PermissionDecision> + Send>>
|
||||
+ Send
|
||||
+ Sync,
|
||||
>;
|
||||
|
||||
/// The core trait that all tools must implement.
|
||||
#[async_trait]
|
||||
pub trait Tool: Send + Sync {
|
||||
fn name(&self) -> &str;
|
||||
fn description(&self) -> &str;
|
||||
fn input_schema(&self) -> ToolInputSchema;
|
||||
|
||||
async fn call(&self, input: Value, context: &ToolUseContext) -> Result<ToolResult, ToolError>;
|
||||
|
||||
fn is_read_only(&self, _input: &Value) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
fn is_concurrency_safe(&self, input: &Value) -> bool {
|
||||
self.is_read_only(input)
|
||||
}
|
||||
}
|
||||
|
||||
/// Errors that can occur during tool execution.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ToolError {
|
||||
#[error("Tool execution failed: {0}")]
|
||||
ExecutionError(String),
|
||||
|
||||
#[error("Invalid input: {0}")]
|
||||
InvalidInput(String),
|
||||
|
||||
#[error("Permission denied: {0}")]
|
||||
PermissionDenied(String),
|
||||
|
||||
#[error("IO error: {0}")]
|
||||
IoError(#[from] std::io::Error),
|
||||
|
||||
#[error("Serialization error: {0}")]
|
||||
SerializationError(#[from] serde_json::Error),
|
||||
|
||||
#[error("Tool aborted")]
|
||||
Aborted,
|
||||
}
|
||||
|
||||
/// Sandbox configuration for restricting tool execution.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct SandboxSettings {
|
||||
/// Whether sandboxing is enabled.
|
||||
#[serde(default)]
|
||||
pub enabled: bool,
|
||||
/// Network restrictions.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub network: Option<SandboxNetworkConfig>,
|
||||
/// Filesystem restrictions.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub filesystem: Option<SandboxFilesystemConfig>,
|
||||
}
|
||||
|
||||
/// Network configuration for the sandbox.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct SandboxNetworkConfig {
|
||||
/// Domains allowed for network access.
|
||||
#[serde(default)]
|
||||
pub allowed_domains: Vec<String>,
|
||||
}
|
||||
|
||||
/// Filesystem configuration for the sandbox.
|
||||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||||
pub struct SandboxFilesystemConfig {
|
||||
/// Paths where writes are denied.
|
||||
#[serde(default)]
|
||||
pub deny_write: Vec<String>,
|
||||
/// Paths where reads are denied.
|
||||
#[serde(default)]
|
||||
pub deny_read: Vec<String>,
|
||||
}
|
||||
|
||||
/// Permission mode for controlling tool access.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum PermissionMode {
|
||||
Default,
|
||||
AcceptEdits,
|
||||
BypassPermissions,
|
||||
Plan,
|
||||
DontAsk,
|
||||
Auto,
|
||||
}
|
||||
|
||||
impl Default for PermissionMode {
|
||||
fn default() -> Self {
|
||||
Self::BypassPermissions
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,222 @@
|
||||
use crate::types::{ContentBlock, Message};
|
||||
use crate::utils::tokens::{estimate_messages_tokens, get_auto_compact_threshold};
|
||||
|
||||
const MICRO_COMPACT_THRESHOLD: usize = 50_000; // chars per tool result
|
||||
|
||||
/// Check if auto-compaction is needed based on message token count.
|
||||
pub fn should_auto_compact(messages: &[Message], model: &str) -> bool {
|
||||
let estimated = estimate_messages_tokens(messages);
|
||||
let threshold = get_auto_compact_threshold(model);
|
||||
estimated > threshold
|
||||
}
|
||||
|
||||
/// Micro-compact: truncate large tool results in messages.
|
||||
pub fn micro_compact_messages(messages: &[Message]) -> Vec<Message> {
|
||||
messages
|
||||
.iter()
|
||||
.map(|msg| {
|
||||
let content = msg
|
||||
.content
|
||||
.iter()
|
||||
.map(|block| match block {
|
||||
crate::types::ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
} => {
|
||||
let compacted_content: Vec<crate::types::ToolResultContentBlock> = content
|
||||
.iter()
|
||||
.map(|c| match c {
|
||||
crate::types::ToolResultContentBlock::Text { text } => {
|
||||
if text.len() > MICRO_COMPACT_THRESHOLD {
|
||||
let truncated =
|
||||
format!("{}... (truncated)", &text[..MICRO_COMPACT_THRESHOLD]);
|
||||
crate::types::ToolResultContentBlock::Text {
|
||||
text: truncated,
|
||||
}
|
||||
} else {
|
||||
c.clone()
|
||||
}
|
||||
}
|
||||
_ => c.clone(),
|
||||
})
|
||||
.collect();
|
||||
crate::types::ContentBlock::ToolResult {
|
||||
tool_use_id: tool_use_id.clone(),
|
||||
content: compacted_content,
|
||||
is_error: *is_error,
|
||||
}
|
||||
}
|
||||
_ => block.clone(),
|
||||
})
|
||||
.collect();
|
||||
Message {
|
||||
role: msg.role.clone(),
|
||||
content,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Create a compact summary prompt for the LLM to summarize the conversation.
|
||||
pub fn create_compact_prompt(messages: &[Message]) -> String {
|
||||
let message_count = messages.len();
|
||||
format!(
|
||||
"Please provide a concise summary of the conversation so far ({} messages). \
|
||||
Focus on: 1) What the user asked for, 2) What has been accomplished, \
|
||||
3) Key decisions made, 4) Current state and any pending work. \
|
||||
Keep the summary under 2000 tokens.",
|
||||
message_count
|
||||
)
|
||||
}
|
||||
|
||||
/// Build the full compaction system prompt for the LLM summarizer.
|
||||
///
|
||||
/// This produces a system prompt that instructs the LLM to produce a structured
|
||||
/// summary suitable for replacing the full conversation history.
|
||||
pub fn build_compaction_prompt() -> String {
|
||||
r#"You are a conversation compaction assistant. Your job is to produce a concise, structured summary of a conversation between a user and an AI coding assistant.
|
||||
|
||||
The summary will REPLACE the original conversation, so it must preserve all information needed to continue the work seamlessly.
|
||||
|
||||
Output format:
|
||||
1. **User Goal**: What the user originally asked for (1-2 sentences).
|
||||
2. **Completed Work**: Bullet list of what has been done (files created/edited, commands run, decisions made).
|
||||
3. **Current State**: The current state of the task - what's working, what's not.
|
||||
4. **Pending Work**: Any remaining tasks or next steps that were discussed but not completed.
|
||||
5. **Key Context**: Important facts, variable names, file paths, or constraints that would be needed to continue.
|
||||
|
||||
Rules:
|
||||
- Be concise but complete. Do not omit file paths, function names, or error messages that are relevant.
|
||||
- Do not include tool call details or raw tool output - summarize the results instead.
|
||||
- Keep the total summary under 2000 tokens.
|
||||
- If the conversation is already short, just summarize it briefly."#
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// Perform full LLM-based conversation compaction.
|
||||
///
|
||||
/// Takes the conversation messages and returns a compacted conversation consisting of
|
||||
/// a single user message (with the compaction prompt) ready to be sent to the LLM for
|
||||
/// summarization. The caller is responsible for actually calling the LLM and using
|
||||
/// the resulting summary.
|
||||
///
|
||||
/// Returns `(system_prompt, user_message)` where:
|
||||
/// - `system_prompt` is the compaction system prompt
|
||||
/// - `user_message` is the serialized conversation for the LLM to summarize
|
||||
pub fn compact_conversation(messages: &[Message]) -> (String, String) {
|
||||
let system_prompt = build_compaction_prompt();
|
||||
|
||||
// Serialize the conversation into a readable text format for the summarizer
|
||||
let mut conversation_text = String::new();
|
||||
for (i, msg) in messages.iter().enumerate() {
|
||||
let role = match &msg.role {
|
||||
crate::types::MessageRole::User => "User",
|
||||
crate::types::MessageRole::Assistant => "Assistant",
|
||||
};
|
||||
conversation_text.push_str(&format!("--- Message {} ({}) ---\n", i + 1, role));
|
||||
for block in &msg.content {
|
||||
match block {
|
||||
ContentBlock::Text { text } => {
|
||||
// Truncate very long text blocks in the compaction input
|
||||
if text.len() > 5000 {
|
||||
conversation_text.push_str(&text[..5000]);
|
||||
conversation_text.push_str("... [truncated]\n");
|
||||
} else {
|
||||
conversation_text.push_str(text);
|
||||
conversation_text.push('\n');
|
||||
}
|
||||
}
|
||||
ContentBlock::ToolUse { name, .. } => {
|
||||
conversation_text.push_str(&format!("[Tool call: {}]\n", name));
|
||||
}
|
||||
ContentBlock::ToolResult { is_error, content, .. } => {
|
||||
let preview = content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
crate::types::ToolResultContentBlock::Text { text } => {
|
||||
if text.len() > 500 {
|
||||
Some(format!("{}...", &text[..500]))
|
||||
} else {
|
||||
Some(text.clone())
|
||||
}
|
||||
}
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("\n");
|
||||
let error_label = if *is_error { " (error)" } else { "" };
|
||||
conversation_text
|
||||
.push_str(&format!("[Tool result{}]: {}\n", error_label, preview));
|
||||
}
|
||||
ContentBlock::Thinking { thinking, .. } => {
|
||||
if thinking.len() > 200 {
|
||||
conversation_text
|
||||
.push_str(&format!("[Thinking]: {}...\n", &thinking[..200]));
|
||||
} else {
|
||||
conversation_text.push_str(&format!("[Thinking]: {}\n", thinking));
|
||||
}
|
||||
}
|
||||
ContentBlock::Image { .. } => {
|
||||
conversation_text.push_str("[Image content]\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
conversation_text.push('\n');
|
||||
}
|
||||
|
||||
let user_message = format!(
|
||||
"Please summarize the following conversation:\n\n{}",
|
||||
conversation_text
|
||||
);
|
||||
|
||||
(system_prompt, user_message)
|
||||
}
|
||||
|
||||
/// Strip image blocks from messages, replacing them with a placeholder text block.
|
||||
///
|
||||
/// This is useful before compaction or when images are no longer needed and consume
|
||||
/// excessive tokens in the context window.
|
||||
pub fn strip_images_from_messages(messages: &[Message]) -> Vec<Message> {
|
||||
messages
|
||||
.iter()
|
||||
.map(|msg| {
|
||||
let content = msg
|
||||
.content
|
||||
.iter()
|
||||
.map(|block| match block {
|
||||
ContentBlock::Image { .. } => ContentBlock::Text {
|
||||
text: "[image removed during compaction]".to_string(),
|
||||
},
|
||||
ContentBlock::ToolResult {
|
||||
tool_use_id,
|
||||
content,
|
||||
is_error,
|
||||
} => {
|
||||
let filtered_content: Vec<crate::types::ToolResultContentBlock> = content
|
||||
.iter()
|
||||
.map(|c| match c {
|
||||
crate::types::ToolResultContentBlock::Image { .. } => {
|
||||
crate::types::ToolResultContentBlock::Text {
|
||||
text: "[image removed during compaction]".to_string(),
|
||||
}
|
||||
}
|
||||
other => other.clone(),
|
||||
})
|
||||
.collect();
|
||||
ContentBlock::ToolResult {
|
||||
tool_use_id: tool_use_id.clone(),
|
||||
content: filtered_content,
|
||||
is_error: *is_error,
|
||||
}
|
||||
}
|
||||
other => other.clone(),
|
||||
})
|
||||
.collect();
|
||||
Message {
|
||||
role: msg.role.clone(),
|
||||
content,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Default maximum number of cache entries.
|
||||
const DEFAULT_MAX_ENTRIES: usize = 100;
|
||||
|
||||
/// Default maximum total cache size in bytes (25 MB).
|
||||
const DEFAULT_MAX_SIZE: usize = 25 * 1024 * 1024;
|
||||
|
||||
/// A cached file entry with content and access metadata.
|
||||
#[derive(Debug, Clone)]
|
||||
struct CacheEntry {
|
||||
content: String,
|
||||
size: usize,
|
||||
last_access: u64,
|
||||
}
|
||||
|
||||
/// A simple LRU file cache for caching file contents in memory.
|
||||
///
|
||||
/// Evicts least-recently-used entries when max_entries or max_size is exceeded.
|
||||
pub struct FileStateCache {
|
||||
entries: HashMap<String, CacheEntry>,
|
||||
max_entries: usize,
|
||||
max_size: usize,
|
||||
current_size: usize,
|
||||
access_counter: u64,
|
||||
}
|
||||
|
||||
impl FileStateCache {
|
||||
/// Create a new file cache with default limits (100 entries, 25 MB).
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
entries: HashMap::new(),
|
||||
max_entries: DEFAULT_MAX_ENTRIES,
|
||||
max_size: DEFAULT_MAX_SIZE,
|
||||
current_size: 0,
|
||||
access_counter: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new file cache with custom limits.
|
||||
pub fn with_limits(max_entries: usize, max_size: usize) -> Self {
|
||||
Self {
|
||||
entries: HashMap::new(),
|
||||
max_entries,
|
||||
max_size,
|
||||
current_size: 0,
|
||||
access_counter: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a cached file's content by path. Returns None if not cached.
|
||||
pub fn get(&mut self, path: &str) -> Option<&str> {
|
||||
self.access_counter += 1;
|
||||
let counter = self.access_counter;
|
||||
if let Some(entry) = self.entries.get_mut(path) {
|
||||
entry.last_access = counter;
|
||||
Some(&entry.content)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Cache file content for a given path. Evicts LRU entries if limits are exceeded.
|
||||
pub fn set(&mut self, path: String, content: String) {
|
||||
let size = content.len();
|
||||
|
||||
// Remove existing entry for same path first
|
||||
if let Some(old) = self.entries.remove(&path) {
|
||||
self.current_size -= old.size;
|
||||
}
|
||||
|
||||
// Evict until we have room
|
||||
while self.entries.len() >= self.max_entries
|
||||
|| (self.current_size + size > self.max_size && !self.entries.is_empty())
|
||||
{
|
||||
self.evict_lru();
|
||||
}
|
||||
|
||||
// If a single entry exceeds max_size, don't cache it
|
||||
if size > self.max_size {
|
||||
return;
|
||||
}
|
||||
|
||||
self.access_counter += 1;
|
||||
self.entries.insert(
|
||||
path,
|
||||
CacheEntry {
|
||||
content,
|
||||
size,
|
||||
last_access: self.access_counter,
|
||||
},
|
||||
);
|
||||
self.current_size += size;
|
||||
}
|
||||
|
||||
/// Remove a cached entry by path.
|
||||
pub fn delete(&mut self, path: &str) -> bool {
|
||||
if let Some(entry) = self.entries.remove(path) {
|
||||
self.current_size -= entry.size;
|
||||
true
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Clear all cached entries.
|
||||
pub fn clear(&mut self) {
|
||||
self.entries.clear();
|
||||
self.current_size = 0;
|
||||
self.access_counter = 0;
|
||||
}
|
||||
|
||||
/// Number of entries currently cached.
|
||||
pub fn len(&self) -> usize {
|
||||
self.entries.len()
|
||||
}
|
||||
|
||||
/// Whether the cache is empty.
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.entries.is_empty()
|
||||
}
|
||||
|
||||
/// Total size of cached content in bytes.
|
||||
pub fn total_size(&self) -> usize {
|
||||
self.current_size
|
||||
}
|
||||
|
||||
/// Evict the least recently used entry.
|
||||
fn evict_lru(&mut self) {
|
||||
if self.entries.is_empty() {
|
||||
return;
|
||||
}
|
||||
|
||||
let lru_key = self
|
||||
.entries
|
||||
.iter()
|
||||
.min_by_key(|(_, entry)| entry.last_access)
|
||||
.map(|(key, _)| key.clone());
|
||||
|
||||
if let Some(key) = lru_key {
|
||||
if let Some(entry) = self.entries.remove(&key) {
|
||||
self.current_size -= entry.size;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for FileStateCache {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_basic_get_set() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set("/foo/bar.rs".to_string(), "hello world".to_string());
|
||||
assert_eq!(cache.get("/foo/bar.rs"), Some("hello world"));
|
||||
assert_eq!(cache.get("/nonexistent"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_delete() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set("/a.rs".to_string(), "content".to_string());
|
||||
assert!(cache.delete("/a.rs"));
|
||||
assert!(!cache.delete("/a.rs"));
|
||||
assert_eq!(cache.get("/a.rs"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clear() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set("/a.rs".to_string(), "a".to_string());
|
||||
cache.set("/b.rs".to_string(), "b".to_string());
|
||||
assert_eq!(cache.len(), 2);
|
||||
cache.clear();
|
||||
assert!(cache.is_empty());
|
||||
assert_eq!(cache.total_size(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lru_eviction_by_entries() {
|
||||
let mut cache = FileStateCache::with_limits(2, 1024 * 1024);
|
||||
cache.set("/a.rs".to_string(), "aaa".to_string());
|
||||
cache.set("/b.rs".to_string(), "bbb".to_string());
|
||||
// Access /a.rs to make it more recent
|
||||
cache.get("/a.rs");
|
||||
// Adding a third entry should evict /b.rs (least recently used)
|
||||
cache.set("/c.rs".to_string(), "ccc".to_string());
|
||||
assert_eq!(cache.len(), 2);
|
||||
assert_eq!(cache.get("/b.rs"), None);
|
||||
assert!(cache.get("/a.rs").is_some());
|
||||
assert!(cache.get("/c.rs").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lru_eviction_by_size() {
|
||||
let mut cache = FileStateCache::with_limits(100, 10);
|
||||
cache.set("a".to_string(), "12345".to_string()); // 5 bytes
|
||||
cache.set("b".to_string(), "12345".to_string()); // 5 bytes, total = 10
|
||||
cache.set("c".to_string(), "12345".to_string()); // would exceed 10, evicts LRU
|
||||
assert_eq!(cache.len(), 2);
|
||||
assert!(cache.total_size() <= 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_oversize_entry_rejected() {
|
||||
let mut cache = FileStateCache::with_limits(100, 5);
|
||||
cache.set("big".to_string(), "123456".to_string()); // 6 bytes > max 5
|
||||
assert!(cache.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_update_existing_key() {
|
||||
let mut cache = FileStateCache::new();
|
||||
cache.set("f".to_string(), "old".to_string());
|
||||
cache.set("f".to_string(), "new content".to_string());
|
||||
assert_eq!(cache.get("f"), Some("new content"));
|
||||
assert_eq!(cache.len(), 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
use crate::types::{ContentBlock, Message, MessageRole};
|
||||
|
||||
/// Create a user message with text content.
|
||||
pub fn create_user_message(text: &str) -> Message {
|
||||
Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: text.to_string(),
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an assistant message with text content.
|
||||
pub fn create_assistant_message(text: &str) -> Message {
|
||||
Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: text.to_string(),
|
||||
}],
|
||||
}
|
||||
}
|
||||
|
||||
/// Ensure messages alternate between user and assistant roles.
|
||||
/// Merges consecutive messages with the same role.
|
||||
pub fn normalize_messages(messages: &[Message]) -> Vec<Message> {
|
||||
if messages.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut normalized = Vec::new();
|
||||
|
||||
for msg in messages {
|
||||
if let Some(last) = normalized.last_mut() {
|
||||
let last: &mut Message = last;
|
||||
if last.role == msg.role {
|
||||
// Merge content blocks
|
||||
last.content.extend(msg.content.clone());
|
||||
continue;
|
||||
}
|
||||
}
|
||||
normalized.push(msg.clone());
|
||||
}
|
||||
|
||||
// Ensure first message is from user
|
||||
if !normalized.is_empty() && normalized[0].role != MessageRole::User {
|
||||
normalized.insert(
|
||||
0,
|
||||
Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "(continuing conversation)".to_string(),
|
||||
}],
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
normalized
|
||||
}
|
||||
|
||||
/// Extract text from a message's content blocks.
|
||||
pub fn extract_text(message: &Message) -> String {
|
||||
message
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|block| match block {
|
||||
ContentBlock::Text { text } => Some(text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("")
|
||||
}
|
||||
|
||||
/// Strip image content blocks from messages (for compaction).
|
||||
pub fn strip_images(messages: &[Message]) -> Vec<Message> {
|
||||
messages
|
||||
.iter()
|
||||
.map(|msg| {
|
||||
let content: Vec<ContentBlock> = msg
|
||||
.content
|
||||
.iter()
|
||||
.filter(|block| !matches!(block, ContentBlock::Image { .. }))
|
||||
.cloned()
|
||||
.collect();
|
||||
Message {
|
||||
role: msg.role.clone(),
|
||||
content,
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Truncate text to a maximum length.
|
||||
pub fn truncate_text(text: &str, max_len: usize) -> String {
|
||||
if text.len() <= max_len {
|
||||
text.to_string()
|
||||
} else {
|
||||
format!("{}... (truncated)", &text[..max_len])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
pub mod compact;
|
||||
pub mod file_cache;
|
||||
pub mod messages;
|
||||
pub mod retry;
|
||||
pub mod tokens;
|
||||
@@ -0,0 +1,103 @@
|
||||
use crate::api::ApiError;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Retry configuration.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RetryConfig {
|
||||
pub max_retries: u32,
|
||||
pub base_delay_ms: u64,
|
||||
pub max_delay_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for RetryConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
max_retries: 3,
|
||||
base_delay_ms: 2000,
|
||||
max_delay_ms: 30_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Check if an API error is retryable.
|
||||
pub fn is_retryable(error: &ApiError) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
ApiError::RateLimitError
|
||||
| ApiError::Timeout
|
||||
| ApiError::NetworkError(_)
|
||||
| ApiError::HttpError { status: 500..=599, .. }
|
||||
)
|
||||
}
|
||||
|
||||
/// Check if an error indicates the prompt is too long.
|
||||
pub fn is_prompt_too_long(error: &ApiError) -> bool {
|
||||
matches!(error, ApiError::PromptTooLong(_))
|
||||
}
|
||||
|
||||
/// Check if an error is an auth error.
|
||||
pub fn is_auth_error(error: &ApiError) -> bool {
|
||||
matches!(error, ApiError::AuthError(_))
|
||||
}
|
||||
|
||||
/// Calculate retry delay with exponential backoff and jitter.
|
||||
pub fn get_retry_delay(config: &RetryConfig, attempt: u32) -> Duration {
|
||||
let base = config.base_delay_ms as f64;
|
||||
let delay = base * 2.0_f64.powi(attempt as i32);
|
||||
let max = config.max_delay_ms as f64;
|
||||
let clamped = delay.min(max);
|
||||
|
||||
// Add jitter (0.5 to 1.5 of calculated delay)
|
||||
let jitter = 0.5 + rand_f64() * 1.0;
|
||||
let final_delay = (clamped * jitter) as u64;
|
||||
|
||||
Duration::from_millis(final_delay)
|
||||
}
|
||||
|
||||
/// Execute an async function with retry logic.
|
||||
pub async fn with_retry<F, Fut, T>(
|
||||
config: &RetryConfig,
|
||||
mut f: F,
|
||||
) -> Result<T, ApiError>
|
||||
where
|
||||
F: FnMut() -> Fut,
|
||||
Fut: std::future::Future<Output = Result<T, ApiError>>,
|
||||
{
|
||||
let mut last_error = None;
|
||||
|
||||
for attempt in 0..=config.max_retries {
|
||||
match f().await {
|
||||
Ok(result) => return Ok(result),
|
||||
Err(error) => {
|
||||
if !is_retryable(&error) || attempt == config.max_retries {
|
||||
return Err(error);
|
||||
}
|
||||
|
||||
let delay = get_retry_delay(config, attempt);
|
||||
tracing::warn!(
|
||||
"Retrying after error (attempt {}/{}): {}. Waiting {:?}",
|
||||
attempt + 1,
|
||||
config.max_retries,
|
||||
error,
|
||||
delay
|
||||
);
|
||||
tokio::time::sleep(delay).await;
|
||||
last_error = Some(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Err(last_error.unwrap_or(ApiError::NetworkError(
|
||||
"Max retries exceeded".to_string(),
|
||||
)))
|
||||
}
|
||||
|
||||
/// Simple pseudo-random f64 in [0, 1) using time-based seed.
|
||||
fn rand_f64() -> f64 {
|
||||
use std::time::SystemTime;
|
||||
let seed = SystemTime::now()
|
||||
.duration_since(SystemTime::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.subsec_nanos();
|
||||
(seed as f64 % 1000.0) / 1000.0
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
use crate::types::{Message, Usage};
|
||||
|
||||
/// Estimate token count from text (~4 characters per token).
|
||||
pub fn estimate_tokens(text: &str) -> u64 {
|
||||
(text.len() as f64 / 4.0).ceil() as u64
|
||||
}
|
||||
|
||||
/// Estimate total tokens for a list of messages.
|
||||
pub fn estimate_messages_tokens(messages: &[Message]) -> u64 {
|
||||
messages
|
||||
.iter()
|
||||
.map(|m| {
|
||||
let content_str = serde_json::to_string(&m.content).unwrap_or_default();
|
||||
estimate_tokens(&content_str) + 4 // overhead per message
|
||||
})
|
||||
.sum()
|
||||
}
|
||||
|
||||
/// Get context window size for a model.
|
||||
pub fn get_context_window_size(model: &str) -> u64 {
|
||||
match model {
|
||||
// Anthropic models
|
||||
m if m.contains("opus") && m.contains("1m") => 1_000_000,
|
||||
m if m.contains("opus") => 200_000,
|
||||
m if m.contains("sonnet") => 200_000,
|
||||
m if m.contains("haiku") => 200_000,
|
||||
// OpenAI models
|
||||
m if m.starts_with("o1") => 200_000,
|
||||
m if m.starts_with("o3") => 200_000,
|
||||
m if m.starts_with("o4-mini") => 200_000,
|
||||
m if m.starts_with("gpt-4o-mini") => 128_000,
|
||||
m if m.starts_with("gpt-4o") => 128_000,
|
||||
// DeepSeek models
|
||||
m if m.starts_with("deepseek") => 128_000,
|
||||
_ => 200_000,
|
||||
}
|
||||
}
|
||||
|
||||
/// Get auto-compact threshold for a model (context window - 13k tokens).
|
||||
pub fn get_auto_compact_threshold(model: &str) -> u64 {
|
||||
let window = get_context_window_size(model);
|
||||
window.saturating_sub(13_000)
|
||||
}
|
||||
|
||||
/// Model pricing per million tokens.
|
||||
pub struct ModelPricing {
|
||||
pub input: f64,
|
||||
pub output: f64,
|
||||
pub cache_read: f64,
|
||||
pub cache_write: f64,
|
||||
}
|
||||
|
||||
/// Get pricing for a model.
|
||||
pub fn get_model_pricing(model: &str) -> ModelPricing {
|
||||
match model {
|
||||
// Anthropic models
|
||||
m if m.contains("opus") => ModelPricing {
|
||||
input: 15.0,
|
||||
output: 75.0,
|
||||
cache_read: 1.5,
|
||||
cache_write: 18.75,
|
||||
},
|
||||
m if m.contains("sonnet") => ModelPricing {
|
||||
input: 3.0,
|
||||
output: 15.0,
|
||||
cache_read: 0.3,
|
||||
cache_write: 3.75,
|
||||
},
|
||||
m if m.contains("haiku") => ModelPricing {
|
||||
input: 0.8,
|
||||
output: 4.0,
|
||||
cache_read: 0.08,
|
||||
cache_write: 1.0,
|
||||
},
|
||||
// OpenAI models
|
||||
m if m.starts_with("gpt-4o-mini") => ModelPricing {
|
||||
input: 0.15,
|
||||
output: 0.60,
|
||||
cache_read: 0.075,
|
||||
cache_write: 0.15,
|
||||
},
|
||||
m if m.starts_with("gpt-4o") => ModelPricing {
|
||||
input: 2.50,
|
||||
output: 10.0,
|
||||
cache_read: 1.25,
|
||||
cache_write: 2.50,
|
||||
},
|
||||
m if m.starts_with("o1") => ModelPricing {
|
||||
input: 15.0,
|
||||
output: 60.0,
|
||||
cache_read: 7.5,
|
||||
cache_write: 15.0,
|
||||
},
|
||||
m if m.starts_with("o3") => ModelPricing {
|
||||
input: 10.0,
|
||||
output: 40.0,
|
||||
cache_read: 5.0,
|
||||
cache_write: 10.0,
|
||||
},
|
||||
m if m.starts_with("o4-mini") => ModelPricing {
|
||||
input: 1.10,
|
||||
output: 4.40,
|
||||
cache_read: 0.55,
|
||||
cache_write: 1.10,
|
||||
},
|
||||
// DeepSeek models
|
||||
m if m.starts_with("deepseek-reasoner") => ModelPricing {
|
||||
input: 0.55,
|
||||
output: 2.19,
|
||||
cache_read: 0.14,
|
||||
cache_write: 0.55,
|
||||
},
|
||||
m if m.starts_with("deepseek-chat") => ModelPricing {
|
||||
input: 0.27,
|
||||
output: 1.10,
|
||||
cache_read: 0.07,
|
||||
cache_write: 0.27,
|
||||
},
|
||||
// Default (sonnet-like pricing)
|
||||
_ => ModelPricing {
|
||||
input: 3.0,
|
||||
output: 15.0,
|
||||
cache_read: 0.3,
|
||||
cache_write: 3.75,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/// Estimate cost in USD for given usage and model.
|
||||
pub fn estimate_cost(model: &str, usage: &Usage) -> f64 {
|
||||
let pricing = get_model_pricing(model);
|
||||
let input_cost = (usage.input_tokens as f64 / 1_000_000.0) * pricing.input;
|
||||
let output_cost = (usage.output_tokens as f64 / 1_000_000.0) * pricing.output;
|
||||
let cache_read_cost =
|
||||
(usage.cache_read_input_tokens as f64 / 1_000_000.0) * pricing.cache_read;
|
||||
let cache_write_cost =
|
||||
(usage.cache_creation_input_tokens as f64 / 1_000_000.0) * pricing.cache_write;
|
||||
|
||||
input_cost + output_cost + cache_read_cost + cache_write_cost
|
||||
}
|
||||
@@ -0,0 +1,226 @@
|
||||
use open_agent_sdk::{Agent, AgentOptions};
|
||||
use open_agent_sdk::types::*;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_creation() {
|
||||
let agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
assert!(agent.get_messages().is_empty());
|
||||
assert!(!agent.session_id().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_custom_model() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
model: Some("claude-haiku-4-5-20251001".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(agent.model(), "claude-haiku-4-5-20251001");
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_custom_cwd() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
cwd: Some("/tmp".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Agent should use /tmp as working directory
|
||||
assert!(!agent.session_id().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_allowed_tools() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
allowed_tools: Some(vec!["Read".to_string(), "Glob".to_string()]),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Agent should only have Read and Glob
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_custom_tools() {
|
||||
use async_trait::async_trait;
|
||||
|
||||
struct TestTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TestTool {
|
||||
fn name(&self) -> &str {
|
||||
"TestTool"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"A test tool"
|
||||
}
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema::default()
|
||||
}
|
||||
async fn call(
|
||||
&self,
|
||||
_input: serde_json::Value,
|
||||
_ctx: &ToolUseContext,
|
||||
) -> Result<ToolResult, ToolError> {
|
||||
Ok(ToolResult::text("test result"))
|
||||
}
|
||||
}
|
||||
|
||||
let agent = Agent::new(AgentOptions {
|
||||
custom_tools: vec![Arc::new(TestTool)],
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!agent.session_id().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_clear() {
|
||||
let mut agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
|
||||
// Manually add a message for testing
|
||||
agent.messages.push(Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "test".to_string(),
|
||||
}],
|
||||
});
|
||||
|
||||
assert!(!agent.get_messages().is_empty());
|
||||
agent.clear();
|
||||
assert!(agent.get_messages().is_empty());
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_set_model() {
|
||||
let mut agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
|
||||
agent.set_model("claude-opus-4-6-20250514");
|
||||
assert_eq!(agent.model(), "claude-opus-4-6-20250514");
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_max_turns() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(5),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(agent.max_turns, 5);
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_max_budget() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
max_budget_usd: Some(1.0),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(agent.max_budget_usd, Some(1.0));
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_thinking() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
thinking: Some(ThinkingConfig::enabled(10000)),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(agent.thinking.is_some());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_system_prompt() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
system_prompt: Some("You are a code reviewer.".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
agent.system_prompt.as_deref(),
|
||||
Some("You are a code reviewer.")
|
||||
);
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_append_system_prompt() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
append_system_prompt: Some("Always respond in JSON.".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
agent.append_system_prompt.as_deref(),
|
||||
Some("Always respond in JSON.")
|
||||
);
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_cost_tracker() {
|
||||
let agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
|
||||
let cost = agent.cost_tracker().total_cost().await;
|
||||
assert_eq!(cost, 0.0);
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_permission_mode() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
permission_mode: Some(PermissionMode::AcceptEdits),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Agent should be created successfully with the permission mode
|
||||
assert!(!agent.session_id().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_disallowed_tools() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
disallowed_tools: Some(vec!["Bash".to_string()]),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Agent should not have Bash tool
|
||||
assert!(!agent.session_id().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
use open_agent_sdk::costtracker::CostTracker;
|
||||
use open_agent_sdk::types::Usage;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_new() {
|
||||
let tracker = CostTracker::new();
|
||||
assert_eq!(tracker.total_cost().await, 0.0);
|
||||
assert_eq!(tracker.total_tokens().await, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_add_usage() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
let usage = Usage {
|
||||
input_tokens: 1000,
|
||||
output_tokens: 500,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
};
|
||||
|
||||
tracker.add_usage("claude-sonnet-4-6-20250514", &usage).await;
|
||||
|
||||
assert_eq!(tracker.total_tokens().await, 1500);
|
||||
assert!(tracker.total_cost().await > 0.0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_multiple_models() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker
|
||||
.add_usage(
|
||||
"claude-sonnet-4-6-20250514",
|
||||
&Usage {
|
||||
input_tokens: 1000,
|
||||
output_tokens: 500,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
tracker
|
||||
.add_usage(
|
||||
"claude-opus-4-6-20250514",
|
||||
&Usage {
|
||||
input_tokens: 500,
|
||||
output_tokens: 200,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(tracker.total_tokens().await, 2200);
|
||||
|
||||
let sonnet_usage = tracker
|
||||
.get_model_usage("claude-sonnet-4-6-20250514")
|
||||
.await;
|
||||
assert!(sonnet_usage.is_some());
|
||||
assert_eq!(sonnet_usage.unwrap().input_tokens, 1000);
|
||||
|
||||
let opus_usage = tracker
|
||||
.get_model_usage("claude-opus-4-6-20250514")
|
||||
.await;
|
||||
assert!(opus_usage.is_some());
|
||||
assert_eq!(opus_usage.unwrap().input_tokens, 500);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_accumulates() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
let usage = Usage {
|
||||
input_tokens: 1000,
|
||||
output_tokens: 500,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
tracker.add_usage("claude-sonnet-4-6-20250514", &usage).await;
|
||||
tracker.add_usage("claude-sonnet-4-6-20250514", &usage).await;
|
||||
|
||||
assert_eq!(tracker.total_tokens().await, 3000);
|
||||
|
||||
let model_usage = tracker
|
||||
.get_model_usage("claude-sonnet-4-6-20250514")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(model_usage.input_tokens, 2000);
|
||||
assert_eq!(model_usage.output_tokens, 1000);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_durations() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker.add_api_duration(100).await;
|
||||
tracker.add_api_duration(200).await;
|
||||
tracker.add_tool_duration(50).await;
|
||||
|
||||
let summary = tracker.summary().await;
|
||||
assert_eq!(summary.api_duration_ms, 300);
|
||||
assert_eq!(summary.tool_duration_ms, 50);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_code_changes() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker.add_code_changes(10, 5).await;
|
||||
tracker.add_code_changes(20, 3).await;
|
||||
|
||||
let summary = tracker.summary().await;
|
||||
assert_eq!(summary.lines_added, 30);
|
||||
assert_eq!(summary.lines_removed, 8);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_web_searches() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker.add_web_search().await;
|
||||
tracker.add_web_search().await;
|
||||
|
||||
let summary = tracker.summary().await;
|
||||
assert_eq!(summary.web_searches, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_calculation_sonnet() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
// 1M input + 1M output for sonnet = $3 + $15 = $18
|
||||
tracker
|
||||
.add_usage(
|
||||
"claude-sonnet-4-6-20250514",
|
||||
&Usage {
|
||||
input_tokens: 1_000_000,
|
||||
output_tokens: 1_000_000,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let cost = tracker.total_cost().await;
|
||||
assert!((cost - 18.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_calculation_opus() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
// 1M input + 1M output for opus = $15 + $75 = $90
|
||||
tracker
|
||||
.add_usage(
|
||||
"claude-opus-4-6-20250514",
|
||||
&Usage {
|
||||
input_tokens: 1_000_000,
|
||||
output_tokens: 1_000_000,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let cost = tracker.total_cost().await;
|
||||
assert!((cost - 90.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_all_usage() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker
|
||||
.add_usage(
|
||||
"model-a",
|
||||
&Usage {
|
||||
input_tokens: 100,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
tracker
|
||||
.add_usage(
|
||||
"model-b",
|
||||
&Usage {
|
||||
input_tokens: 200,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let all = tracker.get_all_usage().await;
|
||||
assert_eq!(all.len(), 2);
|
||||
assert!(all.contains_key("model-a"));
|
||||
assert!(all.contains_key("model-b"));
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
use open_agent_sdk::hooks::*;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_config_default() {
|
||||
let config = HookConfig::default();
|
||||
assert!(config.pre_tool_use.is_empty());
|
||||
assert!(config.post_tool_use.is_empty());
|
||||
assert!(config.post_sampling.is_empty());
|
||||
assert!(config.stop.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pre_tool_use_hook_allow() {
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "*".to_string(),
|
||||
handler: Arc::new(|_input| {
|
||||
Box::pin(async move {
|
||||
HookOutput {
|
||||
blocked: false,
|
||||
message: None,
|
||||
..Default::default()
|
||||
}
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = config.run_pre_tool_use("Bash", &json!({})).await;
|
||||
assert!(result.is_none()); // Not blocked
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pre_tool_use_hook_block() {
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "Bash".to_string(),
|
||||
handler: Arc::new(|_input| {
|
||||
Box::pin(async move {
|
||||
HookOutput {
|
||||
blocked: true,
|
||||
message: Some("Bash is blocked".to_string()),
|
||||
..Default::default()
|
||||
}
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = config.run_pre_tool_use("Bash", &json!({})).await;
|
||||
assert!(result.is_some());
|
||||
assert!(result.unwrap().blocked);
|
||||
|
||||
// Different tool should not be blocked
|
||||
let result = config.run_pre_tool_use("Read", &json!({})).await;
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_matcher_wildcard() {
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "mcp__*".to_string(),
|
||||
handler: Arc::new(|_input| {
|
||||
Box::pin(async move {
|
||||
HookOutput {
|
||||
blocked: true,
|
||||
message: Some("MCP blocked".to_string()),
|
||||
..Default::default()
|
||||
}
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = config
|
||||
.run_pre_tool_use("mcp__server__tool", &json!({}))
|
||||
.await;
|
||||
assert!(result.is_some());
|
||||
|
||||
let result = config.run_pre_tool_use("Read", &json!({})).await;
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_matcher_pipe_or() {
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "Bash|Edit|Write".to_string(),
|
||||
handler: Arc::new(|_input| {
|
||||
Box::pin(async move {
|
||||
HookOutput {
|
||||
blocked: true,
|
||||
message: None,
|
||||
..Default::default()
|
||||
}
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert!(config
|
||||
.run_pre_tool_use("Bash", &json!({}))
|
||||
.await
|
||||
.is_some());
|
||||
assert!(config
|
||||
.run_pre_tool_use("Edit", &json!({}))
|
||||
.await
|
||||
.is_some());
|
||||
assert!(config
|
||||
.run_pre_tool_use("Write", &json!({}))
|
||||
.await
|
||||
.is_some());
|
||||
assert!(config
|
||||
.run_pre_tool_use("Read", &json!({}))
|
||||
.await
|
||||
.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_post_tool_use_hook() {
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
let called = Arc::new(AtomicBool::new(false));
|
||||
let called_clone = called.clone();
|
||||
|
||||
let config = HookConfig {
|
||||
post_tool_use: vec![HookRule {
|
||||
matcher: "Bash".to_string(),
|
||||
handler: Arc::new(move |_input| {
|
||||
let called = called_clone.clone();
|
||||
Box::pin(async move {
|
||||
called.store(true, Ordering::SeqCst);
|
||||
HookOutput::default()
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
config
|
||||
.run_post_tool_use("Bash", &json!({}), "output")
|
||||
.await;
|
||||
assert!(called.load(Ordering::SeqCst));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_input_fields() {
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
let verified = Arc::new(AtomicBool::new(false));
|
||||
let verified_clone = verified.clone();
|
||||
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "*".to_string(),
|
||||
handler: Arc::new(move |input| {
|
||||
let verified = verified_clone.clone();
|
||||
Box::pin(async move {
|
||||
assert_eq!(input.event, HookEvent::PreToolUse);
|
||||
assert_eq!(input.tool_name.as_deref(), Some("Bash"));
|
||||
assert!(input.tool_input.is_some());
|
||||
verified.store(true, Ordering::SeqCst);
|
||||
HookOutput::default()
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
config
|
||||
.run_pre_tool_use("Bash", &json!({"command": "ls"}))
|
||||
.await;
|
||||
assert!(verified.load(Ordering::SeqCst));
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
use open_agent_sdk::permissions::*;
|
||||
use open_agent_sdk::types::PermissionMode;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_bypass_mode_allows_everything() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::BypassPermissions,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "Bash", &json!({"command": "rm -rf /"})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
assert_eq!(
|
||||
check_permission(&config, "Write", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deny_rules() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
deny_rules: vec![Rule {
|
||||
tool_name: "Bash".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "Bash", &json!({})),
|
||||
PermissionResult::Deny("Tool 'Bash' denied by rule".to_string())
|
||||
);
|
||||
// Other tools should not be denied
|
||||
assert_eq!(
|
||||
check_permission(&config, "Read", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allow_rules() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
allow_rules: vec![Rule {
|
||||
tool_name: "Read".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "Read", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wildcard_deny() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
deny_rules: vec![Rule {
|
||||
tool_name: "mcp__*".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
allow_rules: vec![Rule {
|
||||
tool_name: "*".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "mcp__server__tool", &json!({})),
|
||||
PermissionResult::Deny("Tool 'mcp__server__tool' denied by rule".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
check_permission(&config, "Read", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allowed_tools_whitelist() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
allowed_tools: Some(vec!["Read".to_string(), "Glob".to_string()]),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "Read", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
match check_permission(&config, "Bash", &json!({})) {
|
||||
PermissionResult::Deny(_) => {}
|
||||
_ => panic!("Expected deny"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deny_with_pattern() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
deny_rules: vec![Rule {
|
||||
tool_name: "Bash".to_string(),
|
||||
pattern: Some("rm".to_string()),
|
||||
}],
|
||||
allow_rules: vec![Rule {
|
||||
tool_name: "*".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Deny when pattern matches
|
||||
assert_eq!(
|
||||
check_permission(&config, "Bash", &json!({"command": "rm file.txt"})),
|
||||
PermissionResult::Deny("Tool 'Bash' denied by rule".to_string())
|
||||
);
|
||||
|
||||
// Allow when pattern doesn't match
|
||||
assert_eq!(
|
||||
check_permission(&config, "Bash", &json!({"command": "ls"})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
// --- Filesystem Validator Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_is_sensitive_path() {
|
||||
assert!(FilesystemValidator::is_sensitive_path("/home/user/.env"));
|
||||
assert!(FilesystemValidator::is_sensitive_path("/home/user/.env.local"));
|
||||
assert!(FilesystemValidator::is_sensitive_path(
|
||||
"/home/user/.ssh/id_rsa"
|
||||
));
|
||||
assert!(FilesystemValidator::is_sensitive_path(
|
||||
"/app/credentials.json"
|
||||
));
|
||||
assert!(!FilesystemValidator::is_sensitive_path("/app/main.rs"));
|
||||
assert!(!FilesystemValidator::is_sensitive_path(
|
||||
"/app/src/config.rs"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filesystem_validator() {
|
||||
let validator = FilesystemValidator::new("/tmp/project");
|
||||
|
||||
// Paths within working dir should be allowed
|
||||
assert!(validator.validate_path("/tmp/project/src/main.rs", false).is_ok());
|
||||
assert!(validator.validate_path("/tmp/project/src/main.rs", true).is_ok());
|
||||
}
|
||||
@@ -0,0 +1,686 @@
|
||||
use open_agent_sdk::tools::ToolRegistry;
|
||||
use open_agent_sdk::types::{Tool, ToolUseContext};
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
fn create_test_context(dir: &str) -> ToolUseContext {
|
||||
ToolUseContext {
|
||||
working_dir: dir.to_string(),
|
||||
abort_signal: tokio_util::sync::CancellationToken::new(),
|
||||
read_file_state: Arc::new(RwLock::new(std::collections::HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Registry Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_default_registry() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
assert!(registry.len() > 0);
|
||||
|
||||
// Check core tools exist
|
||||
assert!(registry.get("Bash").is_some());
|
||||
assert!(registry.get("Read").is_some());
|
||||
assert!(registry.get("Write").is_some());
|
||||
assert!(registry.get("Edit").is_some());
|
||||
assert!(registry.get("Glob").is_some());
|
||||
assert!(registry.get("Grep").is_some());
|
||||
assert!(registry.get("WebFetch").is_some());
|
||||
assert!(registry.get("WebSearch").is_some());
|
||||
assert!(registry.get("AskUserQuestion").is_some());
|
||||
assert!(registry.get("TaskCreate").is_some());
|
||||
assert!(registry.get("TaskGet").is_some());
|
||||
assert!(registry.get("TaskList").is_some());
|
||||
assert!(registry.get("TaskUpdate").is_some());
|
||||
assert!(registry.get("ToolSearch").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_register_custom() {
|
||||
use async_trait::async_trait;
|
||||
use open_agent_sdk::types::{ToolError, ToolInputSchema, ToolResult};
|
||||
use serde_json::Value;
|
||||
|
||||
struct CustomTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for CustomTool {
|
||||
fn name(&self) -> &str {
|
||||
"CustomTool"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"A test tool"
|
||||
}
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema::default()
|
||||
}
|
||||
async fn call(
|
||||
&self,
|
||||
_input: Value,
|
||||
_ctx: &ToolUseContext,
|
||||
) -> Result<ToolResult, ToolError> {
|
||||
Ok(ToolResult::text("custom result"))
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = ToolRegistry::new();
|
||||
assert!(registry.is_empty());
|
||||
|
||||
registry.register(Arc::new(CustomTool));
|
||||
assert_eq!(registry.len(), 1);
|
||||
assert!(registry.get("CustomTool").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_filter() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read_only = registry.filter(|t| t.is_read_only(&json!({})));
|
||||
assert!(read_only.len() > 0);
|
||||
|
||||
// Read, Glob, Grep should be read-only
|
||||
let names: Vec<&str> = read_only.iter().map(|t| t.name()).collect();
|
||||
assert!(names.contains(&"Read"));
|
||||
assert!(names.contains(&"Glob"));
|
||||
assert!(names.contains(&"Grep"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_retain() {
|
||||
let mut registry = ToolRegistry::default_registry();
|
||||
let initial_count = registry.len();
|
||||
|
||||
registry.retain(&["Read", "Glob", "Grep"]);
|
||||
assert_eq!(registry.len(), 3);
|
||||
assert!(registry.len() < initial_count);
|
||||
|
||||
assert!(registry.get("Read").is_some());
|
||||
assert!(registry.get("Bash").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_remove() {
|
||||
let mut registry = ToolRegistry::default_registry();
|
||||
assert!(registry.get("Bash").is_some());
|
||||
|
||||
registry.remove(&["Bash"]);
|
||||
assert!(registry.get("Bash").is_none());
|
||||
}
|
||||
|
||||
// --- Bash Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_echo() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let bash = registry.get("Bash").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = bash
|
||||
.call(json!({"command": "echo 'hello world'"}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("hello world"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_exit_code() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let bash = registry.get("Bash").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = bash.call(json!({"command": "false"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("Exit code:"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_destructive_detection() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let bash = registry.get("Bash").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = bash
|
||||
.call(json!({"command": "rm -rf /"}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("destructive"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_is_read_only() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let bash = registry.get("Bash").unwrap();
|
||||
|
||||
assert!(bash.is_read_only(&json!({"command": "ls -la"})));
|
||||
assert!(bash.is_read_only(&json!({"command": "git status"})));
|
||||
assert!(!bash.is_read_only(&json!({"command": "rm file.txt"})));
|
||||
assert!(!bash.is_read_only(&json!({"command": "cargo build"})));
|
||||
}
|
||||
|
||||
// --- File Read Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("test.txt");
|
||||
std::fs::write(&file_path, "line 1\nline 2\nline 3\n").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = read
|
||||
.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("line 1"));
|
||||
assert!(text.contains("line 2"));
|
||||
assert!(text.contains("line 3"));
|
||||
// Should have line numbers
|
||||
assert!(text.contains("1\t"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_file_not_found() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = read
|
||||
.call(
|
||||
json!({"file_path": "/tmp/nonexistent_file_12345.txt"}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("not found"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_file_with_offset_and_limit() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("test.txt");
|
||||
let content: String = (1..=100).map(|i| format!("line {}\n", i)).collect();
|
||||
std::fs::write(&file_path, &content).unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = read
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"offset": 10,
|
||||
"limit": 5
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("line 11"));
|
||||
assert!(text.contains("line 15"));
|
||||
assert!(!text.contains("line 16"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_directory_error() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = read.call(json!({"file_path": "/tmp"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("directory"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_is_read_only() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
assert!(read.is_read_only(&json!({})));
|
||||
assert!(read.is_concurrency_safe(&json!({})));
|
||||
}
|
||||
|
||||
// --- File Write Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_write_new_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("new_file.txt");
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let write = registry.get("Write").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = write
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"content": "hello world"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("Created"));
|
||||
assert_eq!(std::fs::read_to_string(&file_path).unwrap(), "hello world");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_write_creates_directories() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("a/b/c/file.txt");
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let write = registry.get("Write").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = write
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"content": "nested"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(file_path.exists());
|
||||
}
|
||||
|
||||
// --- File Edit Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_test.txt");
|
||||
std::fs::write(&file_path, "hello world").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
// Must read first for staleness tracking
|
||||
read.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "hello",
|
||||
"new_string": "goodbye"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&file_path).unwrap(),
|
||||
"goodbye world"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_string_not_found() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_test2.txt");
|
||||
std::fs::write(&file_path, "hello world").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
read.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "xyz",
|
||||
"new_string": "abc"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("not found"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_multiple_occurrences() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_multi.txt");
|
||||
std::fs::write(&file_path, "foo bar foo baz foo").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
read.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Without replace_all, should fail for multiple occurrences
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "foo",
|
||||
"new_string": "qux"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("3 times"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_replace_all() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_all.txt");
|
||||
std::fs::write(&file_path, "foo bar foo baz foo").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
read.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "foo",
|
||||
"new_string": "qux",
|
||||
"replace_all": true
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&file_path).unwrap(),
|
||||
"qux bar qux baz qux"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_same_string_error() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_same.txt");
|
||||
std::fs::write(&file_path, "hello world").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "hello",
|
||||
"new_string": "hello"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("different"));
|
||||
}
|
||||
|
||||
// --- Glob Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_find_files() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("a.rs"), "").unwrap();
|
||||
std::fs::write(dir.path().join("b.rs"), "").unwrap();
|
||||
std::fs::write(dir.path().join("c.txt"), "").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = glob.call(json!({"pattern": "*.rs"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("a.rs"));
|
||||
assert!(text.contains("b.rs"));
|
||||
assert!(!text.contains("c.txt"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_no_matches() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = glob.call(json!({"pattern": "*.xyz"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("No files found"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_dot_lists_current_directory_without_recursing() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("root.rs"), "").unwrap();
|
||||
std::fs::create_dir(dir.path().join("nested")).unwrap();
|
||||
std::fs::write(dir.path().join("nested").join("child.rs"), "").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = glob.call(json!({"pattern": "."}), &ctx).await.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("root.rs"));
|
||||
assert!(!text.contains("child.rs"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_skips_heavy_directories_before_descending() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::create_dir(dir.path().join("node_modules")).unwrap();
|
||||
std::fs::write(dir.path().join("node_modules").join("dep.rs"), "").unwrap();
|
||||
std::fs::create_dir(dir.path().join("src")).unwrap();
|
||||
std::fs::write(dir.path().join("src").join("lib.rs"), "").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = glob
|
||||
.call(json!({"pattern": "**/*.rs"}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("lib.rs"));
|
||||
assert!(!text.contains("dep.rs"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_returns_when_cancelled() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
ctx.abort_signal.cancel();
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
|
||||
let result = glob.call(json!({"pattern": "**/*"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("No files found"));
|
||||
}
|
||||
|
||||
// --- Grep Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_grep_search() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(
|
||||
dir.path().join("test.txt"),
|
||||
"hello world\nfoo bar\nhello again\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let grep = registry.get("Grep").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = grep
|
||||
.call(
|
||||
json!({
|
||||
"pattern": "hello",
|
||||
"path": dir.path().to_str().unwrap(),
|
||||
"output_mode": "content"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("hello world") || text.contains("hello"));
|
||||
}
|
||||
|
||||
// --- Task Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_task_lifecycle() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
// Create a task
|
||||
let create = registry.get("TaskCreate").unwrap();
|
||||
let result = create
|
||||
.call(
|
||||
json!({
|
||||
"subject": "Test task",
|
||||
"description": "A test task"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("Created task"));
|
||||
|
||||
// List tasks
|
||||
let list = registry.get("TaskList").unwrap();
|
||||
let result = list.call(json!({}), &ctx).await.unwrap();
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("Test task"));
|
||||
|
||||
// Get task
|
||||
let get = registry.get("TaskGet").unwrap();
|
||||
let result = get.call(json!({"id": "task_1"}), &ctx).await.unwrap();
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("Test task"));
|
||||
|
||||
// Update task
|
||||
let update = registry.get("TaskUpdate").unwrap();
|
||||
let result = update
|
||||
.call(json!({"id": "task_1", "status": "completed"}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!result.is_error);
|
||||
|
||||
// Verify update
|
||||
let result = get.call(json!({"id": "task_1"}), &ctx).await.unwrap();
|
||||
assert!(result.get_text().to_lowercase().contains("completed"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_task_not_found() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let get = registry.get("TaskGet").unwrap();
|
||||
let result = get.call(json!({"id": "nonexistent"}), &ctx).await.unwrap();
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("not found"));
|
||||
}
|
||||
|
||||
// --- WebFetch Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_webfetch_missing_url() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let webfetch = registry.get("WebFetch").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = webfetch.call(json!({}), &ctx).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
// --- Diff Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_unified_diff() {
|
||||
let old = "line 1\nline 2\nline 3\n";
|
||||
let new = "line 1\nline 2 modified\nline 3\n";
|
||||
let diff = open_agent_sdk::tools::diff::unified_diff(old, new, "test.txt");
|
||||
|
||||
assert!(diff.contains("--- a/test.txt"));
|
||||
assert!(diff.contains("+++ b/test.txt"));
|
||||
assert!(diff.contains("-line 2"));
|
||||
assert!(diff.contains("+line 2 modified"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_count_changes() {
|
||||
let old = "a\nb\nc\n";
|
||||
let new = "a\nx\nc\nd\n";
|
||||
let (added, removed) = open_agent_sdk::tools::diff::count_changes(old, new);
|
||||
assert!(added >= 1);
|
||||
assert!(removed >= 1);
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
use open_agent_sdk::types::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_tool_result_text() {
|
||||
let result = ToolResult::text("hello world");
|
||||
assert!(!result.is_error);
|
||||
assert_eq!(result.get_text(), "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_result_error() {
|
||||
let result = ToolResult::error("something went wrong");
|
||||
assert!(result.is_error);
|
||||
assert_eq!(result.get_text(), "something went wrong");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_result_image() {
|
||||
let result = ToolResult::image("base64data".to_string(), "image/png".to_string());
|
||||
assert!(!result.is_error);
|
||||
assert_eq!(result.content.len(), 1);
|
||||
match &result.content[0] {
|
||||
ToolResultContent::Image { source } => {
|
||||
assert_eq!(source.data, "base64data");
|
||||
assert_eq!(source.media_type, "image/png");
|
||||
}
|
||||
_ => panic!("Expected image content"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_input_schema_default() {
|
||||
let schema = ToolInputSchema::default();
|
||||
assert_eq!(schema.schema_type, "object");
|
||||
assert!(schema.properties.is_empty());
|
||||
assert!(schema.required.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_permission_mode_default() {
|
||||
let mode = PermissionMode::default();
|
||||
assert_eq!(mode, PermissionMode::BypassPermissions);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_permission_mode_serialization() {
|
||||
let mode = PermissionMode::AcceptEdits;
|
||||
let json = serde_json::to_string(&mode).unwrap();
|
||||
assert_eq!(json, "\"acceptEdits\"");
|
||||
|
||||
let deserialized: PermissionMode = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(deserialized, PermissionMode::AcceptEdits);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_message_creation() {
|
||||
let msg = Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
}],
|
||||
};
|
||||
assert_eq!(msg.role, MessageRole::User);
|
||||
assert_eq!(msg.content.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_text_from_message() {
|
||||
let msg = Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::Text {
|
||||
text: "Hello ".to_string(),
|
||||
},
|
||||
ContentBlock::Text {
|
||||
text: "World".to_string(),
|
||||
},
|
||||
],
|
||||
};
|
||||
assert_eq!(extract_text(&msg), "Hello World");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_text_skips_non_text() {
|
||||
let msg = Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::Text {
|
||||
text: "Hello".to_string(),
|
||||
},
|
||||
ContentBlock::ToolUse {
|
||||
id: "123".to_string(),
|
||||
name: "Read".to_string(),
|
||||
input: json!({}),
|
||||
},
|
||||
ContentBlock::Text {
|
||||
text: " World".to_string(),
|
||||
},
|
||||
],
|
||||
};
|
||||
assert_eq!(extract_text(&msg), "Hello World");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_tool_uses() {
|
||||
let msg = Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::Text {
|
||||
text: "Let me read that".to_string(),
|
||||
},
|
||||
ContentBlock::ToolUse {
|
||||
id: "tu_1".to_string(),
|
||||
name: "Read".to_string(),
|
||||
input: json!({"file_path": "/tmp/test.txt"}),
|
||||
},
|
||||
ContentBlock::ToolUse {
|
||||
id: "tu_2".to_string(),
|
||||
name: "Glob".to_string(),
|
||||
input: json!({"pattern": "*.rs"}),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
let tool_uses = extract_tool_uses(&msg);
|
||||
assert_eq!(tool_uses.len(), 2);
|
||||
assert_eq!(tool_uses[0].1, "Read");
|
||||
assert_eq!(tool_uses[1].1, "Glob");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_usage_default() {
|
||||
let usage = Usage::default();
|
||||
assert_eq!(usage.input_tokens, 0);
|
||||
assert_eq!(usage.output_tokens, 0);
|
||||
assert_eq!(usage.cache_creation_input_tokens, 0);
|
||||
assert_eq!(usage.cache_read_input_tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thinking_config() {
|
||||
let config = ThinkingConfig::enabled(10000);
|
||||
assert_eq!(config.thinking_type, "enabled");
|
||||
assert_eq!(config.budget_tokens, Some(10000));
|
||||
|
||||
let disabled = ThinkingConfig::disabled();
|
||||
assert_eq!(disabled.thinking_type, "disabled");
|
||||
assert_eq!(disabled.budget_tokens, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_message_role_serialization() {
|
||||
let user: MessageRole = serde_json::from_str("\"user\"").unwrap();
|
||||
assert_eq!(user, MessageRole::User);
|
||||
|
||||
let assistant: MessageRole = serde_json::from_str("\"assistant\"").unwrap();
|
||||
assert_eq!(assistant, MessageRole::Assistant);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_content_block_text_serialization() {
|
||||
let block = ContentBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
};
|
||||
let json = serde_json::to_value(&block).unwrap();
|
||||
assert_eq!(json["type"], "text");
|
||||
assert_eq!(json["text"], "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_content_block_tool_use_serialization() {
|
||||
let block = ContentBlock::ToolUse {
|
||||
id: "tu_1".to_string(),
|
||||
name: "Read".to_string(),
|
||||
input: json!({"file_path": "/test.txt"}),
|
||||
};
|
||||
let json = serde_json::to_value(&block).unwrap();
|
||||
assert_eq!(json["type"], "tool_use");
|
||||
assert_eq!(json["name"], "Read");
|
||||
assert_eq!(json["input"]["file_path"], "/test.txt");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sdk_message_serialization() {
|
||||
let msg = SDKMessage::System {
|
||||
message: "Agent started".to_string(),
|
||||
};
|
||||
let json = serde_json::to_value(&msg).unwrap();
|
||||
assert_eq!(json["type"], "system");
|
||||
assert_eq!(json["message"], "Agent started");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cache_control() {
|
||||
let cc = CacheControl::ephemeral();
|
||||
assert_eq!(cc.control_type, "ephemeral");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_server_config_stdio() {
|
||||
let config = McpServerConfig::Stdio {
|
||||
command: "npx".to_string(),
|
||||
args: vec!["server".to_string()],
|
||||
env: std::collections::HashMap::new(),
|
||||
};
|
||||
let json = serde_json::to_value(&config).unwrap();
|
||||
assert_eq!(json["transport"], "stdio");
|
||||
assert_eq!(json["command"], "npx");
|
||||
}
|
||||
@@ -0,0 +1,273 @@
|
||||
use open_agent_sdk::types::*;
|
||||
use open_agent_sdk::utils::{compact, messages, retry, tokens};
|
||||
use serde_json::json;
|
||||
|
||||
// --- Token Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_estimate_tokens() {
|
||||
assert_eq!(tokens::estimate_tokens(""), 0);
|
||||
assert_eq!(tokens::estimate_tokens("1234"), 1);
|
||||
assert_eq!(tokens::estimate_tokens("12345678"), 2);
|
||||
assert!(tokens::estimate_tokens("hello world this is a test") > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_messages_tokens() {
|
||||
let msgs = vec![
|
||||
Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
}],
|
||||
},
|
||||
Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "world".to_string(),
|
||||
}],
|
||||
},
|
||||
];
|
||||
let tokens = tokens::estimate_messages_tokens(&msgs);
|
||||
assert!(tokens > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_context_window_size() {
|
||||
assert_eq!(
|
||||
tokens::get_context_window_size("claude-sonnet-4-6-20250514"),
|
||||
200_000
|
||||
);
|
||||
assert_eq!(
|
||||
tokens::get_context_window_size("claude-opus-4-6-20250514"),
|
||||
200_000
|
||||
);
|
||||
assert_eq!(
|
||||
tokens::get_context_window_size("claude-haiku-4-5-20251001"),
|
||||
200_000
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_auto_compact_threshold() {
|
||||
let threshold = tokens::get_auto_compact_threshold("claude-sonnet-4-6-20250514");
|
||||
assert_eq!(threshold, 200_000 - 13_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_cost() {
|
||||
let usage = Usage {
|
||||
input_tokens: 1_000_000,
|
||||
output_tokens: 1_000_000,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
};
|
||||
|
||||
let cost = tokens::estimate_cost("claude-sonnet-4-6-20250514", &usage);
|
||||
// Sonnet: $3/M input + $15/M output = $18
|
||||
assert!((cost - 18.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_cost_with_cache() {
|
||||
let usage = Usage {
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cache_creation_input_tokens: 1_000_000,
|
||||
cache_read_input_tokens: 1_000_000,
|
||||
};
|
||||
|
||||
let cost = tokens::estimate_cost("claude-sonnet-4-6-20250514", &usage);
|
||||
// Sonnet cache: $3.75/M write + $0.3/M read = $4.05
|
||||
assert!((cost - 4.05).abs() < 0.01);
|
||||
}
|
||||
|
||||
// --- Message Utility Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_create_user_message() {
|
||||
let msg = messages::create_user_message("hello");
|
||||
assert_eq!(msg.role, MessageRole::User);
|
||||
assert_eq!(messages::extract_text(&msg), "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_assistant_message() {
|
||||
let msg = messages::create_assistant_message("hi there");
|
||||
assert_eq!(msg.role, MessageRole::Assistant);
|
||||
assert_eq!(messages::extract_text(&msg), "hi there");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_messages_empty() {
|
||||
let normalized = messages::normalize_messages(&[]);
|
||||
assert!(normalized.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_messages_starts_with_user() {
|
||||
let msgs = vec![messages::create_assistant_message("I'm an assistant")];
|
||||
let normalized = messages::normalize_messages(&msgs);
|
||||
assert_eq!(normalized[0].role, MessageRole::User);
|
||||
assert_eq!(normalized.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_messages_merges_consecutive() {
|
||||
let msgs = vec![
|
||||
messages::create_user_message("hello"),
|
||||
messages::create_user_message("world"),
|
||||
messages::create_assistant_message("hi"),
|
||||
];
|
||||
let normalized = messages::normalize_messages(&msgs);
|
||||
assert_eq!(normalized.len(), 2); // Merged two user messages
|
||||
assert_eq!(normalized[0].role, MessageRole::User);
|
||||
assert_eq!(normalized[1].role, MessageRole::Assistant);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_images() {
|
||||
let msgs = vec![Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![
|
||||
ContentBlock::Text {
|
||||
text: "look at this".to_string(),
|
||||
},
|
||||
ContentBlock::Image {
|
||||
source: ImageContentSource {
|
||||
source_type: "base64".to_string(),
|
||||
media_type: "image/png".to_string(),
|
||||
data: "abcdef".to_string(),
|
||||
},
|
||||
},
|
||||
],
|
||||
}];
|
||||
|
||||
let stripped = messages::strip_images(&msgs);
|
||||
assert_eq!(stripped[0].content.len(), 1);
|
||||
match &stripped[0].content[0] {
|
||||
ContentBlock::Text { text } => assert_eq!(text, "look at this"),
|
||||
_ => panic!("Expected text block"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_text() {
|
||||
assert_eq!(messages::truncate_text("hello", 10), "hello");
|
||||
assert_eq!(
|
||||
messages::truncate_text("hello world", 5),
|
||||
"hello... (truncated)"
|
||||
);
|
||||
}
|
||||
|
||||
// --- Compact Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_should_auto_compact_false() {
|
||||
let msgs = vec![messages::create_user_message("short message")];
|
||||
assert!(!compact::should_auto_compact(&msgs, "claude-sonnet-4-6-20250514"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_micro_compact_preserves_short_results() {
|
||||
let msgs = vec![Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::ToolResult {
|
||||
tool_use_id: "tu_1".to_string(),
|
||||
content: vec![ToolResultContentBlock::Text {
|
||||
text: "short result".to_string(),
|
||||
}],
|
||||
is_error: false,
|
||||
}],
|
||||
}];
|
||||
|
||||
let compacted = compact::micro_compact_messages(&msgs);
|
||||
if let ContentBlock::ToolResult { content, .. } = &compacted[0].content[0] {
|
||||
if let ToolResultContentBlock::Text { text } = &content[0] {
|
||||
assert_eq!(text, "short result");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_micro_compact_truncates_long_results() {
|
||||
let long_text = "x".repeat(100_000);
|
||||
let msgs = vec![Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::ToolResult {
|
||||
tool_use_id: "tu_1".to_string(),
|
||||
content: vec![ToolResultContentBlock::Text {
|
||||
text: long_text.clone(),
|
||||
}],
|
||||
is_error: false,
|
||||
}],
|
||||
}];
|
||||
|
||||
let compacted = compact::micro_compact_messages(&msgs);
|
||||
if let ContentBlock::ToolResult { content, .. } = &compacted[0].content[0] {
|
||||
if let ToolResultContentBlock::Text { text } = &content[0] {
|
||||
assert!(text.len() < long_text.len());
|
||||
assert!(text.contains("truncated"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Retry Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_retry_config_default() {
|
||||
let config = retry::RetryConfig::default();
|
||||
assert_eq!(config.max_retries, 3);
|
||||
assert_eq!(config.base_delay_ms, 2000);
|
||||
assert_eq!(config.max_delay_ms, 30_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_retryable() {
|
||||
use open_agent_sdk::api::ApiError;
|
||||
|
||||
assert!(retry::is_retryable(&ApiError::RateLimitError));
|
||||
assert!(retry::is_retryable(&ApiError::Timeout));
|
||||
assert!(retry::is_retryable(&ApiError::NetworkError(
|
||||
"conn reset".to_string()
|
||||
)));
|
||||
assert!(retry::is_retryable(&ApiError::HttpError {
|
||||
status: 500,
|
||||
message: "internal error".to_string(),
|
||||
}));
|
||||
assert!(!retry::is_retryable(&ApiError::AuthError(
|
||||
"bad key".to_string()
|
||||
)));
|
||||
assert!(!retry::is_retryable(&ApiError::PromptTooLong(
|
||||
"too long".to_string()
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_auth_error() {
|
||||
use open_agent_sdk::api::ApiError;
|
||||
|
||||
assert!(retry::is_auth_error(&ApiError::AuthError(
|
||||
"bad key".to_string()
|
||||
)));
|
||||
assert!(!retry::is_auth_error(&ApiError::RateLimitError));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_prompt_too_long() {
|
||||
use open_agent_sdk::api::ApiError;
|
||||
|
||||
assert!(retry::is_prompt_too_long(&ApiError::PromptTooLong(
|
||||
"too long".to_string()
|
||||
)));
|
||||
assert!(!retry::is_prompt_too_long(&ApiError::RateLimitError));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_retry_delay() {
|
||||
let config = retry::RetryConfig::default();
|
||||
let delay = retry::get_retry_delay(&config, 0);
|
||||
// Should be between 1s and 3s (base 2s with jitter)
|
||||
assert!(delay.as_millis() >= 1000);
|
||||
assert!(delay.as_millis() <= 4000);
|
||||
}
|
||||
Reference in New Issue
Block a user