chore: sync current workspace changes
Desktop CI Build / Build macOS ARM64 DMG (push) Has been cancelled
Desktop CI Build / Build Windows x64 (push) Has been cancelled
Desktop CI Build / Build macOS ARM64 DMG (push) Has been cancelled
Desktop CI Build / Build Windows x64 (push) Has been cancelled
This commit is contained in:
Generated
+16
-1
@@ -52,6 +52,20 @@ dependencies = [
|
||||
"subtle",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "agent-core"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"async-trait",
|
||||
"open-agent-sdk",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ahash"
|
||||
version = "0.7.8"
|
||||
@@ -1164,6 +1178,7 @@ dependencies = [
|
||||
name = "boosapi"
|
||||
version = "0.0.1"
|
||||
dependencies = [
|
||||
"agent-core",
|
||||
"async-trait",
|
||||
"base64 0.22.1",
|
||||
"boosapi-agent",
|
||||
@@ -1211,6 +1226,7 @@ dependencies = [
|
||||
name = "boosapi-agent"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"agent-core",
|
||||
"async-trait",
|
||||
"boosapi-core",
|
||||
"boosapi-providers",
|
||||
@@ -5049,7 +5065,6 @@ dependencies = [
|
||||
"serde",
|
||||
"serde_json",
|
||||
"serde_yaml",
|
||||
"tempfile",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tokio-stream",
|
||||
|
||||
@@ -15,6 +15,8 @@ tracing-subscriber = "0.3"
|
||||
anyhow = "1"
|
||||
sqlite-vec = "0.1.8-alpha.1"
|
||||
libsqlite3-sys = "0.30"
|
||||
agent-core = { path = "../../crates/agent-core" }
|
||||
open-agent-sdk = { path = "../../crates/open-agent-sdk" }
|
||||
|
||||
[package]
|
||||
name = "boosapi"
|
||||
@@ -51,7 +53,8 @@ serde_yaml = "0.9"
|
||||
boosapi-core = { path = "crates/core" }
|
||||
boosapi-providers = { path = "crates/providers" }
|
||||
boosapi-agent = { path = "crates/agent" }
|
||||
open-agent-sdk = { path = "crates/open-agent-sdk" }
|
||||
agent-core = { workspace = true }
|
||||
open-agent-sdk = { workspace = true }
|
||||
boosapi-gateway = { path = "crates/gateway" }
|
||||
tokio = { workspace = true }
|
||||
sea-orm = { workspace = true }
|
||||
|
||||
@@ -4,7 +4,8 @@ version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
[dependencies]
|
||||
open-agent-sdk = { path = "../open-agent-sdk" }
|
||||
agent-core = { workspace = true }
|
||||
open-agent-sdk = { workspace = true }
|
||||
boosapi-core = { path = "../core" }
|
||||
boosapi-providers = { path = "../providers" }
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
|
||||
@@ -34,6 +34,29 @@ pub enum PermissionAction {
|
||||
HardDeny,
|
||||
}
|
||||
|
||||
/// The desktop agent deliberately exposes a small, auditable tool surface.
|
||||
/// Keep this list in sync with the tools registered in `commands::agent`.
|
||||
pub const DESKTOP_ALLOWED_TOOLS: &[&str] = &[
|
||||
"Read",
|
||||
"Write",
|
||||
"Edit",
|
||||
"Glob",
|
||||
"Grep",
|
||||
"AskUserQuestion",
|
||||
"Bash",
|
||||
"Skill",
|
||||
];
|
||||
|
||||
pub fn is_desktop_tool_allowed(tool_name: &str) -> bool {
|
||||
DESKTOP_ALLOWED_TOOLS
|
||||
.iter()
|
||||
.any(|allowed| allowed.eq_ignore_ascii_case(tool_name))
|
||||
}
|
||||
|
||||
pub fn is_bash_tool(tool_name: &str) -> bool {
|
||||
tool_name.eq_ignore_ascii_case("bash")
|
||||
}
|
||||
|
||||
/// Classify a tool's risk level based on its name
|
||||
pub fn classify_tool_risk(tool_name: &str) -> RiskLevel {
|
||||
let name_lower = tool_name.to_lowercase();
|
||||
@@ -86,6 +109,12 @@ pub fn decide_permission(
|
||||
risk: RiskLevel,
|
||||
is_always_allowed: bool,
|
||||
) -> PermissionAction {
|
||||
// Host command execution is always a one-shot decision. It must never be
|
||||
// inherited from FullAccess or a remembered approval.
|
||||
if risk == RiskLevel::Execute {
|
||||
return PermissionAction::RequireApproval;
|
||||
}
|
||||
|
||||
// If tool was previously approved with "always allow", auto-allow
|
||||
if is_always_allowed {
|
||||
return PermissionAction::AutoAllow;
|
||||
@@ -102,7 +131,34 @@ pub fn decide_permission(
|
||||
(PermissionMode::AcceptEdits, RiskLevel::Write) => PermissionAction::AutoAllow,
|
||||
(PermissionMode::AcceptEdits, RiskLevel::Execute) => PermissionAction::RequireApproval,
|
||||
|
||||
// Full access: everything auto-allowed
|
||||
// Full access applies to file operations. Execute was handled above.
|
||||
(PermissionMode::FullAccess, _) => PermissionAction::AutoAllow,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn desktop_tool_allowlist_denies_unknown_capabilities() {
|
||||
assert!(is_desktop_tool_allowed("Read"));
|
||||
assert!(is_desktop_tool_allowed("skill"));
|
||||
assert!(!is_desktop_tool_allowed("WebFetch"));
|
||||
assert!(!is_desktop_tool_allowed("NotebookEdit"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bash_always_requires_one_shot_approval() {
|
||||
for mode in [
|
||||
PermissionMode::Default,
|
||||
PermissionMode::AcceptEdits,
|
||||
PermissionMode::FullAccess,
|
||||
] {
|
||||
assert_eq!(
|
||||
decide_permission(mode, RiskLevel::Execute, true),
|
||||
PermissionAction::RequireApproval
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
use sea_orm::entity::prelude::*;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
|
||||
#[sea_orm(table_name = "agent_runs")]
|
||||
pub struct Model {
|
||||
#[sea_orm(primary_key, auto_increment = false)]
|
||||
pub id: String,
|
||||
pub session_id: String,
|
||||
pub conversation_id: String,
|
||||
pub status: String,
|
||||
pub assistant_message_id: Option<String>,
|
||||
pub error_message: Option<String>,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
pub finished_at: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
|
||||
pub enum Relation {}
|
||||
|
||||
impl ActiveModelBehavior for ActiveModel {}
|
||||
@@ -10,6 +10,12 @@ pub struct Model {
|
||||
pub cwd: Option<String>,
|
||||
pub permission_mode: String,
|
||||
pub runtime_status: String,
|
||||
pub current_run_id: Option<String>,
|
||||
pub host_access_enabled: bool,
|
||||
pub revision: i64,
|
||||
pub context_version: i32,
|
||||
pub last_provider_id: Option<String>,
|
||||
pub last_model_id: Option<String>,
|
||||
pub sdk_context_json: Option<String>,
|
||||
pub sdk_context_backup_json: Option<String>,
|
||||
pub total_tokens: i32,
|
||||
|
||||
@@ -40,5 +40,6 @@ pub mod retrieval_hits;
|
||||
pub mod stored_files;
|
||||
|
||||
pub mod agent_sessions;
|
||||
pub mod agent_runs;
|
||||
|
||||
pub use sea_orm;
|
||||
|
||||
@@ -27,3 +27,4 @@ pub mod skill;
|
||||
pub mod stored_file;
|
||||
pub mod tool_execution;
|
||||
pub mod agent_session;
|
||||
pub mod agent_run;
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
use crate::entity::agent_runs;
|
||||
use crate::error::{BoosApiError, Result};
|
||||
use sea_orm::*;
|
||||
|
||||
const TERMINAL_STATUSES: &[&str] = &["succeeded", "failed", "cancelled", "interrupted"];
|
||||
|
||||
fn now() -> String {
|
||||
chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string()
|
||||
}
|
||||
|
||||
pub async fn create_agent_run(
|
||||
db: &DatabaseConnection,
|
||||
run_id: &str,
|
||||
session_id: &str,
|
||||
conversation_id: &str,
|
||||
) -> Result<()> {
|
||||
let timestamp = now();
|
||||
agent_runs::ActiveModel {
|
||||
id: Set(run_id.to_string()),
|
||||
session_id: Set(session_id.to_string()),
|
||||
conversation_id: Set(conversation_id.to_string()),
|
||||
status: Set("starting".to_string()),
|
||||
assistant_message_id: Set(None),
|
||||
error_message: Set(None),
|
||||
created_at: Set(timestamp.clone()),
|
||||
updated_at: Set(timestamp),
|
||||
finished_at: Set(None),
|
||||
}
|
||||
.insert(db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn update_agent_run_status(
|
||||
db: &DatabaseConnection,
|
||||
run_id: &str,
|
||||
status: &str,
|
||||
assistant_message_id: Option<&str>,
|
||||
error_message: Option<&str>,
|
||||
) -> Result<()> {
|
||||
let timestamp = now();
|
||||
let mut update = agent_runs::Entity::update_many()
|
||||
.col_expr(agent_runs::Column::Status, Expr::value(status.to_string()))
|
||||
.col_expr(
|
||||
agent_runs::Column::UpdatedAt,
|
||||
Expr::value(timestamp.clone()),
|
||||
)
|
||||
.filter(agent_runs::Column::Id.eq(run_id))
|
||||
.filter(agent_runs::Column::Status.is_in([
|
||||
"starting",
|
||||
"running",
|
||||
"waiting_approval",
|
||||
"waiting_input",
|
||||
"cancelling",
|
||||
]));
|
||||
if let Some(message_id) = assistant_message_id {
|
||||
update = update.col_expr(
|
||||
agent_runs::Column::AssistantMessageId,
|
||||
Expr::value(Some(message_id.to_string())),
|
||||
);
|
||||
}
|
||||
if let Some(error) = error_message {
|
||||
update = update.col_expr(
|
||||
agent_runs::Column::ErrorMessage,
|
||||
Expr::value(Some(error.to_string())),
|
||||
);
|
||||
}
|
||||
if TERMINAL_STATUSES.contains(&status) {
|
||||
update = update.col_expr(agent_runs::Column::FinishedAt, Expr::value(Some(timestamp)));
|
||||
}
|
||||
let result = update.exec(db).await?;
|
||||
if result.rows_affected != 1 {
|
||||
return Err(BoosApiError::Validation(format!(
|
||||
"Agent run {run_id} is no longer active"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn interrupt_active_runs(db: &DatabaseConnection) -> Result<u64> {
|
||||
let timestamp = now();
|
||||
let result = agent_runs::Entity::update_many()
|
||||
.col_expr(
|
||||
agent_runs::Column::Status,
|
||||
sea_orm::sea_query::Expr::value("interrupted"),
|
||||
)
|
||||
.col_expr(
|
||||
agent_runs::Column::UpdatedAt,
|
||||
sea_orm::sea_query::Expr::value(timestamp.clone()),
|
||||
)
|
||||
.col_expr(
|
||||
agent_runs::Column::FinishedAt,
|
||||
sea_orm::sea_query::Expr::value(Some(timestamp)),
|
||||
)
|
||||
.filter(agent_runs::Column::Status.is_in([
|
||||
"starting",
|
||||
"running",
|
||||
"waiting_approval",
|
||||
"waiting_input",
|
||||
"cancelling",
|
||||
]))
|
||||
.exec(db)
|
||||
.await?;
|
||||
Ok(result.rows_affected)
|
||||
}
|
||||
@@ -12,6 +12,12 @@ fn model_to_agent_session(model: agent_sessions::Model) -> AgentSession {
|
||||
cwd: model.cwd,
|
||||
permission_mode: model.permission_mode,
|
||||
runtime_status: model.runtime_status,
|
||||
current_run_id: model.current_run_id,
|
||||
host_access_enabled: model.host_access_enabled,
|
||||
revision: model.revision,
|
||||
context_version: model.context_version,
|
||||
last_provider_id: model.last_provider_id,
|
||||
last_model_id: model.last_model_id,
|
||||
sdk_context_json: model.sdk_context_json,
|
||||
sdk_context_backup_json: model.sdk_context_backup_json,
|
||||
total_tokens: model.total_tokens,
|
||||
@@ -55,6 +61,12 @@ pub async fn upsert_agent_session(
|
||||
cwd: Set(cwd.map(|s| s.to_string())),
|
||||
permission_mode: Set(permission_mode.unwrap_or("default").to_string()),
|
||||
runtime_status: Set("idle".to_string()),
|
||||
current_run_id: Set(None),
|
||||
host_access_enabled: Set(false),
|
||||
revision: Set(0),
|
||||
context_version: Set(1),
|
||||
last_provider_id: Set(None),
|
||||
last_model_id: Set(None),
|
||||
sdk_context_json: Set(None),
|
||||
sdk_context_backup_json: Set(None),
|
||||
total_tokens: Set(0),
|
||||
@@ -99,6 +111,92 @@ pub async fn update_agent_session_status(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Atomically binds an idle session to a run. The in-process lease is acquired
|
||||
/// before any await; this database guard prevents stale cleanup from taking
|
||||
/// ownership of a newer run after a restart.
|
||||
pub async fn begin_agent_run(
|
||||
db: &DatabaseConnection,
|
||||
id: &str,
|
||||
run_id: &str,
|
||||
provider_id: &str,
|
||||
model_id: &str,
|
||||
) -> Result<()> {
|
||||
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
||||
let result = agent_sessions::Entity::update_many()
|
||||
.col_expr(agent_sessions::Column::RuntimeStatus, Expr::value("active"))
|
||||
.col_expr(
|
||||
agent_sessions::Column::CurrentRunId,
|
||||
Expr::value(Some(run_id.to_string())),
|
||||
)
|
||||
.col_expr(
|
||||
agent_sessions::Column::LastProviderId,
|
||||
Expr::value(Some(provider_id.to_string())),
|
||||
)
|
||||
.col_expr(
|
||||
agent_sessions::Column::LastModelId,
|
||||
Expr::value(Some(model_id.to_string())),
|
||||
)
|
||||
.col_expr(
|
||||
agent_sessions::Column::Revision,
|
||||
Expr::col(agent_sessions::Column::Revision).add(1),
|
||||
)
|
||||
.col_expr(agent_sessions::Column::UpdatedAt, Expr::value(now))
|
||||
.filter(agent_sessions::Column::Id.eq(id))
|
||||
.filter(agent_sessions::Column::CurrentRunId.is_null())
|
||||
.exec(db)
|
||||
.await?;
|
||||
if result.rows_affected != 1 {
|
||||
return Err(BoosApiError::Validation(
|
||||
"Agent session already has an active run".to_string(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn update_agent_session_status_for_run(
|
||||
db: &DatabaseConnection,
|
||||
id: &str,
|
||||
run_id: &str,
|
||||
runtime_status: &str,
|
||||
) -> Result<bool> {
|
||||
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
||||
let result = agent_sessions::Entity::update_many()
|
||||
.col_expr(
|
||||
agent_sessions::Column::RuntimeStatus,
|
||||
Expr::value(runtime_status.to_string()),
|
||||
)
|
||||
.col_expr(agent_sessions::Column::UpdatedAt, Expr::value(now))
|
||||
.filter(agent_sessions::Column::Id.eq(id))
|
||||
.filter(agent_sessions::Column::CurrentRunId.eq(run_id))
|
||||
.exec(db)
|
||||
.await?;
|
||||
Ok(result.rows_affected == 1)
|
||||
}
|
||||
|
||||
pub async fn update_agent_session_host_access(
|
||||
db: &DatabaseConnection,
|
||||
conversation_id: &str,
|
||||
enabled: bool,
|
||||
) -> Result<AgentSession> {
|
||||
let model = agent_sessions::Entity::find()
|
||||
.filter(agent_sessions::Column::ConversationId.eq(conversation_id))
|
||||
.one(db)
|
||||
.await?
|
||||
.ok_or_else(|| BoosApiError::NotFound(format!("AgentSession {conversation_id}")))?;
|
||||
if model.current_run_id.is_some() {
|
||||
return Err(BoosApiError::Validation(
|
||||
"Cannot change host access while an agent is running".to_string(),
|
||||
));
|
||||
}
|
||||
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
||||
let next_revision = model.revision + 1;
|
||||
let mut active: agent_sessions::ActiveModel = model.into();
|
||||
active.host_access_enabled = Set(enabled);
|
||||
active.revision = Set(next_revision);
|
||||
active.updated_at = Set(now);
|
||||
Ok(model_to_agent_session(active.update(db).await?))
|
||||
}
|
||||
|
||||
/// Update the working directory of an agent session.
|
||||
pub async fn update_agent_session_cwd(
|
||||
db: &DatabaseConnection,
|
||||
@@ -137,22 +235,26 @@ pub async fn update_agent_session_permission_mode(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Reset all running/waiting_approval sessions to idle (for app startup recovery).
|
||||
/// Mark sessions interrupted on startup and release any stale run binding.
|
||||
pub async fn reset_running_sessions(db: &DatabaseConnection) -> Result<u64> {
|
||||
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
||||
let result = agent_sessions::Entity::update_many()
|
||||
.col_expr(
|
||||
agent_sessions::Column::RuntimeStatus,
|
||||
Expr::value("idle"),
|
||||
Expr::value("interrupted"),
|
||||
)
|
||||
.col_expr(agent_sessions::Column::CurrentRunId, Expr::value(None::<String>))
|
||||
.col_expr(
|
||||
agent_sessions::Column::UpdatedAt,
|
||||
Expr::value(now),
|
||||
)
|
||||
.filter(
|
||||
Condition::any()
|
||||
.add(agent_sessions::Column::RuntimeStatus.eq("active"))
|
||||
.add(agent_sessions::Column::RuntimeStatus.eq("running"))
|
||||
.add(agent_sessions::Column::RuntimeStatus.eq("waiting_approval")),
|
||||
.add(agent_sessions::Column::RuntimeStatus.eq("waiting_approval"))
|
||||
.add(agent_sessions::Column::RuntimeStatus.eq("waiting_input"))
|
||||
.add(agent_sessions::Column::RuntimeStatus.eq("cancelling")),
|
||||
)
|
||||
.exec(db)
|
||||
.await?;
|
||||
@@ -164,12 +266,14 @@ pub async fn reset_running_sessions(db: &DatabaseConnection) -> Result<u64> {
|
||||
pub async fn update_agent_session_after_query(
|
||||
db: &DatabaseConnection,
|
||||
id: &str,
|
||||
run_id: &str,
|
||||
runtime_status: &str,
|
||||
sdk_context_json: Option<&str>,
|
||||
tokens_delta: i32,
|
||||
cost_delta: f64,
|
||||
) -> Result<()> {
|
||||
let model = agent_sessions::Entity::find_by_id(id)
|
||||
.filter(agent_sessions::Column::CurrentRunId.eq(run_id))
|
||||
.one(db)
|
||||
.await?
|
||||
.ok_or_else(|| BoosApiError::NotFound(format!("AgentSession {}", id)))?;
|
||||
@@ -177,11 +281,13 @@ pub async fn update_agent_session_after_query(
|
||||
let now = chrono::Utc::now().format("%Y-%m-%d %H:%M:%S").to_string();
|
||||
let mut am: agent_sessions::ActiveModel = model.clone().into();
|
||||
am.runtime_status = Set(runtime_status.to_string());
|
||||
am.current_run_id = Set(None);
|
||||
if let Some(ctx) = sdk_context_json {
|
||||
am.sdk_context_json = Set(Some(ctx.to_string()));
|
||||
}
|
||||
am.total_tokens = Set(model.total_tokens + tokens_delta);
|
||||
am.total_cost_usd = Set(model.total_cost_usd + cost_delta);
|
||||
am.revision = Set(model.revision + 1);
|
||||
am.updated_at = Set(now);
|
||||
am.update(db).await?;
|
||||
Ok(())
|
||||
|
||||
@@ -1450,6 +1450,12 @@ pub struct AgentSession {
|
||||
pub cwd: Option<String>,
|
||||
pub permission_mode: String,
|
||||
pub runtime_status: String,
|
||||
pub current_run_id: Option<String>,
|
||||
pub host_access_enabled: bool,
|
||||
pub revision: i64,
|
||||
pub context_version: i32,
|
||||
pub last_provider_id: Option<String>,
|
||||
pub last_model_id: Option<String>,
|
||||
pub sdk_context_json: Option<String>,
|
||||
pub sdk_context_backup_json: Option<String>,
|
||||
pub total_tokens: i32,
|
||||
|
||||
@@ -33,6 +33,7 @@ mod m20260501_000001_add_knowledge_base_rerank_settings;
|
||||
mod m20260504_000001_split_openai_compatible_provider_types;
|
||||
mod m20260515_000001_add_knowledge_base_index_schedule;
|
||||
mod m20260518_000001_add_builtin_model_deletions;
|
||||
mod m20260710_000001_agent_run_state;
|
||||
|
||||
pub struct Migrator;
|
||||
|
||||
@@ -73,6 +74,7 @@ impl MigratorTrait for Migrator {
|
||||
Box::new(m20260504_000001_split_openai_compatible_provider_types::Migration),
|
||||
Box::new(m20260515_000001_add_knowledge_base_index_schedule::Migration),
|
||||
Box::new(m20260518_000001_add_builtin_model_deletions::Migration),
|
||||
Box::new(m20260710_000001_agent_run_state::Migration),
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
use sea_orm_migration::prelude::*;
|
||||
|
||||
#[derive(DeriveMigrationName)]
|
||||
pub struct Migration;
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl MigrationTrait for Migration {
|
||||
async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
|
||||
for column in [
|
||||
ColumnDef::new(Alias::new("current_run_id")).string().null().to_owned(),
|
||||
ColumnDef::new(Alias::new("host_access_enabled"))
|
||||
.boolean()
|
||||
.not_null()
|
||||
.default(false)
|
||||
.to_owned(),
|
||||
ColumnDef::new(Alias::new("revision"))
|
||||
.big_integer()
|
||||
.not_null()
|
||||
.default(0)
|
||||
.to_owned(),
|
||||
ColumnDef::new(Alias::new("context_version"))
|
||||
.integer()
|
||||
.not_null()
|
||||
.default(1)
|
||||
.to_owned(),
|
||||
ColumnDef::new(Alias::new("last_provider_id"))
|
||||
.string()
|
||||
.null()
|
||||
.to_owned(),
|
||||
ColumnDef::new(Alias::new("last_model_id")).string().null().to_owned(),
|
||||
] {
|
||||
manager
|
||||
.alter_table(
|
||||
Table::alter()
|
||||
.table(Alias::new("agent_sessions"))
|
||||
.add_column(column)
|
||||
.to_owned(),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
|
||||
manager
|
||||
.create_table(
|
||||
Table::create()
|
||||
.table(Alias::new("agent_runs"))
|
||||
.if_not_exists()
|
||||
.col(
|
||||
ColumnDef::new(Alias::new("id"))
|
||||
.string()
|
||||
.not_null()
|
||||
.primary_key(),
|
||||
)
|
||||
.col(ColumnDef::new(Alias::new("session_id")).string().not_null())
|
||||
.col(
|
||||
ColumnDef::new(Alias::new("conversation_id"))
|
||||
.string()
|
||||
.not_null(),
|
||||
)
|
||||
.col(
|
||||
ColumnDef::new(Alias::new("status"))
|
||||
.string()
|
||||
.not_null()
|
||||
.default("starting"),
|
||||
)
|
||||
.col(ColumnDef::new(Alias::new("assistant_message_id")).string().null())
|
||||
.col(ColumnDef::new(Alias::new("error_message")).text().null())
|
||||
.col(ColumnDef::new(Alias::new("created_at")).string().not_null())
|
||||
.col(ColumnDef::new(Alias::new("updated_at")).string().not_null())
|
||||
.col(ColumnDef::new(Alias::new("finished_at")).string().null())
|
||||
.foreign_key(
|
||||
ForeignKey::create()
|
||||
.name("fk-agent-runs-session")
|
||||
.from(Alias::new("agent_runs"), Alias::new("session_id"))
|
||||
.to(Alias::new("agent_sessions"), Alias::new("id"))
|
||||
.on_delete(ForeignKeyAction::Cascade),
|
||||
)
|
||||
.to_owned(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
manager
|
||||
.create_index(
|
||||
Index::create()
|
||||
.name("idx-agent-runs-conversation-created")
|
||||
.table(Alias::new("agent_runs"))
|
||||
.col(Alias::new("conversation_id"))
|
||||
.col(Alias::new("created_at"))
|
||||
.to_owned(),
|
||||
)
|
||||
.await?;
|
||||
manager
|
||||
.create_index(
|
||||
Index::create()
|
||||
.name("idx-agent-sessions-conversation-unique")
|
||||
.table(Alias::new("agent_sessions"))
|
||||
.col(Alias::new("conversation_id"))
|
||||
.unique()
|
||||
.to_owned(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
|
||||
manager
|
||||
.drop_table(Table::drop().table(Alias::new("agent_runs")).to_owned())
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -1,21 +0,0 @@
|
||||
/target
|
||||
Cargo.lock
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
|
||||
# Environment
|
||||
.env
|
||||
.env.*
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
# Credentials
|
||||
*.pem
|
||||
*.key
|
||||
@@ -1,78 +0,0 @@
|
||||
[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"
|
||||
@@ -1,265 +0,0 @@
|
||||
# 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
|
||||
@@ -1,77 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,68 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,55 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,37 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,53 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,73 +0,0 @@
|
||||
// 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
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,186 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,34 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,70 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,51 +0,0 @@
|
||||
// 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;
|
||||
}
|
||||
@@ -1,590 +0,0 @@
|
||||
<!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>
|
||||
@@ -1,241 +0,0 @@
|
||||
// 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;
|
||||
@@ -1,408 +0,0 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -1,238 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,4 +0,0 @@
|
||||
mod agent;
|
||||
pub(crate) mod r#loop;
|
||||
|
||||
pub use agent::*;
|
||||
@@ -1,338 +0,0 @@
|
||||
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,
|
||||
}
|
||||
}
|
||||
@@ -1,285 +0,0 @@
|
||||
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(_))
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
pub mod anthropic;
|
||||
mod client;
|
||||
pub mod openai;
|
||||
pub mod provider;
|
||||
|
||||
pub use client::*;
|
||||
pub use provider::{ApiType, LLMProvider, ProviderRequest, ProviderResponse};
|
||||
@@ -1,385 +0,0 @@
|
||||
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,
|
||||
})
|
||||
}
|
||||
@@ -1,84 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,198 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,196 +0,0 @@
|
||||
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,
|
||||
}
|
||||
@@ -1,297 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,59 +0,0 @@
|
||||
//! # 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};
|
||||
@@ -1,341 +0,0 @@
|
||||
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 }
|
||||
}
|
||||
@@ -1,6 +0,0 @@
|
||||
mod client;
|
||||
pub mod shell_path;
|
||||
mod tool;
|
||||
|
||||
pub use client::*;
|
||||
pub use tool::*;
|
||||
@@ -1,185 +0,0 @@
|
||||
//! 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");
|
||||
}
|
||||
}
|
||||
@@ -1,109 +0,0 @@
|
||||
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()
|
||||
}
|
||||
@@ -1,104 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,5 +0,0 @@
|
||||
mod filesystem;
|
||||
mod permissions;
|
||||
|
||||
pub use filesystem::*;
|
||||
pub use permissions::*;
|
||||
@@ -1,119 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,253 +0,0 @@
|
||||
//! # 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,
|
||||
}
|
||||
}
|
||||
@@ -1,159 +0,0 @@
|
||||
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)
|
||||
}
|
||||
@@ -1,5 +0,0 @@
|
||||
pub mod loader;
|
||||
pub mod registry;
|
||||
|
||||
pub use loader::*;
|
||||
pub use registry::*;
|
||||
@@ -1,102 +0,0 @@
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -1,107 +0,0 @@
|
||||
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(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,368 +0,0 @@
|
||||
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()],
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,136 +0,0 @@
|
||||
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))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,247 +0,0 @@
|
||||
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")))
|
||||
}
|
||||
}
|
||||
@@ -1,277 +0,0 @@
|
||||
/// 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
|
||||
}
|
||||
@@ -1,231 +0,0 @@
|
||||
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
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -1,177 +0,0 @@
|
||||
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())
|
||||
}
|
||||
@@ -1,182 +0,0 @@
|
||||
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())
|
||||
}
|
||||
@@ -1,131 +0,0 @@
|
||||
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())
|
||||
}
|
||||
@@ -1,274 +0,0 @@
|
||||
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))
|
||||
}
|
||||
@@ -1,290 +0,0 @@
|
||||
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())
|
||||
}
|
||||
@@ -1,245 +0,0 @@
|
||||
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()
|
||||
}
|
||||
@@ -1,174 +0,0 @@
|
||||
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
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -1,28 +0,0 @@
|
||||
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::*;
|
||||
@@ -1,180 +0,0 @@
|
||||
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
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -1,154 +0,0 @@
|
||||
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))
|
||||
}
|
||||
}
|
||||
@@ -1,181 +0,0 @@
|
||||
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()
|
||||
}
|
||||
}
|
||||
@@ -1,140 +0,0 @@
|
||||
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)))
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,111 +0,0 @@
|
||||
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))
|
||||
}
|
||||
}
|
||||
@@ -1,452 +0,0 @@
|
||||
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))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,188 +0,0 @@
|
||||
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))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,197 +0,0 @@
|
||||
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))),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,148 +0,0 @@
|
||||
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))
|
||||
}
|
||||
}
|
||||
@@ -1,101 +0,0 @@
|
||||
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))
|
||||
}
|
||||
}
|
||||
@@ -1,123 +0,0 @@
|
||||
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(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,255 +0,0 @@
|
||||
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
|
||||
)))
|
||||
}
|
||||
}
|
||||
@@ -1,69 +0,0 @@
|
||||
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,
|
||||
}
|
||||
@@ -1,273 +0,0 @@
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,9 +0,0 @@
|
||||
pub mod message;
|
||||
pub mod mcp;
|
||||
pub mod skill;
|
||||
pub mod tool;
|
||||
|
||||
pub use message::*;
|
||||
pub use mcp::*;
|
||||
pub use skill::*;
|
||||
pub use tool::*;
|
||||
@@ -1,95 +0,0 @@
|
||||
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>,
|
||||
}
|
||||
@@ -1,240 +0,0 @@
|
||||
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
|
||||
}
|
||||
}
|
||||
@@ -1,222 +0,0 @@
|
||||
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()
|
||||
}
|
||||
@@ -1,226 +0,0 @@
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -1,99 +0,0 @@
|
||||
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])
|
||||
}
|
||||
}
|
||||
@@ -1,5 +0,0 @@
|
||||
pub mod compact;
|
||||
pub mod file_cache;
|
||||
pub mod messages;
|
||||
pub mod retry;
|
||||
pub mod tokens;
|
||||
@@ -1,103 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,140 +0,0 @@
|
||||
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
|
||||
}
|
||||
@@ -1,226 +0,0 @@
|
||||
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;
|
||||
}
|
||||
@@ -1,195 +0,0 @@
|
||||
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"));
|
||||
}
|
||||
@@ -1,181 +0,0 @@
|
||||
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));
|
||||
}
|
||||
@@ -1,157 +0,0 @@
|
||||
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());
|
||||
}
|
||||
@@ -1,686 +0,0 @@
|
||||
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);
|
||||
}
|
||||
@@ -1,210 +0,0 @@
|
||||
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");
|
||||
}
|
||||
@@ -1,273 +0,0 @@
|
||||
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);
|
||||
}
|
||||
+1022
-149
File diff suppressed because it is too large
Load Diff
@@ -37,9 +37,9 @@ pub struct AppState {
|
||||
pub agent_cancel_tokens:
|
||||
Arc<Mutex<std::collections::HashMap<String, open_agent_sdk::CancellationToken>>>,
|
||||
pub agent_permission_senders:
|
||||
Arc<Mutex<std::collections::HashMap<String, tokio::sync::oneshot::Sender<String>>>>,
|
||||
Arc<Mutex<std::collections::HashMap<String, commands::agent::PendingAgentPermission>>>,
|
||||
pub agent_ask_senders:
|
||||
Arc<Mutex<std::collections::HashMap<String, tokio::sync::oneshot::Sender<String>>>>,
|
||||
Arc<Mutex<std::collections::HashMap<String, commands::agent::PendingAgentAsk>>>,
|
||||
pub agent_always_allowed:
|
||||
Arc<Mutex<std::collections::HashMap<String, std::collections::HashSet<String>>>>,
|
||||
}
|
||||
@@ -669,6 +669,7 @@ pub fn run() {
|
||||
{
|
||||
let sea_db = app.state::<AppState>().sea_db.clone();
|
||||
let _ = rt.block_on(boosapi_core::repo::agent_session::reset_running_sessions(&sea_db));
|
||||
let _ = rt.block_on(boosapi_core::repo::agent_run::interrupt_active_runs(&sea_db));
|
||||
match rt.block_on(boosapi_core::repo::message::mark_stale_partial_assistant_messages_failed(&sea_db)) {
|
||||
Ok(count) if count > 0 => {
|
||||
tracing::info!(
|
||||
|
||||
@@ -1,526 +0,0 @@
|
||||
import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
|
||||
import { App, Button, Dropdown, Input, Select, Switch, Tag, theme } from 'antd';
|
||||
import type { MenuProps } from 'antd';
|
||||
import {
|
||||
ArrowUp,
|
||||
AtSign,
|
||||
ChevronDown,
|
||||
Clapperboard,
|
||||
ImagePlus,
|
||||
Mic2,
|
||||
Music2,
|
||||
Plus,
|
||||
SlidersHorizontal,
|
||||
Sparkles,
|
||||
UserRound,
|
||||
WandSparkles,
|
||||
Wrench,
|
||||
X,
|
||||
} from 'lucide-react';
|
||||
import { invoke } from '@/lib/invoke';
|
||||
import { findSkillByPreset, getSkillPresetStatus, SKILL_PRESETS, type SkillPreset } from '@/lib/skillPresets';
|
||||
import { useConversationStore, useProviderStore, useSkillStore } from '@/stores';
|
||||
import type { AttachmentInput, Model } from '@/types';
|
||||
|
||||
type CreationMode = 'agent' | 'image' | 'video' | 'music' | 'voice' | 'digital-human' | 'action';
|
||||
|
||||
interface CreationModeOption {
|
||||
key: CreationMode;
|
||||
label: string;
|
||||
hint: string;
|
||||
icon: React.ReactNode;
|
||||
}
|
||||
|
||||
interface AttachmentItem {
|
||||
id: string;
|
||||
name: string;
|
||||
type: string;
|
||||
size: number;
|
||||
data: string;
|
||||
}
|
||||
|
||||
interface ChatModelOption {
|
||||
value: string;
|
||||
label: string;
|
||||
providerId: string;
|
||||
modelId: string;
|
||||
}
|
||||
|
||||
const CREATION_MODES: CreationModeOption[] = [
|
||||
{ key: 'agent', label: 'Agent 模式', hint: '拆解任务、读写文件、执行工具', icon: <WandSparkles size={16} /> },
|
||||
{ key: 'image', label: '图片生成', hint: '即梦图片、参考图与风格控制', icon: <ImagePlus size={16} /> },
|
||||
{ key: 'video', label: '视频生成', hint: '即梦视频、图生视频与分镜', icon: <Clapperboard size={16} /> },
|
||||
{ key: 'music', label: '音乐生成', hint: '配乐、氛围与短片 BGM', icon: <Music2 size={16} /> },
|
||||
{ key: 'voice', label: '配音生成', hint: 'TTS 音色、语速与试听', icon: <Mic2 size={16} /> },
|
||||
{ key: 'digital-human', label: '数字人', hint: '口播、形象与视频工作流', icon: <UserRound size={16} /> },
|
||||
{ key: 'action', label: '动作模仿', hint: '动作迁移与主体参考', icon: <Sparkles size={16} /> },
|
||||
];
|
||||
|
||||
const JIMENG_IMAGE_MODELS = [
|
||||
{ value: 'jimeng_t2i_v40', label: '即梦图片 4.0' },
|
||||
{ value: 'custom', label: '自定义 req_key' },
|
||||
];
|
||||
|
||||
const JIMENG_VIDEO_MODELS = [
|
||||
{ value: 'jimeng_ti2v_v30_pro', label: '即梦图生视频 3.0 Pro' },
|
||||
{ value: 'custom', label: '自定义 req_key' },
|
||||
];
|
||||
|
||||
function formatFileSize(size: number): string {
|
||||
if (size < 1024) return `${size} B`;
|
||||
if (size < 1024 * 1024) return `${(size / 1024).toFixed(1)} KB`;
|
||||
return `${(size / 1024 / 1024).toFixed(1)} MB`;
|
||||
}
|
||||
|
||||
function modelOptionLabel(model: Model, providerName: string): string {
|
||||
return `${model.name || model.model_id} · ${providerName}`;
|
||||
}
|
||||
|
||||
function isUsableChatModel(model: Model): boolean {
|
||||
return model.enabled && (model.model_type === 'Chat' || model.capabilities.includes('TextChat'));
|
||||
}
|
||||
|
||||
function getAttachmentMimeType(fileName: string, mimeType?: string): string {
|
||||
if (mimeType) return mimeType;
|
||||
const extension = fileName.split('.').pop()?.toLowerCase();
|
||||
if (extension === 'pdf') return 'application/pdf';
|
||||
if (extension === 'txt') return 'text/plain';
|
||||
if (extension === 'md') return 'text/markdown';
|
||||
if (extension === 'json') return 'application/json';
|
||||
return 'application/octet-stream';
|
||||
}
|
||||
|
||||
async function fileToAttachmentItem(file: File): Promise<AttachmentItem> {
|
||||
return new Promise((resolve, reject) => {
|
||||
const reader = new FileReader();
|
||||
reader.onerror = () => reject(reader.error ?? new Error(`读取文件失败:${file.name}`));
|
||||
reader.onload = () => {
|
||||
const base64 = (reader.result as string).split(',')[1] || '';
|
||||
resolve({
|
||||
id: `${file.name}-${file.size}-${file.lastModified}-${globalThis.crypto?.randomUUID?.() ?? Date.now()}`,
|
||||
name: file.name,
|
||||
type: getAttachmentMimeType(file.name, file.type),
|
||||
size: file.size,
|
||||
data: base64,
|
||||
});
|
||||
};
|
||||
reader.readAsDataURL(file);
|
||||
});
|
||||
}
|
||||
|
||||
function toAttachmentInput(item: AttachmentItem): AttachmentInput {
|
||||
return {
|
||||
file_name: item.name,
|
||||
file_type: item.type,
|
||||
file_size: item.size,
|
||||
data: item.data,
|
||||
};
|
||||
}
|
||||
|
||||
export function AgentWorkbenchComposer() {
|
||||
const { token } = theme.useToken();
|
||||
const { message } = App.useApp();
|
||||
const fileInputRef = useRef<HTMLInputElement>(null);
|
||||
const providers = useProviderStore((s) => s.providers);
|
||||
const fetchProviders = useProviderStore((s) => s.fetchProviders);
|
||||
const createConversation = useConversationStore((s) => s.createConversation);
|
||||
const updateConversation = useConversationStore((s) => s.updateConversation);
|
||||
const sendAgentMessage = useConversationStore((s) => s.sendAgentMessage);
|
||||
const skills = useSkillStore((s) => s.skills);
|
||||
const loadSkills = useSkillStore((s) => s.loadSkills);
|
||||
const installSkill = useSkillStore((s) => s.installSkill);
|
||||
const toggleSkill = useSkillStore((s) => s.toggleSkill);
|
||||
const [creationMode, setCreationMode] = useState<CreationMode>('agent');
|
||||
const [autoMode, setAutoMode] = useState(true);
|
||||
const [prompt, setPrompt] = useState('');
|
||||
const [agentModel, setAgentModel] = useState('custom');
|
||||
const [customAgentModel, setCustomAgentModel] = useState('');
|
||||
const [imageModel, setImageModel] = useState('jimeng_t2i_v40');
|
||||
const [videoModel, setVideoModel] = useState('jimeng_ti2v_v30_pro');
|
||||
const [customJimengModel, setCustomJimengModel] = useState('');
|
||||
const [attachments, setAttachments] = useState<AttachmentItem[]>([]);
|
||||
const [subjectOpen, setSubjectOpen] = useState(false);
|
||||
const [skillOpen, setSkillOpen] = useState(false);
|
||||
const [submitting, setSubmitting] = useState(false);
|
||||
const [installingSkillKey, setInstallingSkillKey] = useState<string | null>(null);
|
||||
const [result, setResult] = useState('');
|
||||
|
||||
useEffect(() => {
|
||||
if (providers.length === 0) {
|
||||
void fetchProviders().catch((error) => message.error(String(error)));
|
||||
}
|
||||
}, [fetchProviders, message, providers.length]);
|
||||
|
||||
useEffect(() => {
|
||||
void loadSkills();
|
||||
}, [loadSkills]);
|
||||
|
||||
const chatModelOptions = useMemo<ChatModelOption[]>(() => {
|
||||
const rows = providers
|
||||
.filter((provider) => provider.enabled)
|
||||
.flatMap((provider) =>
|
||||
provider.models
|
||||
.filter(isUsableChatModel)
|
||||
.map((model) => ({
|
||||
value: `${provider.id}::${model.model_id}`,
|
||||
label: modelOptionLabel(model, provider.name),
|
||||
providerId: provider.id,
|
||||
modelId: model.model_id,
|
||||
})),
|
||||
);
|
||||
const seen = new Set<string>();
|
||||
return rows.filter((item) => {
|
||||
if (seen.has(item.value)) return false;
|
||||
seen.add(item.value);
|
||||
return true;
|
||||
});
|
||||
}, [providers]);
|
||||
|
||||
useEffect(() => {
|
||||
if (agentModel !== 'custom') return;
|
||||
if (customAgentModel || chatModelOptions.length === 0) return;
|
||||
setAgentModel(chatModelOptions[0].value);
|
||||
}, [agentModel, chatModelOptions, customAgentModel]);
|
||||
|
||||
const currentMode = CREATION_MODES.find((item) => item.key === creationMode) ?? CREATION_MODES[0];
|
||||
const selectedModel = creationMode === 'agent'
|
||||
? (agentModel === 'custom' ? customAgentModel.trim() || '自定义模型' : agentModel)
|
||||
: creationMode === 'image'
|
||||
? (imageModel === 'custom' ? customJimengModel.trim() || '自定义 req_key' : imageModel)
|
||||
: creationMode === 'video'
|
||||
? (videoModel === 'custom' ? customJimengModel.trim() || '自定义 req_key' : videoModel)
|
||||
: '待配置';
|
||||
|
||||
const modeItems: MenuProps['items'] = CREATION_MODES.map((item) => ({
|
||||
key: item.key,
|
||||
label: (
|
||||
<div className="agent-workbench-mode-item">
|
||||
<span>{item.icon}</span>
|
||||
<div>
|
||||
<strong>{item.label}</strong>
|
||||
<small>{item.hint}</small>
|
||||
</div>
|
||||
{item.key === creationMode && <em>✓</em>}
|
||||
</div>
|
||||
),
|
||||
}));
|
||||
|
||||
const selectedChatModel = useMemo(() => (
|
||||
chatModelOptions.find((item) => item.value === agentModel) ?? null
|
||||
), [agentModel, chatModelOptions]);
|
||||
|
||||
const skillMenuPresets = useMemo(() => SKILL_PRESETS.filter((preset) => (
|
||||
preset.category === 'video' || preset.category === 'document' || preset.category === 'media'
|
||||
)), []);
|
||||
|
||||
const handleFilesSelected = useCallback((files: FileList | null) => {
|
||||
if (!files || files.length === 0) return;
|
||||
void Promise.all(Array.from(files).map(fileToAttachmentItem))
|
||||
.then((next) => {
|
||||
setAttachments((current) => [...current, ...next]);
|
||||
})
|
||||
.catch((error) => message.error(String(error)))
|
||||
.finally(() => {
|
||||
if (fileInputRef.current) fileInputRef.current.value = '';
|
||||
});
|
||||
}, [message]);
|
||||
|
||||
const removeAttachment = useCallback((id: string) => {
|
||||
setAttachments((current) => current.filter((item) => item.id !== id));
|
||||
}, []);
|
||||
|
||||
const applySkillPreset = useCallback(async (preset: SkillPreset) => {
|
||||
if (installingSkillKey) return;
|
||||
setInstallingSkillKey(preset.key);
|
||||
try {
|
||||
const existingSkill = findSkillByPreset(useSkillStore.getState().skills, preset);
|
||||
if (!existingSkill) {
|
||||
await installSkill(preset.source);
|
||||
await loadSkills();
|
||||
message.success(`${preset.title} Skill 已安装`);
|
||||
} else if (!existingSkill.enabled) {
|
||||
await toggleSkill(preset.name, true);
|
||||
await loadSkills();
|
||||
message.success(`${preset.title} Skill 已启用`);
|
||||
}
|
||||
setCreationMode('agent');
|
||||
setPrompt((value) => {
|
||||
const text = value.trim();
|
||||
if (text.includes(preset.name) || text.includes(preset.title)) return value;
|
||||
return text ? `${preset.promptPrefix}${text}` : preset.promptPrefix;
|
||||
});
|
||||
setSkillOpen(false);
|
||||
} catch (error) {
|
||||
message.error(`${preset.title} Skill 集成失败:${String(error)}`);
|
||||
} finally {
|
||||
setInstallingSkillKey(null);
|
||||
}
|
||||
}, [installSkill, installingSkillKey, loadSkills, message, toggleSkill]);
|
||||
|
||||
const submit = useCallback(async () => {
|
||||
const text = prompt.trim();
|
||||
if (!text) {
|
||||
message.warning('先输入想法、剧本或任务说明');
|
||||
return;
|
||||
}
|
||||
if (submitting) return;
|
||||
|
||||
if (creationMode !== 'agent') {
|
||||
setResult(`${currentMode.label} 已准备好。模型:${selectedModel};参考素材:${attachments.length} 个。即梦/媒体生成链路下一步接入。`);
|
||||
return;
|
||||
}
|
||||
|
||||
const fallbackProvider = providers.find((provider) => provider.enabled);
|
||||
const providerId = selectedChatModel?.providerId ?? fallbackProvider?.id;
|
||||
const modelId = selectedChatModel?.modelId ?? customAgentModel.trim();
|
||||
if (!providerId || !modelId) {
|
||||
message.warning('先在供应商里启用一个对话模型,或填写自定义模型名');
|
||||
return;
|
||||
}
|
||||
|
||||
setSubmitting(true);
|
||||
try {
|
||||
const title = text.slice(0, 24) || 'Agent 创作';
|
||||
const conversation = await createConversation(title, modelId, providerId);
|
||||
await updateConversation(conversation.id, { mode: 'agent', provider_id: providerId, model_id: modelId });
|
||||
await invoke('agent_update_session', {
|
||||
conversationId: conversation.id,
|
||||
permissionMode: autoMode ? 'accept_edits' : 'default',
|
||||
});
|
||||
await sendAgentMessage(text, attachments.map(toAttachmentInput));
|
||||
setResult(`Agent 已开始执行。模型:${modelId};参考素材:${attachments.length} 个。`);
|
||||
setPrompt('');
|
||||
setAttachments([]);
|
||||
} catch (error) {
|
||||
message.error(String(error));
|
||||
setResult(`Agent 启动失败:${String(error)}`);
|
||||
} finally {
|
||||
setSubmitting(false);
|
||||
}
|
||||
}, [
|
||||
attachments,
|
||||
autoMode,
|
||||
createConversation,
|
||||
creationMode,
|
||||
currentMode.label,
|
||||
customAgentModel,
|
||||
message,
|
||||
prompt,
|
||||
providers,
|
||||
selectedChatModel,
|
||||
selectedModel,
|
||||
sendAgentMessage,
|
||||
submitting,
|
||||
updateConversation,
|
||||
]);
|
||||
|
||||
return (
|
||||
<div className="agent-workbench-desktop" style={{ background: token.colorBgContainer }}>
|
||||
<div className="agent-workbench-desktop-bg" />
|
||||
<section className="agent-workbench-desktop-hero">
|
||||
<h1>你好,想创作什么?</h1>
|
||||
<div className="agent-workbench-desktop-composer">
|
||||
<input
|
||||
ref={fileInputRef}
|
||||
className="agent-workbench-file-input"
|
||||
type="file"
|
||||
multiple
|
||||
accept="image/*,video/*,audio/*,.pdf,.txt,.md,.json"
|
||||
onChange={(event) => handleFilesSelected(event.target.files)}
|
||||
/>
|
||||
<button
|
||||
className="agent-workbench-upload"
|
||||
type="button"
|
||||
onClick={() => fileInputRef.current?.click()}
|
||||
aria-label="上传参考素材"
|
||||
>
|
||||
<Plus size={24} />
|
||||
</button>
|
||||
<Input.TextArea
|
||||
className="agent-workbench-prompt"
|
||||
value={prompt}
|
||||
onChange={(event) => setPrompt(event.target.value)}
|
||||
onPressEnter={(event) => {
|
||||
if (event.metaKey || event.ctrlKey) {
|
||||
event.preventDefault();
|
||||
void submit();
|
||||
}
|
||||
}}
|
||||
placeholder="输入想法、剧本或上传参考,支持 “/” 使用技能,@ 添加主体,和 Agent 一起创作"
|
||||
autoSize={{ minRows: 4, maxRows: 8 }}
|
||||
/>
|
||||
{attachments.length > 0 && (
|
||||
<div className="agent-workbench-attachments">
|
||||
{attachments.map((item) => (
|
||||
<Tag
|
||||
key={item.id}
|
||||
closeIcon={<X size={12} />}
|
||||
onClose={(event) => {
|
||||
event.preventDefault();
|
||||
removeAttachment(item.id);
|
||||
}}
|
||||
>
|
||||
{item.name} <span>{formatFileSize(item.size)}</span>
|
||||
</Tag>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
<div className="agent-workbench-toolbar">
|
||||
<div className="agent-workbench-tool-left">
|
||||
<Dropdown
|
||||
trigger={['click']}
|
||||
menu={{
|
||||
items: modeItems,
|
||||
selectable: true,
|
||||
selectedKeys: [creationMode],
|
||||
onClick: ({ key }) => setCreationMode(key as CreationMode),
|
||||
}}
|
||||
overlayClassName="agent-workbench-mode-menu"
|
||||
>
|
||||
<Button className="agent-workbench-pill active" icon={currentMode.icon}>
|
||||
{currentMode.label} <ChevronDown size={14} />
|
||||
</Button>
|
||||
</Dropdown>
|
||||
<Button className="agent-workbench-pill" icon={<SlidersHorizontal size={15} />} onClick={() => setAutoMode((value) => !value)}>
|
||||
自动
|
||||
<Switch size="small" checked={autoMode} onChange={setAutoMode} onClick={(_, event) => event.stopPropagation()} />
|
||||
</Button>
|
||||
{creationMode === 'agent' && (
|
||||
<>
|
||||
<Select
|
||||
className="agent-workbench-model-select"
|
||||
value={agentModel}
|
||||
onChange={(value) => {
|
||||
setAgentModel(value);
|
||||
if (value !== 'custom') setCustomAgentModel('');
|
||||
}}
|
||||
options={[
|
||||
...chatModelOptions,
|
||||
{ value: 'custom', label: '自定义模型' },
|
||||
]}
|
||||
popupMatchSelectWidth={false}
|
||||
/>
|
||||
{agentModel === 'custom' && (
|
||||
<Input
|
||||
className="agent-workbench-model-input"
|
||||
value={customAgentModel}
|
||||
onChange={(event) => setCustomAgentModel(event.target.value)}
|
||||
placeholder="输入 BoosAPI 模型名"
|
||||
/>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
{creationMode === 'image' && (
|
||||
<>
|
||||
<Select
|
||||
className="agent-workbench-model-select"
|
||||
value={imageModel}
|
||||
onChange={setImageModel}
|
||||
options={JIMENG_IMAGE_MODELS}
|
||||
popupMatchSelectWidth={false}
|
||||
/>
|
||||
{imageModel === 'custom' && (
|
||||
<Input
|
||||
className="agent-workbench-model-input"
|
||||
value={customJimengModel}
|
||||
onChange={(event) => setCustomJimengModel(event.target.value)}
|
||||
placeholder="输入火山 req_key"
|
||||
/>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
{creationMode === 'video' && (
|
||||
<>
|
||||
<Select
|
||||
className="agent-workbench-model-select"
|
||||
value={videoModel}
|
||||
onChange={setVideoModel}
|
||||
options={JIMENG_VIDEO_MODELS}
|
||||
popupMatchSelectWidth={false}
|
||||
/>
|
||||
{videoModel === 'custom' && (
|
||||
<Input
|
||||
className="agent-workbench-model-input"
|
||||
value={customJimengModel}
|
||||
onChange={(event) => setCustomJimengModel(event.target.value)}
|
||||
placeholder="输入火山 req_key"
|
||||
/>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
<Dropdown
|
||||
open={skillOpen}
|
||||
onOpenChange={setSkillOpen}
|
||||
trigger={['click']}
|
||||
dropdownRender={() => (
|
||||
<div className="agent-workbench-popover">
|
||||
<Input placeholder="搜索技能" prefix={<Wrench size={15} />} />
|
||||
{skillMenuPresets.map((preset) => {
|
||||
const status = getSkillPresetStatus(skills, preset);
|
||||
return (
|
||||
<button key={preset.key} type="button" onClick={() => void applySkillPreset(preset)}>
|
||||
{preset.key === 'jianying' ? <Clapperboard size={16} /> : <Wrench size={16} />}
|
||||
{preset.title}
|
||||
<small>
|
||||
{installingSkillKey === preset.key
|
||||
? '正在安装/启用...'
|
||||
: status === 'enabled'
|
||||
? `已启用,${preset.description}`
|
||||
: status === 'installed'
|
||||
? '已安装,点击启用并写入提示'
|
||||
: `安装/启用 ${preset.name} Skill`}
|
||||
</small>
|
||||
</button>
|
||||
);
|
||||
})}
|
||||
<button type="button"><Plus size={16} /> 用 Agent 创建技能</button>
|
||||
</div>
|
||||
)}
|
||||
>
|
||||
<Button className="agent-workbench-pill" loading={Boolean(installingSkillKey)} icon={<Wrench size={15} />}>使用技能</Button>
|
||||
</Dropdown>
|
||||
<Dropdown
|
||||
open={subjectOpen}
|
||||
onOpenChange={setSubjectOpen}
|
||||
trigger={['click']}
|
||||
dropdownRender={() => (
|
||||
<div className="agent-workbench-subject-popover">
|
||||
<span>可能@的内容</span>
|
||||
<div>
|
||||
<AtSign size={36} />
|
||||
<strong>你还没有创建过主体</strong>
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
setPrompt((value) => `${value}${value && !value.endsWith(' ') ? ' ' : ''}@主体 `);
|
||||
setSubjectOpen(false);
|
||||
}}
|
||||
>
|
||||
+ 创建主体
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
>
|
||||
<Button className="agent-workbench-pill square" icon={<AtSign size={16} />} />
|
||||
</Dropdown>
|
||||
</div>
|
||||
<Button
|
||||
className="agent-workbench-submit"
|
||||
type="primary"
|
||||
icon={<ArrowUp size={18} />}
|
||||
loading={submitting}
|
||||
onClick={() => void submit()}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
{result && (
|
||||
<div className="agent-workbench-result">
|
||||
<strong>{currentMode.label}</strong>
|
||||
<p>{result}</p>
|
||||
</div>
|
||||
)}
|
||||
</section>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -1,5 +1,5 @@
|
||||
import React, { useState, useMemo } from 'react';
|
||||
import { Card, Input, Button, Space, Typography, Checkbox, theme } from 'antd';
|
||||
import { Alert, Card, Input, Button, Space, Typography, Checkbox, theme } from 'antd';
|
||||
import { MessageCircleQuestion, CheckCircle2 } from 'lucide-react';
|
||||
import { useAgentStore } from '@/stores';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
@@ -21,6 +21,7 @@ const AskUserCard: React.FC<AskUserCardProps> = ({ askId, question, options }) =
|
||||
const [selectedOptions, setSelectedOptions] = useState<string[]>([]);
|
||||
const [submitting, setSubmitting] = useState(false);
|
||||
const [submitted, setSubmitted] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const respondAskUser = useAgentStore((s) => s.respondAskUser);
|
||||
|
||||
const hasOptions = options && options.length > 0;
|
||||
@@ -51,10 +52,13 @@ const AskUserCard: React.FC<AskUserCardProps> = ({ askId, question, options }) =
|
||||
const finalAnswer = buildAnswer();
|
||||
if (!finalAnswer || submitting || submitted) return;
|
||||
setSubmitting(true);
|
||||
setError(null);
|
||||
try {
|
||||
await respondAskUser(askId, finalAnswer);
|
||||
setSubmitted(true);
|
||||
} catch {
|
||||
} catch (submitError) {
|
||||
setError(submitError instanceof Error ? submitError.message : String(submitError));
|
||||
} finally {
|
||||
setSubmitting(false);
|
||||
}
|
||||
};
|
||||
@@ -114,6 +118,7 @@ const AskUserCard: React.FC<AskUserCardProps> = ({ askId, question, options }) =
|
||||
}
|
||||
autoSize={{ minRows: 1, maxRows: 4 }}
|
||||
disabled={submitting || submitted}
|
||||
autoFocus
|
||||
style={{ borderColor: token.colorPrimary }}
|
||||
onPressEnter={(e) => {
|
||||
if (!e.shiftKey) {
|
||||
@@ -122,6 +127,7 @@ const AskUserCard: React.FC<AskUserCardProps> = ({ askId, question, options }) =
|
||||
}
|
||||
}}
|
||||
/>
|
||||
{error && <Alert type="error" showIcon message={error} aria-live="polite" />}
|
||||
<div style={{ display: 'flex', justifyContent: 'flex-end' }}>
|
||||
{submitted ? (
|
||||
<Space size={4} style={{ color: token.colorSuccess }}>
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user