Michael Neale
1479 lines
56 KiB
Rust
1479 lines
56 KiB
Rust
use anyhow::Result;
|
|
use async_stream::try_stream;
|
|
use async_trait::async_trait;
|
|
use futures::future::BoxFuture;
|
|
use rmcp::model::{Role, Tool};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::{json, Value};
|
|
use std::collections::HashMap;
|
|
use std::io::Write;
|
|
use std::path::{Path, PathBuf};
|
|
use std::process::Stdio;
|
|
use std::sync::Arc;
|
|
use tempfile::NamedTempFile;
|
|
use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader};
|
|
use tokio::process::Command;
|
|
use tokio::sync::oneshot;
|
|
|
|
use super::base::{
|
|
stream_from_single_message, ConfigKey, MessageStream, PermissionRouting, Provider, ProviderDef,
|
|
ProviderMetadata, ProviderUsage, Usage,
|
|
};
|
|
use super::errors::ProviderError;
|
|
use super::utils::filter_extensions_from_system_prompt;
|
|
use crate::config::base::ClaudeCodeCommand;
|
|
use crate::config::paths::Paths;
|
|
use crate::config::search_path::SearchPaths;
|
|
use crate::config::{Config, ExtensionConfig, GooseMode};
|
|
use crate::conversation::message::{Message, MessageContent};
|
|
use crate::model::ModelConfig;
|
|
use crate::permission::permission_confirmation::PrincipalType;
|
|
use crate::permission::{Permission, PermissionConfirmation};
|
|
use crate::subprocess::configure_subprocess;
|
|
|
|
use super::cli_common::{error_from_event, extract_usage_tokens};
|
|
|
|
const CLAUDE_CODE_PROVIDER_NAME: &str = "claude-code";
|
|
pub const CLAUDE_CODE_DEFAULT_MODEL: &str = "default";
|
|
pub const CLAUDE_CODE_DOC_URL: &str = "https://code.claude.com/docs/en/setup";
|
|
|
|
// https://github.com/anthropics/claude-agent-sdk-python/blob/0e9397e/src/claude_agent_sdk/types.py#L857-L859
|
|
#[derive(Serialize)]
|
|
struct ControlResponse<T: Serialize> {
|
|
#[serde(rename = "type")]
|
|
msg_type: &'static str,
|
|
response: ControlResponseBody<T>,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct ControlResponseBody<T: Serialize> {
|
|
subtype: &'static str,
|
|
request_id: String,
|
|
response: T,
|
|
}
|
|
|
|
// https://github.com/anthropics/claude-agent-sdk-python/blob/0e9397e/src/claude_agent_sdk/types.py#L135-L153
|
|
#[derive(Serialize)]
|
|
#[serde(tag = "behavior")]
|
|
enum PermissionResponse {
|
|
#[serde(rename = "allow")]
|
|
Allow {
|
|
#[serde(rename = "updatedInput")]
|
|
updated_input: serde_json::Map<String, Value>,
|
|
#[serde(rename = "toolUseID")]
|
|
tool_use_id: String,
|
|
},
|
|
#[serde(rename = "deny")]
|
|
Deny { message: String },
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
struct ControlRequest {
|
|
#[serde(rename = "type")]
|
|
msg_type: &'static str,
|
|
request_id: String,
|
|
request: ControlRequestBody,
|
|
}
|
|
|
|
#[derive(Serialize)]
|
|
#[serde(tag = "subtype")]
|
|
enum ControlRequestBody {
|
|
#[serde(rename = "initialize")]
|
|
Initialize,
|
|
#[serde(rename = "set_model")]
|
|
SetModel { model: String },
|
|
}
|
|
|
|
impl ControlRequestBody {
|
|
fn label(&self) -> &'static str {
|
|
match self {
|
|
Self::Initialize => "initialize",
|
|
Self::SetModel { .. } => "set_model",
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct IncomingControlResponse {
|
|
response: IncomingControlResponseBody,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[serde(tag = "subtype")]
|
|
enum IncomingControlResponseBody {
|
|
#[serde(rename = "success")]
|
|
Success {
|
|
request_id: String,
|
|
#[serde(default)]
|
|
response: Option<Value>,
|
|
},
|
|
#[serde(rename = "error")]
|
|
Error {
|
|
request_id: String,
|
|
#[serde(default)]
|
|
error: String,
|
|
},
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
struct IncomingControlRequest {
|
|
request_id: String,
|
|
request: IncomingRequestBody,
|
|
}
|
|
|
|
#[derive(Deserialize)]
|
|
#[serde(tag = "subtype")]
|
|
enum IncomingRequestBody {
|
|
#[serde(rename = "can_use_tool")]
|
|
CanUseTool {
|
|
tool_name: String,
|
|
#[serde(default)]
|
|
input: serde_json::Map<String, Value>,
|
|
#[serde(default)]
|
|
tool_use_id: String,
|
|
},
|
|
}
|
|
|
|
impl<T: Serialize> ControlResponse<T> {
|
|
fn success(request_id: String, response: T) -> Self {
|
|
Self {
|
|
msg_type: "control_response",
|
|
response: ControlResponseBody {
|
|
subtype: "success",
|
|
request_id,
|
|
response,
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
struct CliProcess {
|
|
child: tokio::process::Child,
|
|
stdin: Box<dyn tokio::io::AsyncWrite + Unpin + Send>,
|
|
reader: BufReader<Box<dyn tokio::io::AsyncRead + Unpin + Send>>,
|
|
#[allow(dead_code)]
|
|
stderr_handle: tokio::task::JoinHandle<String>,
|
|
current_model: String,
|
|
log_model_update: bool,
|
|
next_request_id: u64,
|
|
needs_drain: bool,
|
|
}
|
|
|
|
impl std::fmt::Debug for CliProcess {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("CliProcess")
|
|
.field("current_model", &self.current_model)
|
|
.field("next_request_id", &self.next_request_id)
|
|
.finish_non_exhaustive()
|
|
}
|
|
}
|
|
|
|
impl CliProcess {
|
|
fn next_request_id(&mut self) -> String {
|
|
let id = self.next_request_id;
|
|
self.next_request_id += 1;
|
|
format!("req_{id}")
|
|
}
|
|
|
|
async fn send_control_request(
|
|
&mut self,
|
|
body: ControlRequestBody,
|
|
) -> Result<Option<Value>, ProviderError> {
|
|
let request_id = self.next_request_id();
|
|
exchange_control(&mut self.stdin, &mut self.reader, &request_id, body).await
|
|
}
|
|
|
|
async fn send_set_model(&mut self, model: &str) -> Result<(), ProviderError> {
|
|
if model == self.current_model {
|
|
return Ok(());
|
|
}
|
|
self.send_control_request(ControlRequestBody::SetModel {
|
|
model: model.to_string(),
|
|
})
|
|
.await?;
|
|
self.current_model = model.to_string();
|
|
self.log_model_update = true;
|
|
Ok(())
|
|
}
|
|
|
|
async fn drain_pending_response(&mut self) {
|
|
if !self.needs_drain {
|
|
return;
|
|
}
|
|
tracing::debug!("Draining cancelled response from CLI process");
|
|
|
|
let drain = async {
|
|
let mut line = String::new();
|
|
loop {
|
|
line.clear();
|
|
match self.reader.read_line(&mut line).await {
|
|
Ok(0) => break,
|
|
Ok(_) => {
|
|
let trimmed = line.trim();
|
|
if trimmed.is_empty() {
|
|
continue;
|
|
}
|
|
if let Ok(parsed) = serde_json::from_str::<Value>(trimmed) {
|
|
match parsed.get("type").and_then(|t| t.as_str()) {
|
|
Some("result") | Some("error") => break,
|
|
_ => continue,
|
|
}
|
|
} else {
|
|
tracing::trace!(line = trimmed, "Non-JSON line during drain");
|
|
}
|
|
}
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
};
|
|
|
|
const DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
|
|
if tokio::time::timeout(DRAIN_TIMEOUT, drain).await.is_err() {
|
|
// CLI is still producing the old response. Leave needs_drain
|
|
// true so the next call retries — by then the old response
|
|
// likely completed and drain will succeed quickly.
|
|
tracing::warn!(
|
|
"Drain did not complete in {DRAIN_TIMEOUT:?}; \
|
|
will retry on next request"
|
|
);
|
|
return;
|
|
}
|
|
|
|
self.needs_drain = false;
|
|
tracing::debug!("Drain complete, protocol re-synced");
|
|
}
|
|
}
|
|
|
|
impl Drop for CliProcess {
|
|
fn drop(&mut self) {
|
|
self.stderr_handle.abort();
|
|
let _ = self.child.start_kill();
|
|
}
|
|
}
|
|
|
|
/// Spawns the Claude Code CLI (`claude`) as a persistent child process using
|
|
/// `--input-format stream-json --output-format stream-json`. The CLI stays alive
|
|
/// across turns, maintaining conversation state internally. Messages are sent as
|
|
/// NDJSON on stdin with content arrays supporting text and image blocks. Responses
|
|
/// are NDJSON on stdout (`assistant` + `result` events per turn).
|
|
#[derive(Debug, serde::Serialize)]
|
|
pub struct ClaudeCodeProvider {
|
|
command: PathBuf,
|
|
model: ModelConfig,
|
|
#[serde(skip)]
|
|
name: String,
|
|
/// Temp file holding MCP config JSON (auto-deleted on drop).
|
|
#[serde(skip)]
|
|
mcp_config_file: Option<NamedTempFile>,
|
|
#[serde(skip)]
|
|
cli_process: tokio::sync::OnceCell<Arc<tokio::sync::Mutex<CliProcess>>>,
|
|
#[serde(skip)]
|
|
pending_confirmations:
|
|
Arc<tokio::sync::Mutex<HashMap<String, oneshot::Sender<PermissionConfirmation>>>>,
|
|
#[serde(skip)]
|
|
initial_mode: tokio::sync::Mutex<Option<GooseMode>>,
|
|
}
|
|
|
|
impl ClaudeCodeProvider {
|
|
/// Build content blocks from the last user message only. The CLI maintains
|
|
/// conversation context internally per session_id.
|
|
fn last_user_content_blocks(&self, messages: &[Message]) -> Vec<Value> {
|
|
let msgs = match messages.iter().rev().find(|m| m.role == Role::User) {
|
|
Some(msg) => std::slice::from_ref(msg),
|
|
None => messages,
|
|
};
|
|
let mut blocks: Vec<Value> = Vec::new();
|
|
for message in msgs {
|
|
let prefix = match message.role {
|
|
Role::User => "Human: ",
|
|
Role::Assistant => "Assistant: ",
|
|
};
|
|
let mut text_parts = Vec::new();
|
|
for content in &message.content {
|
|
match content {
|
|
MessageContent::Text(t) => text_parts.push(t.text.clone()),
|
|
MessageContent::Image(img) => {
|
|
if !text_parts.is_empty() {
|
|
blocks.push(json!({"type":"text","text":format!("{}{}", prefix, text_parts.join("\n"))}));
|
|
text_parts.clear();
|
|
}
|
|
blocks.push(json!({"type":"image","source":{"type":"base64","media_type":img.mime_type,"data":img.data}}));
|
|
}
|
|
MessageContent::ToolRequest(req) => {
|
|
if let Ok(call) = &req.tool_call {
|
|
text_parts.push(format!("[tool_use: {} id={}]", call.name, req.id));
|
|
}
|
|
}
|
|
MessageContent::ToolResponse(resp) => {
|
|
if let Ok(result) = &resp.tool_result {
|
|
let text: String = result
|
|
.content
|
|
.iter()
|
|
.filter_map(|c| match &c.raw {
|
|
rmcp::model::RawContent::Text(t) => Some(t.text.as_str()),
|
|
_ => None,
|
|
})
|
|
.collect::<Vec<&str>>()
|
|
.join("\n");
|
|
text_parts.push(format!("[tool_result id={}] {}", resp.id, text));
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
if !text_parts.is_empty() {
|
|
blocks.push(
|
|
json!({"type":"text","text":format!("{}{}", prefix, text_parts.join("\n"))}),
|
|
);
|
|
}
|
|
}
|
|
blocks
|
|
}
|
|
|
|
fn build_stream_json_command(&self) -> Command {
|
|
let mut cmd = Command::new(&self.command);
|
|
configure_subprocess(&mut cmd);
|
|
// Allow goose to run inside a Claude Code session.
|
|
cmd.env_remove("CLAUDECODE");
|
|
cmd.arg("--input-format")
|
|
.arg("stream-json")
|
|
.arg("--output-format")
|
|
.arg("stream-json")
|
|
.arg("--verbose")
|
|
.stdin(Stdio::piped())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
cmd
|
|
}
|
|
|
|
/// Returns true when the control protocol is enabled.
|
|
fn apply_permission_flags(cmd: &mut Command) -> Result<bool, ProviderError> {
|
|
let config = Config::global();
|
|
let goose_mode = config.get_goose_mode().unwrap_or(GooseMode::Auto);
|
|
|
|
match goose_mode {
|
|
GooseMode::Auto => {
|
|
cmd.arg("--dangerously-skip-permissions");
|
|
Ok(false)
|
|
}
|
|
GooseMode::SmartApprove | GooseMode::Approve => {
|
|
cmd.arg("--permission-prompt-tool").arg("stdio");
|
|
Ok(true)
|
|
}
|
|
GooseMode::Chat => Ok(false),
|
|
}
|
|
}
|
|
|
|
async fn spawn_process(&self, filtered_system: &str) -> Result<CliProcess, ProviderError> {
|
|
let mut cmd = self.build_stream_json_command();
|
|
|
|
if let Some(f) = &self.mcp_config_file {
|
|
cmd.arg("--mcp-config").arg(f.path());
|
|
cmd.arg("--strict-mcp-config");
|
|
}
|
|
|
|
cmd.arg("--include-partial-messages")
|
|
.arg("--system-prompt")
|
|
.arg(filtered_system)
|
|
.arg("--model")
|
|
.arg(&self.model.model_name);
|
|
|
|
let control_protocol_enabled = Self::apply_permission_flags(&mut cmd)?;
|
|
|
|
let mut child = cmd.spawn().map_err(|e| {
|
|
ProviderError::RequestFailed(format!(
|
|
"Failed to spawn Claude CLI command '{:?}': {}.",
|
|
self.command, e
|
|
))
|
|
})?;
|
|
|
|
let stdin = child
|
|
.stdin
|
|
.take()
|
|
.ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdin".to_string()))?;
|
|
let stdout = child
|
|
.stdout
|
|
.take()
|
|
.ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdout".to_string()))?;
|
|
|
|
let stderr = child.stderr.take();
|
|
let stderr_handle = tokio::spawn(async move {
|
|
let mut output = String::new();
|
|
if let Some(mut stderr) = stderr {
|
|
use tokio::io::AsyncReadExt;
|
|
let _ = stderr.read_to_string(&mut output).await;
|
|
}
|
|
output
|
|
});
|
|
|
|
let mut process = CliProcess {
|
|
child,
|
|
stdin: Box::new(stdin),
|
|
reader: BufReader::new(Box::new(stdout)),
|
|
stderr_handle,
|
|
current_model: self.model.model_name.clone(),
|
|
log_model_update: false,
|
|
next_request_id: 0,
|
|
needs_drain: false,
|
|
};
|
|
|
|
if control_protocol_enabled {
|
|
process
|
|
.send_control_request(ControlRequestBody::Initialize)
|
|
.await?;
|
|
}
|
|
|
|
Ok(process)
|
|
}
|
|
|
|
async fn get_or_init_process(
|
|
&self,
|
|
filtered_system: &str,
|
|
) -> Result<&Arc<tokio::sync::Mutex<CliProcess>>, ProviderError> {
|
|
self.cli_process
|
|
.get_or_try_init(|| async {
|
|
Ok(Arc::new(tokio::sync::Mutex::new(
|
|
self.spawn_process(filtered_system).await?,
|
|
)))
|
|
})
|
|
.await
|
|
}
|
|
}
|
|
|
|
async fn exchange_control(
|
|
stdin: &mut (impl AsyncWrite + Unpin),
|
|
reader: &mut (impl AsyncBufRead + Unpin),
|
|
request_id: &str,
|
|
body: ControlRequestBody,
|
|
) -> Result<Option<Value>, ProviderError> {
|
|
let label = body.label();
|
|
let req = ControlRequest {
|
|
msg_type: "control_request",
|
|
request_id: request_id.to_string(),
|
|
request: body,
|
|
};
|
|
let mut req_str = serde_json::to_string(&req).map_err(|e| {
|
|
ProviderError::RequestFailed(format!("Failed to serialize {label} request: {e}"))
|
|
})?;
|
|
req_str.push('\n');
|
|
stdin.write_all(req_str.as_bytes()).await.map_err(|e| {
|
|
ProviderError::RequestFailed(format!("Failed to write {label} request: {e}"))
|
|
})?;
|
|
|
|
let mut line = String::new();
|
|
loop {
|
|
line.clear();
|
|
match reader.read_line(&mut line).await {
|
|
Ok(0) => {
|
|
return Err(ProviderError::RequestFailed(format!(
|
|
"CLI process terminated while waiting for {label} response"
|
|
)));
|
|
}
|
|
Ok(_) => {
|
|
let trimmed = line.trim();
|
|
if trimmed.is_empty() {
|
|
continue;
|
|
}
|
|
if let Ok(msg) = serde_json::from_str::<IncomingControlResponse>(trimmed) {
|
|
match msg.response {
|
|
IncomingControlResponseBody::Success {
|
|
request_id: ref rid,
|
|
response,
|
|
} if rid == request_id => return Ok(response),
|
|
IncomingControlResponseBody::Error {
|
|
request_id: ref rid,
|
|
error,
|
|
} if rid == request_id => {
|
|
return Err(ProviderError::RequestFailed(format!(
|
|
"{label} failed: {error}"
|
|
)));
|
|
}
|
|
_ => continue,
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
return Err(ProviderError::RequestFailed(format!(
|
|
"Failed to read {label} response: {e}"
|
|
)));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn extract_model_aliases(response: Option<&Value>) -> Vec<String> {
|
|
response
|
|
.and_then(|v| v.get("models")?.as_array())
|
|
.map(|arr| {
|
|
arr.iter()
|
|
.filter_map(|m| m.get("value")?.as_str().map(String::from))
|
|
.collect()
|
|
})
|
|
.unwrap_or_default()
|
|
}
|
|
|
|
fn build_stream_json_input(content_blocks: &[Value], session_id: &str) -> String {
|
|
let msg = json!({"type":"user","session_id":session_id,"message":{"role":"user","content":content_blocks}});
|
|
serde_json::to_string(&msg).expect("serializing JSON content blocks cannot fail")
|
|
}
|
|
|
|
fn claude_mcp_config_json(extensions: &[ExtensionConfig]) -> Option<String> {
|
|
let mut mcp_servers = serde_json::Map::new();
|
|
|
|
for extension in extensions {
|
|
match extension {
|
|
ExtensionConfig::StreamableHttp { uri, headers, .. } => {
|
|
let key = extension.key();
|
|
let mut config = serde_json::Map::new();
|
|
config.insert("type".to_string(), json!("http"));
|
|
config.insert("url".to_string(), json!(uri));
|
|
if !headers.is_empty() {
|
|
config.insert("headers".to_string(), json!(headers));
|
|
}
|
|
mcp_servers.insert(key, Value::Object(config));
|
|
}
|
|
ExtensionConfig::Stdio {
|
|
cmd, args, envs, ..
|
|
} => {
|
|
let key = extension.key();
|
|
let mut config = serde_json::Map::new();
|
|
config.insert("type".to_string(), json!("stdio"));
|
|
config.insert("command".to_string(), json!(cmd));
|
|
if !args.is_empty() {
|
|
config.insert("args".to_string(), json!(args));
|
|
}
|
|
let env_map = envs.get_env();
|
|
if !env_map.is_empty() {
|
|
config.insert("env".to_string(), json!(env_map));
|
|
}
|
|
mcp_servers.insert(key, Value::Object(config));
|
|
}
|
|
ExtensionConfig::Sse { name, .. } => {
|
|
tracing::debug!(name, "skipping SSE extension, migrate to streamable_http");
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
if mcp_servers.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
serde_json::to_string(&json!({ "mcpServers": mcp_servers })).ok()
|
|
}
|
|
|
|
/// Write the MCP config JSON to a temp file with restricted permissions
|
|
/// so secrets (headers, env vars) are not leaked via process argv.
|
|
fn write_mcp_config_file(state_dir: &Path, json: &str) -> Result<NamedTempFile, anyhow::Error> {
|
|
let dir = state_dir.join("claude-code");
|
|
std::fs::create_dir_all(&dir)?;
|
|
let prefix = format!("mcp-config-{}_", chrono::Utc::now().format("%Y%m%d"));
|
|
let mut tmp = tempfile::Builder::new()
|
|
.prefix(&prefix)
|
|
.suffix(".json")
|
|
.tempfile_in(&dir)?;
|
|
tmp.write_all(json.as_bytes())?;
|
|
#[cfg(unix)]
|
|
{
|
|
use std::os::unix::fs::PermissionsExt;
|
|
tmp.as_file()
|
|
.set_permissions(std::fs::Permissions::from_mode(0o600))?;
|
|
}
|
|
Ok(tmp)
|
|
}
|
|
|
|
impl ProviderDef for ClaudeCodeProvider {
|
|
type Provider = Self;
|
|
|
|
fn metadata() -> ProviderMetadata {
|
|
ProviderMetadata::new(
|
|
CLAUDE_CODE_PROVIDER_NAME,
|
|
"Claude Code CLI",
|
|
"[Deprecated: use claude-acp instead] Requires claude CLI installed, no MCPs. Use claude-acp for ACP support with extensions.",
|
|
CLAUDE_CODE_DEFAULT_MODEL,
|
|
// Only a few agentic choices; fetched dynamically via fetch_supported_models.
|
|
vec![],
|
|
CLAUDE_CODE_DOC_URL,
|
|
vec![ConfigKey::from_value_type::<ClaudeCodeCommand>(
|
|
true, false, true,
|
|
)],
|
|
)
|
|
}
|
|
|
|
fn from_env(
|
|
model: ModelConfig,
|
|
extensions: Vec<ExtensionConfig>,
|
|
) -> BoxFuture<'static, Result<Self::Provider>> {
|
|
Box::pin(async move {
|
|
let config = crate::config::Config::global();
|
|
let command: String = config.get_claude_code_command().unwrap_or_default().into();
|
|
let resolved_command = SearchPaths::builder().with_npm().resolve(command)?;
|
|
|
|
let mut resolved = Vec::with_capacity(extensions.len());
|
|
for ext in extensions {
|
|
resolved.push(ext.resolve(config).await?);
|
|
}
|
|
|
|
let mcp_config_file = claude_mcp_config_json(&resolved)
|
|
.map(|json| write_mcp_config_file(&Paths::state_dir(), &json))
|
|
.transpose()?;
|
|
|
|
Ok(Self {
|
|
command: resolved_command,
|
|
model,
|
|
name: CLAUDE_CODE_PROVIDER_NAME.to_string(),
|
|
mcp_config_file,
|
|
cli_process: tokio::sync::OnceCell::new(),
|
|
pending_confirmations: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
|
initial_mode: tokio::sync::Mutex::new(None),
|
|
})
|
|
})
|
|
}
|
|
}
|
|
|
|
#[async_trait]
|
|
impl Provider for ClaudeCodeProvider {
|
|
fn get_name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
fn manages_own_context(&self) -> bool {
|
|
true
|
|
}
|
|
|
|
fn get_model_config(&self) -> ModelConfig {
|
|
self.model.clone()
|
|
}
|
|
|
|
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
|
|
// Uses a separate short-lived process because --system-prompt is a CLI-only
|
|
// flag with no NDJSON equivalent. The persistent process needs it at spawn,
|
|
// but it's unavailable during model listing.
|
|
// See: https://code.claude.com/docs/en/cli-reference#system-prompt-flags
|
|
let mut cmd = self.build_stream_json_command();
|
|
let mut child = cmd.spawn().map_err(|e| {
|
|
ProviderError::RequestFailed(format!("Failed to spawn CLI for model listing: {e}"))
|
|
})?;
|
|
|
|
let mut stdin = child
|
|
.stdin
|
|
.take()
|
|
.ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdin".to_string()))?;
|
|
let stdout = child
|
|
.stdout
|
|
.take()
|
|
.ok_or_else(|| ProviderError::RequestFailed("Failed to capture stdout".to_string()))?;
|
|
|
|
let mut reader = BufReader::new(stdout);
|
|
let response = exchange_control(
|
|
&mut stdin,
|
|
&mut reader,
|
|
"model_list",
|
|
ControlRequestBody::Initialize,
|
|
)
|
|
.await;
|
|
let _ = child.kill().await;
|
|
Ok(extract_model_aliases(response.ok().flatten().as_ref()))
|
|
}
|
|
|
|
async fn update_mode(&self, _session_id: &str, mode: GooseMode) -> Result<(), ProviderError> {
|
|
// Mode is baked into the subprocess at spawn; claude-acp replaces
|
|
// this provider (#7801).
|
|
let mut guard = self.initial_mode.lock().await;
|
|
let current = *guard.get_or_insert(mode);
|
|
if current != mode {
|
|
return Err(ProviderError::RequestFailed(format!(
|
|
"Mode change not supported: session is {current}, requested {mode}",
|
|
)));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn permission_routing(&self) -> PermissionRouting {
|
|
PermissionRouting::ActionRequired
|
|
}
|
|
|
|
async fn handle_permission_confirmation(
|
|
&self,
|
|
request_id: &str,
|
|
confirmation: &PermissionConfirmation,
|
|
) -> bool {
|
|
let mut pending = self.pending_confirmations.lock().await;
|
|
if let Some(tx) = pending.remove(request_id) {
|
|
let _ = tx.send(confirmation.clone());
|
|
return true;
|
|
}
|
|
false
|
|
}
|
|
|
|
async fn stream(
|
|
&self,
|
|
model_config: &ModelConfig,
|
|
session_id: &str,
|
|
system: &str,
|
|
messages: &[Message],
|
|
_tools: &[Tool],
|
|
) -> Result<MessageStream, ProviderError> {
|
|
if super::cli_common::is_session_description_request(system) {
|
|
let (message, usage) = super::cli_common::generate_simple_session_description(
|
|
&model_config.model_name,
|
|
messages,
|
|
)?;
|
|
return Ok(stream_from_single_message(message, usage));
|
|
}
|
|
|
|
let filtered_system = filter_extensions_from_system_prompt(system);
|
|
let process_arc = Arc::clone(self.get_or_init_process(&filtered_system).await?);
|
|
|
|
// Prepare the payload outside the lock — these don't need the process.
|
|
let blocks = self.last_user_content_blocks(messages);
|
|
let ndjson_line = build_stream_json_input(&blocks, session_id);
|
|
let model_name = model_config.model_name.clone();
|
|
let message_id = uuid::Uuid::new_v4().to_string();
|
|
let pending_confirmations = Arc::clone(&self.pending_confirmations);
|
|
|
|
Ok(Box::pin(try_stream! {
|
|
// Single lock acquisition covers write-to-stdin and read-from-stdout,
|
|
// eliminating the race window between the two.
|
|
let mut process = process_arc.lock_owned().await;
|
|
|
|
// Clean up pending permissions from a cancelled stream
|
|
{
|
|
let mut pending = pending_confirmations.lock().await;
|
|
for (req_id, tx) in pending.drain() {
|
|
drop(tx);
|
|
let resp = ControlResponse::success(
|
|
req_id,
|
|
PermissionResponse::Deny { message: "Stream cancelled".to_string() },
|
|
);
|
|
let mut s = serde_json::to_string(&resp).map_err(|e| {
|
|
ProviderError::RequestFailed(format!("Failed to serialize cleanup deny response: {e}"))
|
|
})?;
|
|
s.push('\n');
|
|
let _ = process.stdin.write_all(s.as_bytes()).await;
|
|
}
|
|
}
|
|
|
|
process.drain_pending_response().await;
|
|
process.send_set_model(&model_name).await?;
|
|
|
|
process
|
|
.stdin
|
|
.write_all(ndjson_line.as_bytes())
|
|
.await
|
|
.map_err(|e| {
|
|
ProviderError::RequestFailed(format!("Failed to write to stdin: {}", e))
|
|
})?;
|
|
process.stdin.write_all(b"\n").await.map_err(|e| {
|
|
ProviderError::RequestFailed(format!("Failed to write newline to stdin: {}", e))
|
|
})?;
|
|
|
|
process.needs_drain = true;
|
|
let mut line = String::new();
|
|
let mut accumulated_usage = Usage::default();
|
|
let mut stream_error: Option<ProviderError> = None;
|
|
let stream_timestamp = chrono::Utc::now().timestamp();
|
|
|
|
loop {
|
|
line.clear();
|
|
match process.reader.read_line(&mut line).await {
|
|
Ok(0) => {
|
|
process.needs_drain = false;
|
|
stream_error = Some(ProviderError::RequestFailed(
|
|
"Claude CLI process terminated unexpectedly".to_string(),
|
|
));
|
|
break;
|
|
}
|
|
Ok(_) => {
|
|
let trimmed = line.trim();
|
|
if trimmed.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
if let Ok(parsed) = serde_json::from_str::<Value>(trimmed) {
|
|
match parsed.get("type").and_then(|t| t.as_str()) {
|
|
Some("stream_event") => {
|
|
if let Some(event) = parsed.get("event") {
|
|
match event.get("type").and_then(|t| t.as_str()) {
|
|
Some("content_block_delta") => {
|
|
if let Some(text) = event
|
|
.get("delta")
|
|
.filter(|d| {
|
|
d.get("type").and_then(|t| t.as_str())
|
|
== Some("text_delta")
|
|
})
|
|
.and_then(|d| d.get("text"))
|
|
.and_then(|t| t.as_str())
|
|
{
|
|
if !text.is_empty() {
|
|
let mut partial_message = Message::new(
|
|
Role::Assistant,
|
|
stream_timestamp,
|
|
vec![MessageContent::text(text)],
|
|
);
|
|
partial_message.id =
|
|
Some(message_id.clone());
|
|
yield (Some(partial_message), None);
|
|
}
|
|
}
|
|
}
|
|
Some("message_start") => {
|
|
if let Some(usage_info) = event
|
|
.get("message")
|
|
.and_then(|m| m.get("usage"))
|
|
{
|
|
let new = extract_usage_tokens(usage_info);
|
|
if let Some(i) = new.input_tokens {
|
|
accumulated_usage.input_tokens = Some(i);
|
|
}
|
|
}
|
|
}
|
|
Some("message_delta") => {
|
|
if let Some(usage_info) = event.get("usage") {
|
|
let new = extract_usage_tokens(usage_info);
|
|
if let Some(o) = new.output_tokens {
|
|
accumulated_usage.output_tokens = Some(o);
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
Some("result") => {
|
|
process.needs_drain = false;
|
|
if let Some(usage_info) = parsed.get("usage") {
|
|
let new = extract_usage_tokens(usage_info);
|
|
accumulated_usage = Usage::new(
|
|
new.input_tokens.or(accumulated_usage.input_tokens),
|
|
new.output_tokens.or(accumulated_usage.output_tokens),
|
|
None,
|
|
);
|
|
}
|
|
break;
|
|
}
|
|
Some("error") => {
|
|
process.needs_drain = false;
|
|
stream_error = Some(error_from_event("Claude CLI", &parsed));
|
|
break;
|
|
}
|
|
Some("control_request") => {
|
|
if let Ok(IncomingControlRequest {
|
|
request_id,
|
|
request: IncomingRequestBody::CanUseTool { tool_name, input, tool_use_id },
|
|
}) = serde_json::from_str::<IncomingControlRequest>(trimmed) {
|
|
tracing::debug!(raw = %parsed, "can_use_tool control_request received");
|
|
|
|
let (tx, rx) = oneshot::channel();
|
|
pending_confirmations.lock().await.insert(request_id.clone(), tx);
|
|
|
|
let action_msg = Message::assistant().with_action_required(
|
|
request_id.clone(), tool_name, input.clone(), None,
|
|
);
|
|
yield (Some(action_msg), None);
|
|
|
|
let confirmation = rx.await.unwrap_or(PermissionConfirmation {
|
|
principal_type: PrincipalType::Tool,
|
|
permission: Permission::Cancel,
|
|
});
|
|
pending_confirmations.lock().await.remove(&request_id);
|
|
|
|
let perm_resp = match confirmation.permission {
|
|
Permission::AlwaysAllow | Permission::AllowOnce => {
|
|
PermissionResponse::Allow {
|
|
updated_input: input,
|
|
tool_use_id,
|
|
}
|
|
}
|
|
_ => PermissionResponse::Deny {
|
|
message: "User denied the tool call".to_string(),
|
|
},
|
|
};
|
|
let resp = ControlResponse::success(request_id, perm_resp);
|
|
let mut resp_str = serde_json::to_string(&resp).map_err(|e| {
|
|
ProviderError::RequestFailed(format!("Failed to serialize permission response: {e}"))
|
|
})?;
|
|
tracing::debug!(json = %resp_str, "can_use_tool control_response sent");
|
|
resp_str.push('\n');
|
|
process.stdin.write_all(resp_str.as_bytes()).await.map_err(|e| {
|
|
ProviderError::RequestFailed(format!("Failed to write permission response: {e}"))
|
|
})?;
|
|
}
|
|
}
|
|
Some("system") if process.log_model_update => {
|
|
if let Some(resolved) = parsed.get("model").and_then(|m| m.as_str()) {
|
|
tracing::debug!(
|
|
from = %process.current_model,
|
|
to = %resolved,
|
|
"set_model resolved"
|
|
);
|
|
}
|
|
process.log_model_update = false;
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
process.needs_drain = false;
|
|
stream_error = Some(ProviderError::RequestFailed(format!(
|
|
"Failed to read streaming output: {e}"
|
|
)));
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if let Some(err) = stream_error {
|
|
Err(err)?;
|
|
}
|
|
|
|
let provider_usage = ProviderUsage::new(model_name, accumulated_usage);
|
|
yield (None, Some(provider_usage));
|
|
}))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::agents::extension::Envs;
|
|
use chrono::Utc;
|
|
use goose_test_support::session::TEST_SESSION_ID;
|
|
use serde_json::json;
|
|
use std::collections::HashMap;
|
|
use std::fs;
|
|
use tempfile::tempdir;
|
|
use test_case::test_case;
|
|
|
|
#[test_case(
|
|
json!({"input_tokens": 100, "output_tokens": 50}),
|
|
Some(100), Some(50)
|
|
; "both_tokens"
|
|
)]
|
|
#[test_case(json!({"input_tokens": 100}), Some(100), None ; "input_only")]
|
|
#[test_case(json!({}), None, None ; "empty_usage")]
|
|
fn test_extract_usage_tokens(
|
|
usage_json: Value,
|
|
expected_input: Option<i32>,
|
|
expected_output: Option<i32>,
|
|
) {
|
|
let usage = extract_usage_tokens(&usage_json);
|
|
assert_eq!(usage.input_tokens, expected_input);
|
|
assert_eq!(usage.output_tokens, expected_output);
|
|
}
|
|
|
|
#[test_case(
|
|
r#"{"type":"error","error":"context window exceeded"}"#,
|
|
true
|
|
; "context_exceeded"
|
|
)]
|
|
#[test_case(
|
|
r#"{"type":"error","error":"Model not supported"}"#,
|
|
false
|
|
; "generic_error_from_event"
|
|
)]
|
|
#[test_case(r#"{"type":"error"}"#, false ; "missing_error_field")]
|
|
fn test_error_from_event(line: &str, is_context_exceeded: bool) {
|
|
let parsed: Value = serde_json::from_str(line).unwrap();
|
|
let err = error_from_event("Claude CLI", &parsed);
|
|
if is_context_exceeded {
|
|
assert!(matches!(err, ProviderError::ContextLengthExceeded(_)));
|
|
} else {
|
|
assert!(matches!(err, ProviderError::RequestFailed(_)));
|
|
}
|
|
}
|
|
|
|
/// (role, text, optional (image_data, mime_type))
|
|
type MsgSpec<'a> = (&'a str, &'a str, Option<(&'a str, &'a str)>);
|
|
|
|
fn build_messages(specs: &[MsgSpec]) -> Vec<Message> {
|
|
specs
|
|
.iter()
|
|
.map(|(role, text, image)| {
|
|
let role = if *role == "user" {
|
|
Role::User
|
|
} else {
|
|
Role::Assistant
|
|
};
|
|
let mut msg = Message::new(role, 0, vec![]);
|
|
if !text.is_empty() {
|
|
msg = Message::new(msg.role.clone(), 0, vec![MessageContent::text(*text)]);
|
|
}
|
|
if let Some((data, mime)) = image {
|
|
msg.content.push(MessageContent::image(*data, *mime));
|
|
}
|
|
msg
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
#[test_case(
|
|
build_messages(&[]),
|
|
&[]
|
|
; "empty"
|
|
)]
|
|
#[test_case(
|
|
build_messages(&[("user", "Hello", None)]),
|
|
&[json!({"type":"text","text":"Human: Hello"})]
|
|
; "single_user"
|
|
)]
|
|
#[test_case(
|
|
build_messages(&[("user", "Hello", None), ("assistant", "Hi there!", None)]),
|
|
&[json!({"type":"text","text":"Human: Hello"})]
|
|
; "picks_last_user_ignores_assistant"
|
|
)]
|
|
#[test_case(
|
|
build_messages(&[("user", "First", None), ("assistant", "Reply", None), ("user", "Second", None)]),
|
|
&[json!({"type":"text","text":"Human: Second"})]
|
|
; "multi_turn_picks_last_user"
|
|
)]
|
|
#[test_case(
|
|
build_messages(&[("user", "Describe this", Some(("base64data", "image/png")))]),
|
|
&[json!({"type":"text","text":"Human: Describe this"}),
|
|
json!({"type":"image","source":{"type":"base64","media_type":"image/png","data":"base64data"}})]
|
|
; "user_with_image"
|
|
)]
|
|
#[test_case(
|
|
build_messages(&[("user", "", Some(("iVBORw0KGgo", "image/png")))]),
|
|
&[json!({"type":"image","source":{"type":"base64","media_type":"image/png","data":"iVBORw0KGgo"}})]
|
|
; "image_only"
|
|
)]
|
|
#[test_case(
|
|
vec![Message::new(Role::Assistant, 0, vec![
|
|
MessageContent::tool_request("call_123", Ok(rmcp::model::CallToolRequestParams::new("developer__shell").with_arguments(serde_json::from_value(json!({"cmd": "ls"})).unwrap())))
|
|
])],
|
|
&[json!({"type":"text","text":"Assistant: [tool_use: developer__shell id=call_123]"})]
|
|
; "tool_request_no_user_fallback"
|
|
)]
|
|
#[test_case(
|
|
vec![Message::new(Role::User, 0, vec![
|
|
MessageContent::tool_response("call_123", Ok(rmcp::model::CallToolResult::success(vec![rmcp::model::Content::text("file1.txt\nfile2.txt")])))
|
|
])],
|
|
&[json!({"type":"text","text":"Human: [tool_result id=call_123] file1.txt\nfile2.txt"})]
|
|
; "tool_response"
|
|
)]
|
|
#[test_case(
|
|
vec![Message::new(Role::User, 0, vec![MessageContent::text("hidden input")]).user_only()],
|
|
&[json!({"type":"text","text":"Human: hidden input"})]
|
|
; "user_only_message_not_dropped"
|
|
)]
|
|
fn test_last_user_content_blocks(messages: Vec<Message>, expected: &[Value]) {
|
|
let provider = make_provider();
|
|
let blocks = provider.last_user_content_blocks(&messages);
|
|
assert_eq!(blocks, expected);
|
|
}
|
|
|
|
#[test_case(
|
|
&[json!({"type":"text","text":"Hello"})],
|
|
json!({"type":"user","session_id":TEST_SESSION_ID,"message":{"role":"user","content":[{"type":"text","text":"Hello"}]}})
|
|
; "text_block"
|
|
)]
|
|
#[test_case(
|
|
&[json!({"type":"text","text":"Look"}), json!({"type":"image","source":{"type":"base64","media_type":"image/png","data":"abc"}})],
|
|
json!({"type":"user","session_id":TEST_SESSION_ID,"message":{"role":"user","content":[{"type":"text","text":"Look"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"abc"}}]}})
|
|
; "text_and_image_blocks"
|
|
)]
|
|
fn test_build_stream_json_input(blocks: &[Value], expected: Value) {
|
|
let line = build_stream_json_input(blocks, TEST_SESSION_ID);
|
|
let parsed: Value = serde_json::from_str(&line).unwrap();
|
|
assert_eq!(parsed, expected);
|
|
}
|
|
|
|
#[test_case(
|
|
Some(json!({"models":[{"value":"default","displayName":"Default"},{"value":"sonnet","displayName":"Sonnet"},{"value":"haiku","displayName":"Haiku"}]})),
|
|
vec!["default".into(), "sonnet".into(), "haiku".into()]
|
|
; "success"
|
|
)]
|
|
#[test_case(
|
|
Some(json!({"models":[{"value":"default","displayName":"Default"},{"value":null,"displayName":"Bad"}]})),
|
|
vec!["default".into()]
|
|
; "filters_null_values"
|
|
)]
|
|
#[test_case(
|
|
None,
|
|
vec![]
|
|
; "none_input"
|
|
)]
|
|
#[test_case(
|
|
Some(json!({"other":"data"})),
|
|
vec![]
|
|
; "no_models_key"
|
|
)]
|
|
fn test_extract_model_aliases(response: Option<Value>, expected: Vec<String>) {
|
|
assert_eq!(extract_model_aliases(response.as_ref()), expected);
|
|
}
|
|
|
|
#[test_case(
|
|
vec![],
|
|
None
|
|
; "empty_extensions_returns_none"
|
|
)]
|
|
#[test_case(
|
|
vec![ExtensionConfig::Sse {
|
|
name: "legacy".into(),
|
|
description: String::new(),
|
|
uri: Some("http://localhost/sse".into()),
|
|
}],
|
|
None
|
|
; "sse_only_returns_none"
|
|
)]
|
|
#[test_case(
|
|
vec![ExtensionConfig::Stdio {
|
|
name: "lookup".into(),
|
|
description: String::new(),
|
|
cmd: "node".into(),
|
|
args: vec!["server.js".into()],
|
|
envs: Envs::new([("API_KEY".into(), "secret".into())].into()),
|
|
env_keys: vec![],
|
|
timeout: None,
|
|
bundled: Some(false),
|
|
available_tools: vec![],
|
|
}],
|
|
Some(json!({ "mcpServers": {
|
|
"lookup": {
|
|
"type": "stdio",
|
|
"command": "node",
|
|
"args": ["server.js"],
|
|
"env": { "API_KEY": "secret" }
|
|
}
|
|
}}))
|
|
; "stdio_converts_to_mcp_config_json"
|
|
)]
|
|
#[test_case(
|
|
vec![ExtensionConfig::StreamableHttp {
|
|
name: "lookup".into(),
|
|
description: String::new(),
|
|
uri: "http://localhost/mcp".into(),
|
|
envs: Envs::default(),
|
|
env_keys: vec![],
|
|
headers: HashMap::from([("Authorization".into(), "Bearer token".into())]),
|
|
timeout: None,
|
|
bundled: Some(false),
|
|
available_tools: vec![],
|
|
}],
|
|
Some(json!({ "mcpServers": {
|
|
"lookup": {
|
|
"type": "http",
|
|
"url": "http://localhost/mcp",
|
|
"headers": { "Authorization": "Bearer token" }
|
|
}
|
|
}}))
|
|
; "streamable_http_converts_to_mcp_config_json"
|
|
)]
|
|
#[test_case(
|
|
vec![ExtensionConfig::StreamableHttp {
|
|
name: "mcp_kiwi_com".into(),
|
|
description: String::new(),
|
|
uri: "https://mcp.kiwi.com".into(),
|
|
envs: Envs::default(),
|
|
env_keys: vec![],
|
|
headers: HashMap::new(),
|
|
timeout: None,
|
|
bundled: None,
|
|
available_tools: vec![],
|
|
}],
|
|
Some(json!({ "mcpServers": {
|
|
"mcp_kiwi_com": {
|
|
"type": "http",
|
|
"url": "https://mcp.kiwi.com"
|
|
}
|
|
}}))
|
|
; "resolved_name_used_as_key"
|
|
)]
|
|
fn test_claude_mcp_config_json(extensions: Vec<ExtensionConfig>, expected: Option<Value>) {
|
|
let result = claude_mcp_config_json(&extensions)
|
|
.map(|json| serde_json::from_str::<Value>(&json).unwrap());
|
|
assert_eq!(result, expected);
|
|
}
|
|
|
|
#[test]
|
|
fn test_write_mcp_config_file() {
|
|
let state_dir = tempdir().unwrap();
|
|
let json = r#"{"mcpServers":{}}"#;
|
|
|
|
let tmp = write_mcp_config_file(state_dir.path(), json).unwrap();
|
|
|
|
assert_eq!(fs::read_to_string(tmp.path()).unwrap(), json);
|
|
|
|
let norm_path = tmp.path().to_string_lossy().replace('\\', "/");
|
|
let expected_prefix = format!("claude-code/mcp-config-{}_", Utc::now().format("%Y%m%d"));
|
|
assert!(norm_path.contains(&expected_prefix));
|
|
assert!(norm_path.ends_with(".json"));
|
|
}
|
|
|
|
#[test]
|
|
fn test_write_mcp_config_file_invalid_state_dir() {
|
|
assert!(write_mcp_config_file(Path::new("/dev/null"), "{}").is_err());
|
|
}
|
|
|
|
fn make_provider() -> ClaudeCodeProvider {
|
|
ClaudeCodeProvider {
|
|
command: PathBuf::from("claude"),
|
|
model: ModelConfig::new(CLAUDE_CODE_DEFAULT_MODEL)
|
|
.unwrap()
|
|
.with_canonical_limits(CLAUDE_CODE_PROVIDER_NAME),
|
|
name: "claude-code".to_string(),
|
|
mcp_config_file: None,
|
|
cli_process: tokio::sync::OnceCell::new(),
|
|
pending_confirmations: Arc::new(tokio::sync::Mutex::new(HashMap::new())),
|
|
initial_mode: tokio::sync::Mutex::new(None),
|
|
}
|
|
}
|
|
|
|
fn make_test_process(canned_stdout: &str) -> (CliProcess, tokio::io::DuplexStream) {
|
|
let child = tokio::process::Command::new("true")
|
|
.spawn()
|
|
.expect("failed to spawn `true`");
|
|
let (stdin_writer, stdin_reader) = tokio::io::duplex(1024);
|
|
let process = CliProcess {
|
|
child,
|
|
stdin: Box::new(stdin_writer),
|
|
reader: BufReader::new(Box::new(std::io::Cursor::new(
|
|
canned_stdout.as_bytes().to_vec(),
|
|
))),
|
|
stderr_handle: tokio::spawn(async { String::new() }),
|
|
current_model: String::new(),
|
|
log_model_update: false,
|
|
next_request_id: 0,
|
|
needs_drain: false,
|
|
};
|
|
(process, stdin_reader)
|
|
}
|
|
|
|
async fn stream_with_canned_stdout(
|
|
canned_lines: &[&str],
|
|
) -> (ClaudeCodeProvider, MessageStream, tokio::io::DuplexStream) {
|
|
let canned_stdout = canned_lines.join("\n");
|
|
let (process, stdin_reader) = make_test_process(&canned_stdout);
|
|
let provider = make_provider();
|
|
let process_arc = Arc::new(tokio::sync::Mutex::new(process));
|
|
provider.cli_process.set(process_arc).unwrap();
|
|
|
|
let messages = vec![Message::user().with_text("test")];
|
|
let stream = provider
|
|
.stream(&provider.model, "test-session", "", &messages, &[])
|
|
.await
|
|
.unwrap();
|
|
(provider, stream, stdin_reader)
|
|
}
|
|
|
|
async fn capture_stdin(
|
|
provider: &ClaudeCodeProvider,
|
|
mut reader: tokio::io::DuplexStream,
|
|
) -> String {
|
|
use tokio::io::AsyncReadExt;
|
|
provider.cli_process.get().unwrap().lock().await.stdin = Box::new(tokio::io::sink());
|
|
let mut buf = Vec::new();
|
|
reader.read_to_end(&mut buf).await.unwrap();
|
|
String::from_utf8(buf).unwrap()
|
|
}
|
|
|
|
fn extract_permission_response(stdin_str: &str, request_id: &str) -> Value {
|
|
let line = stdin_str
|
|
.lines()
|
|
.find(|l| l.contains(request_id) && l.contains("control_response"))
|
|
.unwrap();
|
|
let json: Value = serde_json::from_str(line).unwrap();
|
|
json.pointer("/response/response").unwrap().clone()
|
|
}
|
|
|
|
#[test_case(
|
|
&[r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#],
|
|
Some("default"), "sonnet",
|
|
Ok(()),
|
|
"{\"type\":\"control_request\",\"request_id\":\"req_0\",\"request\":{\"subtype\":\"set_model\",\"model\":\"sonnet\"}}\n"
|
|
; "default_to_sonnet"
|
|
)]
|
|
#[test_case(
|
|
&[r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#],
|
|
Some("sonnet"), "default",
|
|
Ok(()),
|
|
"{\"type\":\"control_request\",\"request_id\":\"req_0\",\"request\":{\"subtype\":\"set_model\",\"model\":\"default\"}}\n"
|
|
; "sonnet_to_default"
|
|
)]
|
|
#[test_case(
|
|
&[r#"{"type":"control_response","response":{"subtype":"error","request_id":"req_0","error":"bad model"}}"#],
|
|
None, "bad",
|
|
Err(ProviderError::RequestFailed("set_model failed: bad model".into())),
|
|
"{\"type\":\"control_request\",\"request_id\":\"req_0\",\"request\":{\"subtype\":\"set_model\",\"model\":\"bad\"}}\n"
|
|
; "failure"
|
|
)]
|
|
#[test_case(
|
|
&[],
|
|
Some("sonnet"), "sonnet",
|
|
Ok(()), ""
|
|
; "skip_when_same_model"
|
|
)]
|
|
#[test_case(
|
|
&[],
|
|
None, "sonnet",
|
|
Err(ProviderError::RequestFailed("CLI process terminated while waiting for set_model response".into())),
|
|
"{\"type\":\"control_request\",\"request_id\":\"req_0\",\"request\":{\"subtype\":\"set_model\",\"model\":\"sonnet\"}}\n"
|
|
; "eof"
|
|
)]
|
|
#[tokio::test]
|
|
async fn test_send_set_model(
|
|
lines: &[&str],
|
|
initial_model: Option<&str>,
|
|
target_model: &str,
|
|
expected: Result<(), ProviderError>,
|
|
expected_stdin: &str,
|
|
) {
|
|
use tokio::io::AsyncReadExt;
|
|
|
|
let stdout = lines.join("\n");
|
|
let (mut process, mut stdin_reader) = make_test_process(&stdout);
|
|
if let Some(m) = initial_model {
|
|
process.current_model = m.to_string();
|
|
}
|
|
|
|
let result = process.send_set_model(target_model).await;
|
|
process.stdin = Box::new(tokio::io::sink());
|
|
let mut stdin_bytes = Vec::new();
|
|
stdin_reader.read_to_end(&mut stdin_bytes).await.unwrap();
|
|
|
|
assert_eq!(result, expected);
|
|
if expected.is_ok() {
|
|
assert_eq!(process.current_model, target_model);
|
|
}
|
|
assert_eq!(String::from_utf8(stdin_bytes).unwrap(), expected_stdin);
|
|
}
|
|
|
|
#[test_case(
|
|
Permission::AllowOnce,
|
|
json!({"behavior":"allow","updatedInput":{"path":"foo.txt","content":"hello"},"toolUseID":"tu_1"})
|
|
; "allow"
|
|
)]
|
|
#[test_case(
|
|
Permission::DenyOnce,
|
|
json!({"behavior":"deny","message":"User denied the tool call"})
|
|
; "deny"
|
|
)]
|
|
#[tokio::test]
|
|
async fn test_can_use_tool(permission: Permission, expected_response: Value) {
|
|
use futures::StreamExt;
|
|
|
|
let (provider, mut stream, stdin_reader) = stream_with_canned_stdout(&[
|
|
r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#,
|
|
r#"{"type":"control_request","request_id":"perm_1","request":{"subtype":"can_use_tool","tool_name":"Write","input":{"path":"foo.txt","content":"hello"},"tool_use_id":"tu_1"}}"#,
|
|
r#"{"type":"result","result":"Done","usage":{"input_tokens":10,"output_tokens":5}}"#,
|
|
]).await;
|
|
|
|
let (first_msg, _) = stream.next().await.unwrap().unwrap();
|
|
let first_msg = first_msg.unwrap();
|
|
let ar = first_msg
|
|
.content
|
|
.iter()
|
|
.find_map(|c| c.as_action_required())
|
|
.unwrap();
|
|
match &ar.data {
|
|
crate::conversation::message::ActionRequiredData::ToolConfirmation {
|
|
id,
|
|
tool_name,
|
|
..
|
|
} => {
|
|
assert_eq!(id, "perm_1");
|
|
assert_eq!(tool_name, "Write");
|
|
}
|
|
_ => panic!("expected ToolConfirmation"),
|
|
}
|
|
|
|
let handled = provider
|
|
.handle_permission_confirmation(
|
|
"perm_1",
|
|
&PermissionConfirmation {
|
|
principal_type: PrincipalType::Tool,
|
|
permission: permission.clone(),
|
|
},
|
|
)
|
|
.await;
|
|
assert!(handled);
|
|
assert!(provider.pending_confirmations.lock().await.is_empty());
|
|
|
|
while let Some(item) = stream.next().await {
|
|
item.unwrap();
|
|
}
|
|
drop(stream);
|
|
|
|
let stdin_str = capture_stdin(&provider, stdin_reader).await;
|
|
let response_data = extract_permission_response(&stdin_str, "perm_1");
|
|
assert_eq!(response_data, expected_response);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_can_use_tool_cancel_on_drop() {
|
|
use futures::StreamExt;
|
|
|
|
let (provider, mut stream, stdin_reader) = stream_with_canned_stdout(&[
|
|
r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#,
|
|
r#"{"type":"control_request","request_id":"perm_1","request":{"subtype":"can_use_tool","tool_name":"Write","input":{"path":"foo.txt"},"tool_use_id":"tu_1"}}"#,
|
|
r#"{"type":"result","result":"Done","usage":{"input_tokens":10,"output_tokens":5}}"#,
|
|
]).await;
|
|
|
|
let pending = Arc::clone(&provider.pending_confirmations);
|
|
|
|
let (first_msg, _) = stream.next().await.unwrap().unwrap();
|
|
assert!(first_msg
|
|
.unwrap()
|
|
.content
|
|
.iter()
|
|
.any(|c| c.as_action_required().is_some()));
|
|
|
|
let tx = pending.lock().await.remove("perm_1").unwrap();
|
|
drop(tx);
|
|
|
|
while let Some(item) = stream.next().await {
|
|
item.unwrap();
|
|
}
|
|
drop(stream);
|
|
|
|
let stdin_str = capture_stdin(&provider, stdin_reader).await;
|
|
let response_data = extract_permission_response(&stdin_str, "perm_1");
|
|
assert_eq!(
|
|
response_data,
|
|
json!({"behavior":"deny","message":"User denied the tool call"})
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_pending_permissions_cleaned_on_new_stream() {
|
|
use futures::StreamExt;
|
|
|
|
let canned_stdout = [
|
|
r#"{"type":"control_response","response":{"subtype":"success","request_id":"req_0"}}"#,
|
|
r#"{"type":"result","result":"Done","usage":{"input_tokens":10,"output_tokens":5}}"#,
|
|
]
|
|
.join("\n");
|
|
|
|
let (process, stdin_reader) = make_test_process(&canned_stdout);
|
|
let provider = make_provider();
|
|
let process_arc = Arc::new(tokio::sync::Mutex::new(process));
|
|
provider.cli_process.set(process_arc).unwrap();
|
|
|
|
let (tx, _rx) = oneshot::channel();
|
|
provider
|
|
.pending_confirmations
|
|
.lock()
|
|
.await
|
|
.insert("stale_1".to_string(), tx);
|
|
|
|
let messages = vec![Message::user().with_text("test")];
|
|
let mut stream = provider
|
|
.stream(&provider.model, "test-session", "", &messages, &[])
|
|
.await
|
|
.unwrap();
|
|
|
|
while let Some(item) = stream.next().await {
|
|
item.unwrap();
|
|
}
|
|
drop(stream);
|
|
|
|
assert!(provider.pending_confirmations.lock().await.is_empty());
|
|
|
|
let stdin_str = capture_stdin(&provider, stdin_reader).await;
|
|
let response_data = extract_permission_response(&stdin_str, "stale_1");
|
|
assert_eq!(response_data["behavior"], "deny");
|
|
}
|
|
}
|