Restore codebase deleted by bad auto-commit e4227cf
Commit e4227cf (a story creation auto-commit) erroneously deleted 175
files from master's tree, likely due to a race condition between
concurrent git operations. This commit re-adds all files from the
working directory.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
dave
@@ -0,0 +1,73 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use portable_pty::ChildKiller;
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use crate::agent_log::AgentLogWriter;
|
||||
use crate::io::watcher::WatcherEvent;
|
||||
|
||||
use super::{AgentEvent, AgentRuntime, RuntimeContext, RuntimeResult, RuntimeStatus};
|
||||
|
||||
/// Agent runtime that spawns the `claude` CLI in a PTY and streams JSON events.
|
||||
///
|
||||
/// This is the default runtime (`runtime = "claude-code"` in project.toml).
|
||||
/// It wraps the existing PTY-based execution logic, preserving all streaming,
|
||||
/// token tracking, and inactivity timeout behaviour.
|
||||
pub struct ClaudeCodeRuntime {
|
||||
child_killers: Arc<Mutex<HashMap<String, Box<dyn ChildKiller + Send + Sync>>>>,
|
||||
watcher_tx: broadcast::Sender<WatcherEvent>,
|
||||
}
|
||||
|
||||
impl ClaudeCodeRuntime {
|
||||
pub fn new(
|
||||
child_killers: Arc<Mutex<HashMap<String, Box<dyn ChildKiller + Send + Sync>>>>,
|
||||
watcher_tx: broadcast::Sender<WatcherEvent>,
|
||||
) -> Self {
|
||||
Self {
|
||||
child_killers,
|
||||
watcher_tx,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AgentRuntime for ClaudeCodeRuntime {
|
||||
async fn start(
|
||||
&self,
|
||||
ctx: RuntimeContext,
|
||||
tx: broadcast::Sender<AgentEvent>,
|
||||
event_log: Arc<Mutex<Vec<AgentEvent>>>,
|
||||
log_writer: Option<Arc<Mutex<AgentLogWriter>>>,
|
||||
) -> Result<RuntimeResult, String> {
|
||||
let pty_result = super::super::pty::run_agent_pty_streaming(
|
||||
&ctx.story_id,
|
||||
&ctx.agent_name,
|
||||
&ctx.command,
|
||||
&ctx.args,
|
||||
&ctx.prompt,
|
||||
&ctx.cwd,
|
||||
&tx,
|
||||
&event_log,
|
||||
log_writer,
|
||||
ctx.inactivity_timeout_secs,
|
||||
Arc::clone(&self.child_killers),
|
||||
self.watcher_tx.clone(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(RuntimeResult {
|
||||
session_id: pty_result.session_id,
|
||||
token_usage: pty_result.token_usage,
|
||||
})
|
||||
}
|
||||
|
||||
fn stop(&self) {
|
||||
// Stopping is handled externally by the pool via kill_child_for_key().
|
||||
// The ChildKillerGuard in pty.rs deregisters automatically on process exit.
|
||||
}
|
||||
|
||||
fn get_status(&self) -> RuntimeStatus {
|
||||
// Lifecycle status is tracked by the pool; the runtime itself is stateless.
|
||||
RuntimeStatus::Idle
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,809 @@
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use reqwest::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use crate::agent_log::AgentLogWriter;
|
||||
use crate::slog;
|
||||
|
||||
use super::super::{AgentEvent, TokenUsage};
|
||||
use super::{AgentRuntime, RuntimeContext, RuntimeResult, RuntimeStatus};
|
||||
|
||||
// ── Public runtime struct ────────────────────────────────────────────
|
||||
|
||||
/// Agent runtime that drives a Gemini model through the Google AI
|
||||
/// `generateContent` REST API.
|
||||
///
|
||||
/// The runtime:
|
||||
/// 1. Fetches MCP tool definitions from storkit's MCP server.
|
||||
/// 2. Converts them to Gemini function-calling format.
|
||||
/// 3. Sends the agent prompt + tools to the Gemini API.
|
||||
/// 4. Executes any requested function calls via MCP `tools/call`.
|
||||
/// 5. Loops until the model produces a text-only response or an error.
|
||||
/// 6. Tracks token usage from the API response metadata.
|
||||
pub struct GeminiRuntime {
|
||||
/// Whether a stop has been requested.
|
||||
cancelled: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl GeminiRuntime {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
cancelled: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AgentRuntime for GeminiRuntime {
|
||||
async fn start(
|
||||
&self,
|
||||
ctx: RuntimeContext,
|
||||
tx: broadcast::Sender<AgentEvent>,
|
||||
event_log: Arc<Mutex<Vec<AgentEvent>>>,
|
||||
log_writer: Option<Arc<Mutex<AgentLogWriter>>>,
|
||||
) -> Result<RuntimeResult, String> {
|
||||
let api_key = std::env::var("GOOGLE_AI_API_KEY").map_err(|_| {
|
||||
"GOOGLE_AI_API_KEY environment variable is not set. \
|
||||
Set it to your Google AI API key to use the Gemini runtime."
|
||||
.to_string()
|
||||
})?;
|
||||
|
||||
let model = if ctx.command.starts_with("gemini") {
|
||||
// The pool puts the model into `command` for non-CLI runtimes,
|
||||
// but also check args for a --model flag.
|
||||
ctx.command.clone()
|
||||
} else {
|
||||
// Fall back to args: look for --model <value>
|
||||
ctx.args
|
||||
.iter()
|
||||
.position(|a| a == "--model")
|
||||
.and_then(|i| ctx.args.get(i + 1))
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "gemini-2.5-pro".to_string())
|
||||
};
|
||||
|
||||
let mcp_port = ctx.mcp_port;
|
||||
let mcp_base = format!("http://localhost:{mcp_port}/mcp");
|
||||
|
||||
let client = Client::new();
|
||||
let cancelled = Arc::clone(&self.cancelled);
|
||||
|
||||
// Step 1: Fetch MCP tool definitions and convert to Gemini format.
|
||||
let gemini_tools = fetch_and_convert_mcp_tools(&client, &mcp_base).await?;
|
||||
|
||||
// Step 2: Build the initial conversation contents.
|
||||
let system_instruction = build_system_instruction(&ctx);
|
||||
let mut contents: Vec<Value> = vec![json!({
|
||||
"role": "user",
|
||||
"parts": [{ "text": ctx.prompt }]
|
||||
})];
|
||||
|
||||
let mut total_usage = TokenUsage {
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
total_cost_usd: 0.0,
|
||||
};
|
||||
|
||||
let emit = |event: AgentEvent| {
|
||||
super::super::pty::emit_event(
|
||||
event,
|
||||
&tx,
|
||||
&event_log,
|
||||
log_writer.as_ref().map(|w| w.as_ref()),
|
||||
);
|
||||
};
|
||||
|
||||
emit(AgentEvent::Status {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
status: "running".to_string(),
|
||||
});
|
||||
|
||||
// Step 3: Conversation loop.
|
||||
let mut turn = 0u32;
|
||||
let max_turns = 200; // Safety limit
|
||||
|
||||
loop {
|
||||
if cancelled.load(Ordering::Relaxed) {
|
||||
emit(AgentEvent::Error {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
message: "Agent was stopped by user".to_string(),
|
||||
});
|
||||
return Ok(RuntimeResult {
|
||||
session_id: None,
|
||||
token_usage: Some(total_usage),
|
||||
});
|
||||
}
|
||||
|
||||
turn += 1;
|
||||
if turn > max_turns {
|
||||
emit(AgentEvent::Error {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
message: format!("Exceeded maximum turns ({max_turns})"),
|
||||
});
|
||||
return Ok(RuntimeResult {
|
||||
session_id: None,
|
||||
token_usage: Some(total_usage),
|
||||
});
|
||||
}
|
||||
|
||||
slog!("[gemini] Turn {turn} for {}:{}", ctx.story_id, ctx.agent_name);
|
||||
|
||||
let request_body = build_generate_content_request(
|
||||
&system_instruction,
|
||||
&contents,
|
||||
&gemini_tools,
|
||||
);
|
||||
|
||||
let url = format!(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
|
||||
);
|
||||
|
||||
let response = client
|
||||
.post(&url)
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("Gemini API request failed: {e}"))?;
|
||||
|
||||
let status = response.status();
|
||||
let body: Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to parse Gemini API response: {e}"))?;
|
||||
|
||||
if !status.is_success() {
|
||||
let error_msg = body["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap_or("Unknown API error");
|
||||
let err = format!("Gemini API error ({status}): {error_msg}");
|
||||
emit(AgentEvent::Error {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
message: err.clone(),
|
||||
});
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
// Accumulate token usage.
|
||||
if let Some(usage) = parse_usage_metadata(&body) {
|
||||
total_usage.input_tokens += usage.input_tokens;
|
||||
total_usage.output_tokens += usage.output_tokens;
|
||||
}
|
||||
|
||||
// Extract the candidate response.
|
||||
let candidate = body["candidates"]
|
||||
.as_array()
|
||||
.and_then(|c| c.first())
|
||||
.ok_or_else(|| "No candidates in Gemini response".to_string())?;
|
||||
|
||||
let parts = candidate["content"]["parts"]
|
||||
.as_array()
|
||||
.ok_or_else(|| "No parts in Gemini response candidate".to_string())?;
|
||||
|
||||
// Check finish reason.
|
||||
let finish_reason = candidate["finishReason"].as_str().unwrap_or("");
|
||||
|
||||
// Separate text parts and function call parts.
|
||||
let mut text_parts: Vec<String> = Vec::new();
|
||||
let mut function_calls: Vec<GeminiFunctionCall> = Vec::new();
|
||||
|
||||
for part in parts {
|
||||
if let Some(text) = part["text"].as_str() {
|
||||
text_parts.push(text.to_string());
|
||||
}
|
||||
if let Some(fc) = part.get("functionCall")
|
||||
&& let (Some(name), Some(args)) =
|
||||
(fc["name"].as_str(), fc.get("args"))
|
||||
{
|
||||
function_calls.push(GeminiFunctionCall {
|
||||
name: name.to_string(),
|
||||
args: args.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Emit any text output.
|
||||
for text in &text_parts {
|
||||
if !text.is_empty() {
|
||||
emit(AgentEvent::Output {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
text: text.clone(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// If no function calls, the model is done.
|
||||
if function_calls.is_empty() {
|
||||
emit(AgentEvent::Done {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
session_id: None,
|
||||
});
|
||||
return Ok(RuntimeResult {
|
||||
session_id: None,
|
||||
token_usage: Some(total_usage),
|
||||
});
|
||||
}
|
||||
|
||||
// Add the model's response to the conversation.
|
||||
let model_parts: Vec<Value> = parts.to_vec();
|
||||
contents.push(json!({
|
||||
"role": "model",
|
||||
"parts": model_parts
|
||||
}));
|
||||
|
||||
// Execute function calls via MCP and build response parts.
|
||||
let mut response_parts: Vec<Value> = Vec::new();
|
||||
|
||||
for fc in &function_calls {
|
||||
if cancelled.load(Ordering::Relaxed) {
|
||||
break;
|
||||
}
|
||||
|
||||
slog!(
|
||||
"[gemini] Calling MCP tool '{}' for {}:{}",
|
||||
fc.name,
|
||||
ctx.story_id,
|
||||
ctx.agent_name
|
||||
);
|
||||
|
||||
emit(AgentEvent::Output {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
text: format!("\n[Tool call: {}]\n", fc.name),
|
||||
});
|
||||
|
||||
let tool_result =
|
||||
call_mcp_tool(&client, &mcp_base, &fc.name, &fc.args).await;
|
||||
|
||||
let response_value = match &tool_result {
|
||||
Ok(result) => {
|
||||
emit(AgentEvent::Output {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
text: format!(
|
||||
"[Tool result: {} chars]\n",
|
||||
result.len()
|
||||
),
|
||||
});
|
||||
json!({ "result": result })
|
||||
}
|
||||
Err(e) => {
|
||||
emit(AgentEvent::Output {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
text: format!("[Tool error: {e}]\n"),
|
||||
});
|
||||
json!({ "error": e })
|
||||
}
|
||||
};
|
||||
|
||||
response_parts.push(json!({
|
||||
"functionResponse": {
|
||||
"name": fc.name,
|
||||
"response": response_value
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
// Add function responses to the conversation.
|
||||
contents.push(json!({
|
||||
"role": "user",
|
||||
"parts": response_parts
|
||||
}));
|
||||
|
||||
// If the model indicated it's done despite having function calls,
|
||||
// respect the finish reason.
|
||||
if finish_reason == "STOP" && function_calls.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
emit(AgentEvent::Done {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
session_id: None,
|
||||
});
|
||||
|
||||
Ok(RuntimeResult {
|
||||
session_id: None,
|
||||
token_usage: Some(total_usage),
|
||||
})
|
||||
}
|
||||
|
||||
fn stop(&self) {
|
||||
self.cancelled.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
fn get_status(&self) -> RuntimeStatus {
|
||||
if self.cancelled.load(Ordering::Relaxed) {
|
||||
RuntimeStatus::Failed
|
||||
} else {
|
||||
RuntimeStatus::Idle
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Internal types ───────────────────────────────────────────────────
|
||||
|
||||
struct GeminiFunctionCall {
|
||||
name: String,
|
||||
args: Value,
|
||||
}
|
||||
|
||||
// ── Gemini API types (for serde) ─────────────────────────────────────
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct GeminiFunctionDeclaration {
|
||||
name: String,
|
||||
description: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
parameters: Option<Value>,
|
||||
}
|
||||
|
||||
// ── Helper functions ─────────────────────────────────────────────────
|
||||
|
||||
/// Build the system instruction content from the RuntimeContext.
|
||||
fn build_system_instruction(ctx: &RuntimeContext) -> Value {
|
||||
// Use system_prompt from args if provided via --append-system-prompt,
|
||||
// otherwise use a sensible default.
|
||||
let system_text = ctx
|
||||
.args
|
||||
.iter()
|
||||
.position(|a| a == "--append-system-prompt")
|
||||
.and_then(|i| ctx.args.get(i + 1))
|
||||
.cloned()
|
||||
.unwrap_or_else(|| {
|
||||
format!(
|
||||
"You are an AI coding agent working on story {}. \
|
||||
You have access to tools via function calling. \
|
||||
Use them to complete the task. \
|
||||
Work in the directory: {}",
|
||||
ctx.story_id, ctx.cwd
|
||||
)
|
||||
});
|
||||
|
||||
json!({
|
||||
"parts": [{ "text": system_text }]
|
||||
})
|
||||
}
|
||||
|
||||
/// Build the full `generateContent` request body.
|
||||
fn build_generate_content_request(
|
||||
system_instruction: &Value,
|
||||
contents: &[Value],
|
||||
gemini_tools: &[GeminiFunctionDeclaration],
|
||||
) -> Value {
|
||||
let mut body = json!({
|
||||
"system_instruction": system_instruction,
|
||||
"contents": contents,
|
||||
"generationConfig": {
|
||||
"temperature": 0.2,
|
||||
"maxOutputTokens": 65536,
|
||||
}
|
||||
});
|
||||
|
||||
if !gemini_tools.is_empty() {
|
||||
body["tools"] = json!([{
|
||||
"functionDeclarations": gemini_tools
|
||||
}]);
|
||||
}
|
||||
|
||||
body
|
||||
}
|
||||
|
||||
/// Fetch MCP tool definitions from storkit's MCP server and convert
|
||||
/// them to Gemini function declaration format.
|
||||
async fn fetch_and_convert_mcp_tools(
|
||||
client: &Client,
|
||||
mcp_base: &str,
|
||||
) -> Result<Vec<GeminiFunctionDeclaration>, String> {
|
||||
let request = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/list",
|
||||
"params": {}
|
||||
});
|
||||
|
||||
let response = client
|
||||
.post(mcp_base)
|
||||
.json(&request)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to fetch MCP tools: {e}"))?;
|
||||
|
||||
let body: Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to parse MCP tools response: {e}"))?;
|
||||
|
||||
let tools = body["result"]["tools"]
|
||||
.as_array()
|
||||
.ok_or_else(|| "No tools array in MCP response".to_string())?;
|
||||
|
||||
let mut declarations = Vec::new();
|
||||
|
||||
for tool in tools {
|
||||
let name = tool["name"].as_str().unwrap_or("").to_string();
|
||||
let description = tool["description"].as_str().unwrap_or("").to_string();
|
||||
|
||||
if name.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Convert MCP inputSchema (JSON Schema) to Gemini parameters
|
||||
// (OpenAPI-subset schema). They are structurally compatible for
|
||||
// simple object schemas.
|
||||
let parameters = convert_mcp_schema_to_gemini(tool.get("inputSchema"));
|
||||
|
||||
declarations.push(GeminiFunctionDeclaration {
|
||||
name,
|
||||
description,
|
||||
parameters,
|
||||
});
|
||||
}
|
||||
|
||||
slog!("[gemini] Loaded {} MCP tools as function declarations", declarations.len());
|
||||
Ok(declarations)
|
||||
}
|
||||
|
||||
/// Convert an MCP inputSchema (JSON Schema) to a Gemini-compatible
|
||||
/// OpenAPI-subset parameter schema.
|
||||
///
|
||||
/// Gemini function calling expects parameters in OpenAPI format, which
|
||||
/// is structurally similar to JSON Schema for simple object types.
|
||||
/// We strip unsupported fields and ensure the type is "object".
|
||||
fn convert_mcp_schema_to_gemini(schema: Option<&Value>) -> Option<Value> {
|
||||
let schema = schema?;
|
||||
|
||||
// If the schema has no properties (empty tool), return None.
|
||||
let properties = schema.get("properties")?;
|
||||
if properties.as_object().is_some_and(|p| p.is_empty()) {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut result = json!({
|
||||
"type": "object",
|
||||
"properties": clean_schema_properties(properties),
|
||||
});
|
||||
|
||||
// Preserve required fields if present.
|
||||
if let Some(required) = schema.get("required") {
|
||||
result["required"] = required.clone();
|
||||
}
|
||||
|
||||
Some(result)
|
||||
}
|
||||
|
||||
/// Recursively clean schema properties to be Gemini-compatible.
|
||||
/// Removes unsupported JSON Schema keywords.
|
||||
fn clean_schema_properties(properties: &Value) -> Value {
|
||||
let Some(obj) = properties.as_object() else {
|
||||
return properties.clone();
|
||||
};
|
||||
|
||||
let mut cleaned = serde_json::Map::new();
|
||||
for (key, value) in obj {
|
||||
let mut prop = value.clone();
|
||||
// Remove JSON Schema keywords not supported by Gemini
|
||||
if let Some(p) = prop.as_object_mut() {
|
||||
p.remove("$schema");
|
||||
p.remove("additionalProperties");
|
||||
|
||||
// Recursively clean nested object properties
|
||||
if let Some(nested_props) = p.get("properties").cloned() {
|
||||
p.insert(
|
||||
"properties".to_string(),
|
||||
clean_schema_properties(&nested_props),
|
||||
);
|
||||
}
|
||||
|
||||
// Clean items schema for arrays
|
||||
if let Some(items) = p.get("items").cloned()
|
||||
&& let Some(items_obj) = items.as_object()
|
||||
{
|
||||
let mut cleaned_items = items_obj.clone();
|
||||
cleaned_items.remove("$schema");
|
||||
cleaned_items.remove("additionalProperties");
|
||||
p.insert("items".to_string(), Value::Object(cleaned_items));
|
||||
}
|
||||
}
|
||||
cleaned.insert(key.clone(), prop);
|
||||
}
|
||||
Value::Object(cleaned)
|
||||
}
|
||||
|
||||
/// Call an MCP tool via storkit's MCP server.
|
||||
async fn call_mcp_tool(
|
||||
client: &Client,
|
||||
mcp_base: &str,
|
||||
tool_name: &str,
|
||||
args: &Value,
|
||||
) -> Result<String, String> {
|
||||
let request = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": tool_name,
|
||||
"arguments": args
|
||||
}
|
||||
});
|
||||
|
||||
let response = client
|
||||
.post(mcp_base)
|
||||
.json(&request)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("MCP tool call failed: {e}"))?;
|
||||
|
||||
let body: Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to parse MCP tool response: {e}"))?;
|
||||
|
||||
if let Some(error) = body.get("error") {
|
||||
let msg = error["message"].as_str().unwrap_or("Unknown MCP error");
|
||||
return Err(format!("MCP tool '{tool_name}' error: {msg}"));
|
||||
}
|
||||
|
||||
// MCP tools/call returns { result: { content: [{ type: "text", text: "..." }] } }
|
||||
let content = &body["result"]["content"];
|
||||
if let Some(arr) = content.as_array() {
|
||||
let texts: Vec<&str> = arr
|
||||
.iter()
|
||||
.filter_map(|c| c["text"].as_str())
|
||||
.collect();
|
||||
if !texts.is_empty() {
|
||||
return Ok(texts.join("\n"));
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to serializing the entire result.
|
||||
Ok(body["result"].to_string())
|
||||
}
|
||||
|
||||
/// Parse token usage metadata from a Gemini API response.
|
||||
fn parse_usage_metadata(response: &Value) -> Option<TokenUsage> {
|
||||
let metadata = response.get("usageMetadata")?;
|
||||
Some(TokenUsage {
|
||||
input_tokens: metadata
|
||||
.get("promptTokenCount")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0),
|
||||
output_tokens: metadata
|
||||
.get("candidatesTokenCount")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0),
|
||||
// Gemini doesn't have cache token fields, but we keep the struct uniform.
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
// Google AI API doesn't report cost; leave at 0.
|
||||
total_cost_usd: 0.0,
|
||||
})
|
||||
}
|
||||
|
||||
// ── Tests ────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_simple_object() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"story_id": {
|
||||
"type": "string",
|
||||
"description": "Story identifier"
|
||||
}
|
||||
},
|
||||
"required": ["story_id"]
|
||||
});
|
||||
|
||||
let result = convert_mcp_schema_to_gemini(Some(&schema)).unwrap();
|
||||
assert_eq!(result["type"], "object");
|
||||
assert!(result["properties"]["story_id"].is_object());
|
||||
assert_eq!(result["required"][0], "story_id");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_empty_properties_returns_none() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
});
|
||||
|
||||
assert!(convert_mcp_schema_to_gemini(Some(&schema)).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_none_returns_none() {
|
||||
assert!(convert_mcp_schema_to_gemini(None).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_strips_additional_properties() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {
|
||||
"type": "string",
|
||||
"additionalProperties": false,
|
||||
"$schema": "http://json-schema.org/draft-07/schema#"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let result = convert_mcp_schema_to_gemini(Some(&schema)).unwrap();
|
||||
let name_prop = &result["properties"]["name"];
|
||||
assert!(name_prop.get("additionalProperties").is_none());
|
||||
assert!(name_prop.get("$schema").is_none());
|
||||
assert_eq!(name_prop["type"], "string");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_with_nested_objects() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"config": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"key": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let result = convert_mcp_schema_to_gemini(Some(&schema)).unwrap();
|
||||
assert!(result["properties"]["config"]["properties"]["key"].is_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_with_array_items() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": { "type": "string" }
|
||||
},
|
||||
"additionalProperties": false
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let result = convert_mcp_schema_to_gemini(Some(&schema)).unwrap();
|
||||
let items_schema = &result["properties"]["items"]["items"];
|
||||
assert!(items_schema.get("additionalProperties").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_system_instruction_uses_args() {
|
||||
let ctx = RuntimeContext {
|
||||
story_id: "42_story_test".to_string(),
|
||||
agent_name: "coder-1".to_string(),
|
||||
command: "gemini-2.5-pro".to_string(),
|
||||
args: vec![
|
||||
"--append-system-prompt".to_string(),
|
||||
"Custom system prompt".to_string(),
|
||||
],
|
||||
prompt: "Do the thing".to_string(),
|
||||
cwd: "/tmp/wt".to_string(),
|
||||
inactivity_timeout_secs: 300,
|
||||
mcp_port: 3001,
|
||||
};
|
||||
|
||||
let instruction = build_system_instruction(&ctx);
|
||||
assert_eq!(instruction["parts"][0]["text"], "Custom system prompt");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_system_instruction_default() {
|
||||
let ctx = RuntimeContext {
|
||||
story_id: "42_story_test".to_string(),
|
||||
agent_name: "coder-1".to_string(),
|
||||
command: "gemini-2.5-pro".to_string(),
|
||||
args: vec![],
|
||||
prompt: "Do the thing".to_string(),
|
||||
cwd: "/tmp/wt".to_string(),
|
||||
inactivity_timeout_secs: 300,
|
||||
mcp_port: 3001,
|
||||
};
|
||||
|
||||
let instruction = build_system_instruction(&ctx);
|
||||
let text = instruction["parts"][0]["text"].as_str().unwrap();
|
||||
assert!(text.contains("42_story_test"));
|
||||
assert!(text.contains("/tmp/wt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_generate_content_request_includes_tools() {
|
||||
let system = json!({"parts": [{"text": "system"}]});
|
||||
let contents = vec![json!({"role": "user", "parts": [{"text": "hello"}]})];
|
||||
let tools = vec![GeminiFunctionDeclaration {
|
||||
name: "my_tool".to_string(),
|
||||
description: "A tool".to_string(),
|
||||
parameters: Some(json!({"type": "object", "properties": {"x": {"type": "string"}}})),
|
||||
}];
|
||||
|
||||
let body = build_generate_content_request(&system, &contents, &tools);
|
||||
assert!(body["tools"][0]["functionDeclarations"].is_array());
|
||||
assert_eq!(body["tools"][0]["functionDeclarations"][0]["name"], "my_tool");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_generate_content_request_no_tools() {
|
||||
let system = json!({"parts": [{"text": "system"}]});
|
||||
let contents = vec![json!({"role": "user", "parts": [{"text": "hello"}]})];
|
||||
let tools: Vec<GeminiFunctionDeclaration> = vec![];
|
||||
|
||||
let body = build_generate_content_request(&system, &contents, &tools);
|
||||
assert!(body.get("tools").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_usage_metadata_valid() {
|
||||
let response = json!({
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 100,
|
||||
"candidatesTokenCount": 50,
|
||||
"totalTokenCount": 150
|
||||
}
|
||||
});
|
||||
|
||||
let usage = parse_usage_metadata(&response).unwrap();
|
||||
assert_eq!(usage.input_tokens, 100);
|
||||
assert_eq!(usage.output_tokens, 50);
|
||||
assert_eq!(usage.cache_creation_input_tokens, 0);
|
||||
assert_eq!(usage.total_cost_usd, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_usage_metadata_missing() {
|
||||
let response = json!({"candidates": []});
|
||||
assert!(parse_usage_metadata(&response).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gemini_runtime_stop_sets_cancelled() {
|
||||
let runtime = GeminiRuntime::new();
|
||||
assert_eq!(runtime.get_status(), RuntimeStatus::Idle);
|
||||
runtime.stop();
|
||||
assert_eq!(runtime.get_status(), RuntimeStatus::Failed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_extraction_from_command() {
|
||||
// When command starts with "gemini", use it as model name
|
||||
let ctx = RuntimeContext {
|
||||
story_id: "1".to_string(),
|
||||
agent_name: "coder".to_string(),
|
||||
command: "gemini-2.5-pro".to_string(),
|
||||
args: vec![],
|
||||
prompt: "test".to_string(),
|
||||
cwd: "/tmp".to_string(),
|
||||
inactivity_timeout_secs: 300,
|
||||
mcp_port: 3001,
|
||||
};
|
||||
|
||||
// The model extraction logic is inside start(), but we test the
|
||||
// condition here.
|
||||
assert!(ctx.command.starts_with("gemini"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
mod claude_code;
|
||||
mod gemini;
|
||||
mod openai;
|
||||
|
||||
pub use claude_code::ClaudeCodeRuntime;
|
||||
pub use gemini::GeminiRuntime;
|
||||
pub use openai::OpenAiRuntime;
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use crate::agent_log::AgentLogWriter;
|
||||
|
||||
use super::{AgentEvent, TokenUsage};
|
||||
|
||||
/// Context passed to a runtime when launching an agent session.
|
||||
pub struct RuntimeContext {
|
||||
pub story_id: String,
|
||||
pub agent_name: String,
|
||||
pub command: String,
|
||||
pub args: Vec<String>,
|
||||
pub prompt: String,
|
||||
pub cwd: String,
|
||||
pub inactivity_timeout_secs: u64,
|
||||
/// Port of the storkit MCP server, used by API-based runtimes (Gemini, OpenAI)
|
||||
/// to call back for tool execution.
|
||||
pub mcp_port: u16,
|
||||
}
|
||||
|
||||
/// Result returned by a runtime after the agent session completes.
|
||||
pub struct RuntimeResult {
|
||||
pub session_id: Option<String>,
|
||||
pub token_usage: Option<TokenUsage>,
|
||||
}
|
||||
|
||||
/// Runtime status reported by the backend.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
#[allow(dead_code)]
|
||||
pub enum RuntimeStatus {
|
||||
Idle,
|
||||
Running,
|
||||
Completed,
|
||||
Failed,
|
||||
}
|
||||
|
||||
/// Abstraction over different agent execution backends.
|
||||
///
|
||||
/// Implementations:
|
||||
/// - [`ClaudeCodeRuntime`]: spawns the `claude` CLI via a PTY (default, `runtime = "claude-code"`)
|
||||
///
|
||||
/// Future implementations could include OpenAI and Gemini API runtimes.
|
||||
#[allow(dead_code)]
|
||||
pub trait AgentRuntime: Send + Sync {
|
||||
/// Start the agent and drive it to completion, streaming events through
|
||||
/// the provided broadcast sender and event log.
|
||||
///
|
||||
/// Returns when the agent session finishes (success or error).
|
||||
async fn start(
|
||||
&self,
|
||||
ctx: RuntimeContext,
|
||||
tx: broadcast::Sender<AgentEvent>,
|
||||
event_log: Arc<Mutex<Vec<AgentEvent>>>,
|
||||
log_writer: Option<Arc<Mutex<AgentLogWriter>>>,
|
||||
) -> Result<RuntimeResult, String>;
|
||||
|
||||
/// Stop the running agent.
|
||||
fn stop(&self);
|
||||
|
||||
/// Get the current runtime status.
|
||||
fn get_status(&self) -> RuntimeStatus;
|
||||
|
||||
/// Return any events buffered outside the broadcast channel.
|
||||
///
|
||||
/// PTY-based runtimes stream directly to the broadcast channel; this
|
||||
/// returns empty by default. API-based runtimes may buffer events here.
|
||||
fn stream_events(&self) -> Vec<AgentEvent> {
|
||||
vec![]
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn runtime_context_fields() {
|
||||
let ctx = RuntimeContext {
|
||||
story_id: "42_story_foo".to_string(),
|
||||
agent_name: "coder-1".to_string(),
|
||||
command: "claude".to_string(),
|
||||
args: vec!["--model".to_string(), "sonnet".to_string()],
|
||||
prompt: "Do the thing".to_string(),
|
||||
cwd: "/tmp/wt".to_string(),
|
||||
inactivity_timeout_secs: 300,
|
||||
mcp_port: 3001,
|
||||
};
|
||||
assert_eq!(ctx.story_id, "42_story_foo");
|
||||
assert_eq!(ctx.agent_name, "coder-1");
|
||||
assert_eq!(ctx.command, "claude");
|
||||
assert_eq!(ctx.args.len(), 2);
|
||||
assert_eq!(ctx.prompt, "Do the thing");
|
||||
assert_eq!(ctx.cwd, "/tmp/wt");
|
||||
assert_eq!(ctx.inactivity_timeout_secs, 300);
|
||||
assert_eq!(ctx.mcp_port, 3001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_result_fields() {
|
||||
let result = RuntimeResult {
|
||||
session_id: Some("sess-123".to_string()),
|
||||
token_usage: Some(TokenUsage {
|
||||
input_tokens: 100,
|
||||
output_tokens: 50,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
total_cost_usd: 0.01,
|
||||
}),
|
||||
};
|
||||
assert_eq!(result.session_id, Some("sess-123".to_string()));
|
||||
assert!(result.token_usage.is_some());
|
||||
let usage = result.token_usage.unwrap();
|
||||
assert_eq!(usage.input_tokens, 100);
|
||||
assert_eq!(usage.output_tokens, 50);
|
||||
assert_eq!(usage.total_cost_usd, 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_result_no_usage() {
|
||||
let result = RuntimeResult {
|
||||
session_id: None,
|
||||
token_usage: None,
|
||||
};
|
||||
assert!(result.session_id.is_none());
|
||||
assert!(result.token_usage.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn runtime_status_variants() {
|
||||
assert_eq!(RuntimeStatus::Idle, RuntimeStatus::Idle);
|
||||
assert_ne!(RuntimeStatus::Running, RuntimeStatus::Completed);
|
||||
assert_ne!(RuntimeStatus::Failed, RuntimeStatus::Idle);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_code_runtime_get_status_returns_idle() {
|
||||
use std::collections::HashMap;
|
||||
use crate::io::watcher::WatcherEvent;
|
||||
let killers = Arc::new(Mutex::new(HashMap::new()));
|
||||
let (watcher_tx, _) = broadcast::channel::<WatcherEvent>(16);
|
||||
let runtime = ClaudeCodeRuntime::new(killers, watcher_tx);
|
||||
assert_eq!(runtime.get_status(), RuntimeStatus::Idle);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_code_runtime_stream_events_empty() {
|
||||
use std::collections::HashMap;
|
||||
use crate::io::watcher::WatcherEvent;
|
||||
let killers = Arc::new(Mutex::new(HashMap::new()));
|
||||
let (watcher_tx, _) = broadcast::channel::<WatcherEvent>(16);
|
||||
let runtime = ClaudeCodeRuntime::new(killers, watcher_tx);
|
||||
assert!(runtime.stream_events().is_empty());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,704 @@
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use reqwest::Client;
|
||||
use serde_json::{json, Value};
|
||||
use tokio::sync::broadcast;
|
||||
|
||||
use crate::agent_log::AgentLogWriter;
|
||||
use crate::slog;
|
||||
|
||||
use super::super::{AgentEvent, TokenUsage};
|
||||
use super::{AgentRuntime, RuntimeContext, RuntimeResult, RuntimeStatus};
|
||||
|
||||
// ── Public runtime struct ────────────────────────────────────────────
|
||||
|
||||
/// Agent runtime that drives an OpenAI model (GPT-4o, o3, etc.) through
|
||||
/// the OpenAI Chat Completions API.
|
||||
///
|
||||
/// The runtime:
|
||||
/// 1. Fetches MCP tool definitions from storkit's MCP server.
|
||||
/// 2. Converts them to OpenAI function-calling format.
|
||||
/// 3. Sends the agent prompt + tools to the Chat Completions API.
|
||||
/// 4. Executes any requested tool calls via MCP `tools/call`.
|
||||
/// 5. Loops until the model produces a response with no tool calls.
|
||||
/// 6. Tracks token usage from the API response.
|
||||
pub struct OpenAiRuntime {
|
||||
/// Whether a stop has been requested.
|
||||
cancelled: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl OpenAiRuntime {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
cancelled: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AgentRuntime for OpenAiRuntime {
|
||||
async fn start(
|
||||
&self,
|
||||
ctx: RuntimeContext,
|
||||
tx: broadcast::Sender<AgentEvent>,
|
||||
event_log: Arc<Mutex<Vec<AgentEvent>>>,
|
||||
log_writer: Option<Arc<Mutex<AgentLogWriter>>>,
|
||||
) -> Result<RuntimeResult, String> {
|
||||
let api_key = std::env::var("OPENAI_API_KEY").map_err(|_| {
|
||||
"OPENAI_API_KEY environment variable is not set. \
|
||||
Set it to your OpenAI API key to use the OpenAI runtime."
|
||||
.to_string()
|
||||
})?;
|
||||
|
||||
let model = if ctx.command.starts_with("gpt") || ctx.command.starts_with("o") {
|
||||
// The pool puts the model into `command` for non-CLI runtimes.
|
||||
ctx.command.clone()
|
||||
} else {
|
||||
// Fall back to args: look for --model <value>
|
||||
ctx.args
|
||||
.iter()
|
||||
.position(|a| a == "--model")
|
||||
.and_then(|i| ctx.args.get(i + 1))
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "gpt-4o".to_string())
|
||||
};
|
||||
|
||||
let mcp_port = ctx.mcp_port;
|
||||
let mcp_base = format!("http://localhost:{mcp_port}/mcp");
|
||||
|
||||
let client = Client::new();
|
||||
let cancelled = Arc::clone(&self.cancelled);
|
||||
|
||||
// Step 1: Fetch MCP tool definitions and convert to OpenAI format.
|
||||
let openai_tools = fetch_and_convert_mcp_tools(&client, &mcp_base).await?;
|
||||
|
||||
// Step 2: Build the initial conversation messages.
|
||||
let system_text = build_system_text(&ctx);
|
||||
let mut messages: Vec<Value> = vec![
|
||||
json!({ "role": "system", "content": system_text }),
|
||||
json!({ "role": "user", "content": ctx.prompt }),
|
||||
];
|
||||
|
||||
let mut total_usage = TokenUsage {
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
total_cost_usd: 0.0,
|
||||
};
|
||||
|
||||
let emit = |event: AgentEvent| {
|
||||
super::super::pty::emit_event(
|
||||
event,
|
||||
&tx,
|
||||
&event_log,
|
||||
log_writer.as_ref().map(|w| w.as_ref()),
|
||||
);
|
||||
};
|
||||
|
||||
emit(AgentEvent::Status {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
status: "running".to_string(),
|
||||
});
|
||||
|
||||
// Step 3: Conversation loop.
|
||||
let mut turn = 0u32;
|
||||
let max_turns = 200; // Safety limit
|
||||
|
||||
loop {
|
||||
if cancelled.load(Ordering::Relaxed) {
|
||||
emit(AgentEvent::Error {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
message: "Agent was stopped by user".to_string(),
|
||||
});
|
||||
return Ok(RuntimeResult {
|
||||
session_id: None,
|
||||
token_usage: Some(total_usage),
|
||||
});
|
||||
}
|
||||
|
||||
turn += 1;
|
||||
if turn > max_turns {
|
||||
emit(AgentEvent::Error {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
message: format!("Exceeded maximum turns ({max_turns})"),
|
||||
});
|
||||
return Ok(RuntimeResult {
|
||||
session_id: None,
|
||||
token_usage: Some(total_usage),
|
||||
});
|
||||
}
|
||||
|
||||
slog!(
|
||||
"[openai] Turn {turn} for {}:{}",
|
||||
ctx.story_id,
|
||||
ctx.agent_name
|
||||
);
|
||||
|
||||
let mut request_body = json!({
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"temperature": 0.2,
|
||||
});
|
||||
|
||||
if !openai_tools.is_empty() {
|
||||
request_body["tools"] = json!(openai_tools);
|
||||
}
|
||||
|
||||
let response = client
|
||||
.post("https://api.openai.com/v1/chat/completions")
|
||||
.bearer_auth(&api_key)
|
||||
.json(&request_body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("OpenAI API request failed: {e}"))?;
|
||||
|
||||
let status = response.status();
|
||||
let body: Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to parse OpenAI API response: {e}"))?;
|
||||
|
||||
if !status.is_success() {
|
||||
let error_msg = body["error"]["message"]
|
||||
.as_str()
|
||||
.unwrap_or("Unknown API error");
|
||||
let err = format!("OpenAI API error ({status}): {error_msg}");
|
||||
emit(AgentEvent::Error {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
message: err.clone(),
|
||||
});
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
// Accumulate token usage.
|
||||
if let Some(usage) = parse_usage(&body) {
|
||||
total_usage.input_tokens += usage.input_tokens;
|
||||
total_usage.output_tokens += usage.output_tokens;
|
||||
}
|
||||
|
||||
// Extract the first choice.
|
||||
let choice = body["choices"]
|
||||
.as_array()
|
||||
.and_then(|c| c.first())
|
||||
.ok_or_else(|| "No choices in OpenAI response".to_string())?;
|
||||
|
||||
let message = &choice["message"];
|
||||
let content = message["content"].as_str().unwrap_or("");
|
||||
|
||||
// Emit any text content.
|
||||
if !content.is_empty() {
|
||||
emit(AgentEvent::Output {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
text: content.to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
// Check for tool calls.
|
||||
let tool_calls = message["tool_calls"].as_array();
|
||||
|
||||
if tool_calls.is_none() || tool_calls.is_some_and(|tc| tc.is_empty()) {
|
||||
// No tool calls — model is done.
|
||||
emit(AgentEvent::Done {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
session_id: None,
|
||||
});
|
||||
return Ok(RuntimeResult {
|
||||
session_id: None,
|
||||
token_usage: Some(total_usage),
|
||||
});
|
||||
}
|
||||
|
||||
let tool_calls = tool_calls.unwrap();
|
||||
|
||||
// Add the assistant message (with tool_calls) to the conversation.
|
||||
messages.push(message.clone());
|
||||
|
||||
// Execute each tool call via MCP and add results.
|
||||
for tc in tool_calls {
|
||||
if cancelled.load(Ordering::Relaxed) {
|
||||
break;
|
||||
}
|
||||
|
||||
let call_id = tc["id"].as_str().unwrap_or("");
|
||||
let function = &tc["function"];
|
||||
let tool_name = function["name"].as_str().unwrap_or("");
|
||||
let arguments_str = function["arguments"].as_str().unwrap_or("{}");
|
||||
|
||||
let args: Value = serde_json::from_str(arguments_str).unwrap_or(json!({}));
|
||||
|
||||
slog!(
|
||||
"[openai] Calling MCP tool '{}' for {}:{}",
|
||||
tool_name,
|
||||
ctx.story_id,
|
||||
ctx.agent_name
|
||||
);
|
||||
|
||||
emit(AgentEvent::Output {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
text: format!("\n[Tool call: {tool_name}]\n"),
|
||||
});
|
||||
|
||||
let tool_result = call_mcp_tool(&client, &mcp_base, tool_name, &args).await;
|
||||
|
||||
let result_content = match &tool_result {
|
||||
Ok(result) => {
|
||||
emit(AgentEvent::Output {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
text: format!("[Tool result: {} chars]\n", result.len()),
|
||||
});
|
||||
result.clone()
|
||||
}
|
||||
Err(e) => {
|
||||
emit(AgentEvent::Output {
|
||||
story_id: ctx.story_id.clone(),
|
||||
agent_name: ctx.agent_name.clone(),
|
||||
text: format!("[Tool error: {e}]\n"),
|
||||
});
|
||||
format!("Error: {e}")
|
||||
}
|
||||
};
|
||||
|
||||
// OpenAI expects tool results as role=tool messages with
|
||||
// the matching tool_call_id.
|
||||
messages.push(json!({
|
||||
"role": "tool",
|
||||
"tool_call_id": call_id,
|
||||
"content": result_content,
|
||||
}));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn stop(&self) {
|
||||
self.cancelled.store(true, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
fn get_status(&self) -> RuntimeStatus {
|
||||
if self.cancelled.load(Ordering::Relaxed) {
|
||||
RuntimeStatus::Failed
|
||||
} else {
|
||||
RuntimeStatus::Idle
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Helper functions ─────────────────────────────────────────────────
|
||||
|
||||
/// Build the system message text from the RuntimeContext.
|
||||
fn build_system_text(ctx: &RuntimeContext) -> String {
|
||||
ctx.args
|
||||
.iter()
|
||||
.position(|a| a == "--append-system-prompt")
|
||||
.and_then(|i| ctx.args.get(i + 1))
|
||||
.cloned()
|
||||
.unwrap_or_else(|| {
|
||||
format!(
|
||||
"You are an AI coding agent working on story {}. \
|
||||
You have access to tools via function calling. \
|
||||
Use them to complete the task. \
|
||||
Work in the directory: {}",
|
||||
ctx.story_id, ctx.cwd
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Fetch MCP tool definitions from storkit's MCP server and convert
|
||||
/// them to OpenAI function-calling format.
|
||||
async fn fetch_and_convert_mcp_tools(
|
||||
client: &Client,
|
||||
mcp_base: &str,
|
||||
) -> Result<Vec<Value>, String> {
|
||||
let request = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/list",
|
||||
"params": {}
|
||||
});
|
||||
|
||||
let response = client
|
||||
.post(mcp_base)
|
||||
.json(&request)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to fetch MCP tools: {e}"))?;
|
||||
|
||||
let body: Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to parse MCP tools response: {e}"))?;
|
||||
|
||||
let tools = body["result"]["tools"]
|
||||
.as_array()
|
||||
.ok_or_else(|| "No tools array in MCP response".to_string())?;
|
||||
|
||||
let mut openai_tools = Vec::new();
|
||||
|
||||
for tool in tools {
|
||||
let name = tool["name"].as_str().unwrap_or("").to_string();
|
||||
let description = tool["description"].as_str().unwrap_or("").to_string();
|
||||
|
||||
if name.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// OpenAI function calling uses JSON Schema natively for parameters,
|
||||
// so the MCP inputSchema can be used with minimal cleanup.
|
||||
let parameters = convert_mcp_schema_to_openai(tool.get("inputSchema"));
|
||||
|
||||
openai_tools.push(json!({
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"description": description,
|
||||
"parameters": parameters.unwrap_or_else(|| json!({"type": "object", "properties": {}})),
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
slog!(
|
||||
"[openai] Loaded {} MCP tools as function definitions",
|
||||
openai_tools.len()
|
||||
);
|
||||
Ok(openai_tools)
|
||||
}
|
||||
|
||||
/// Convert an MCP inputSchema (JSON Schema) to OpenAI-compatible
|
||||
/// function parameters.
|
||||
///
|
||||
/// OpenAI uses JSON Schema natively, so less transformation is needed
|
||||
/// compared to Gemini. We still strip `$schema` to keep payloads clean.
|
||||
fn convert_mcp_schema_to_openai(schema: Option<&Value>) -> Option<Value> {
|
||||
let schema = schema?;
|
||||
|
||||
let mut result = json!({
|
||||
"type": "object",
|
||||
});
|
||||
|
||||
if let Some(properties) = schema.get("properties") {
|
||||
result["properties"] = clean_schema_properties(properties);
|
||||
} else {
|
||||
result["properties"] = json!({});
|
||||
}
|
||||
|
||||
if let Some(required) = schema.get("required") {
|
||||
result["required"] = required.clone();
|
||||
}
|
||||
|
||||
// OpenAI recommends additionalProperties: false for strict mode.
|
||||
result["additionalProperties"] = json!(false);
|
||||
|
||||
Some(result)
|
||||
}
|
||||
|
||||
/// Recursively clean schema properties, removing unsupported keywords.
|
||||
fn clean_schema_properties(properties: &Value) -> Value {
|
||||
let Some(obj) = properties.as_object() else {
|
||||
return properties.clone();
|
||||
};
|
||||
|
||||
let mut cleaned = serde_json::Map::new();
|
||||
for (key, value) in obj {
|
||||
let mut prop = value.clone();
|
||||
if let Some(p) = prop.as_object_mut() {
|
||||
p.remove("$schema");
|
||||
|
||||
// Recursively clean nested object properties.
|
||||
if let Some(nested_props) = p.get("properties").cloned() {
|
||||
p.insert(
|
||||
"properties".to_string(),
|
||||
clean_schema_properties(&nested_props),
|
||||
);
|
||||
}
|
||||
|
||||
// Clean items schema for arrays.
|
||||
if let Some(items) = p.get("items").cloned()
|
||||
&& let Some(items_obj) = items.as_object()
|
||||
{
|
||||
let mut cleaned_items = items_obj.clone();
|
||||
cleaned_items.remove("$schema");
|
||||
p.insert("items".to_string(), Value::Object(cleaned_items));
|
||||
}
|
||||
}
|
||||
cleaned.insert(key.clone(), prop);
|
||||
}
|
||||
Value::Object(cleaned)
|
||||
}
|
||||
|
||||
/// Call an MCP tool via storkit's MCP server.
|
||||
async fn call_mcp_tool(
|
||||
client: &Client,
|
||||
mcp_base: &str,
|
||||
tool_name: &str,
|
||||
args: &Value,
|
||||
) -> Result<String, String> {
|
||||
let request = json!({
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": tool_name,
|
||||
"arguments": args
|
||||
}
|
||||
});
|
||||
|
||||
let response = client
|
||||
.post(mcp_base)
|
||||
.json(&request)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("MCP tool call failed: {e}"))?;
|
||||
|
||||
let body: Value = response
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| format!("Failed to parse MCP tool response: {e}"))?;
|
||||
|
||||
if let Some(error) = body.get("error") {
|
||||
let msg = error["message"].as_str().unwrap_or("Unknown MCP error");
|
||||
return Err(format!("MCP tool '{tool_name}' error: {msg}"));
|
||||
}
|
||||
|
||||
// MCP tools/call returns { result: { content: [{ type: "text", text: "..." }] } }
|
||||
let content = &body["result"]["content"];
|
||||
if let Some(arr) = content.as_array() {
|
||||
let texts: Vec<&str> = arr
|
||||
.iter()
|
||||
.filter_map(|c| c["text"].as_str())
|
||||
.collect();
|
||||
if !texts.is_empty() {
|
||||
return Ok(texts.join("\n"));
|
||||
}
|
||||
}
|
||||
|
||||
// Fall back to serializing the entire result.
|
||||
Ok(body["result"].to_string())
|
||||
}
|
||||
|
||||
/// Parse token usage from an OpenAI API response.
|
||||
fn parse_usage(response: &Value) -> Option<TokenUsage> {
|
||||
let usage = response.get("usage")?;
|
||||
Some(TokenUsage {
|
||||
input_tokens: usage
|
||||
.get("prompt_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0),
|
||||
output_tokens: usage
|
||||
.get("completion_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0),
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
// OpenAI API doesn't report cost directly; leave at 0.
|
||||
total_cost_usd: 0.0,
|
||||
})
|
||||
}
|
||||
|
||||
// ── Tests ────────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_simple_object() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"story_id": {
|
||||
"type": "string",
|
||||
"description": "Story identifier"
|
||||
}
|
||||
},
|
||||
"required": ["story_id"]
|
||||
});
|
||||
|
||||
let result = convert_mcp_schema_to_openai(Some(&schema)).unwrap();
|
||||
assert_eq!(result["type"], "object");
|
||||
assert!(result["properties"]["story_id"].is_object());
|
||||
assert_eq!(result["required"][0], "story_id");
|
||||
assert_eq!(result["additionalProperties"], false);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_empty_properties() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
});
|
||||
|
||||
let result = convert_mcp_schema_to_openai(Some(&schema)).unwrap();
|
||||
assert_eq!(result["type"], "object");
|
||||
assert!(result["properties"].as_object().unwrap().is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_none_returns_none() {
|
||||
assert!(convert_mcp_schema_to_openai(None).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_strips_dollar_schema() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {
|
||||
"type": "string",
|
||||
"$schema": "http://json-schema.org/draft-07/schema#"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let result = convert_mcp_schema_to_openai(Some(&schema)).unwrap();
|
||||
let name_prop = &result["properties"]["name"];
|
||||
assert!(name_prop.get("$schema").is_none());
|
||||
assert_eq!(name_prop["type"], "string");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_with_nested_objects() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"config": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"key": { "type": "string" }
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let result = convert_mcp_schema_to_openai(Some(&schema)).unwrap();
|
||||
assert!(result["properties"]["config"]["properties"]["key"].is_object());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn convert_mcp_schema_with_array_items() {
|
||||
let schema = json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": { "type": "string" }
|
||||
},
|
||||
"$schema": "http://json-schema.org/draft-07/schema#"
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let result = convert_mcp_schema_to_openai(Some(&schema)).unwrap();
|
||||
let items_schema = &result["properties"]["items"]["items"];
|
||||
assert!(items_schema.get("$schema").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_system_text_uses_args() {
|
||||
let ctx = RuntimeContext {
|
||||
story_id: "42_story_test".to_string(),
|
||||
agent_name: "coder-1".to_string(),
|
||||
command: "gpt-4o".to_string(),
|
||||
args: vec![
|
||||
"--append-system-prompt".to_string(),
|
||||
"Custom system prompt".to_string(),
|
||||
],
|
||||
prompt: "Do the thing".to_string(),
|
||||
cwd: "/tmp/wt".to_string(),
|
||||
inactivity_timeout_secs: 300,
|
||||
mcp_port: 3001,
|
||||
};
|
||||
|
||||
assert_eq!(build_system_text(&ctx), "Custom system prompt");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_system_text_default() {
|
||||
let ctx = RuntimeContext {
|
||||
story_id: "42_story_test".to_string(),
|
||||
agent_name: "coder-1".to_string(),
|
||||
command: "gpt-4o".to_string(),
|
||||
args: vec![],
|
||||
prompt: "Do the thing".to_string(),
|
||||
cwd: "/tmp/wt".to_string(),
|
||||
inactivity_timeout_secs: 300,
|
||||
mcp_port: 3001,
|
||||
};
|
||||
|
||||
let text = build_system_text(&ctx);
|
||||
assert!(text.contains("42_story_test"));
|
||||
assert!(text.contains("/tmp/wt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_usage_valid() {
|
||||
let response = json!({
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"total_tokens": 150
|
||||
}
|
||||
});
|
||||
|
||||
let usage = parse_usage(&response).unwrap();
|
||||
assert_eq!(usage.input_tokens, 100);
|
||||
assert_eq!(usage.output_tokens, 50);
|
||||
assert_eq!(usage.cache_creation_input_tokens, 0);
|
||||
assert_eq!(usage.total_cost_usd, 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_usage_missing() {
|
||||
let response = json!({"choices": []});
|
||||
assert!(parse_usage(&response).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_runtime_stop_sets_cancelled() {
|
||||
let runtime = OpenAiRuntime::new();
|
||||
assert_eq!(runtime.get_status(), RuntimeStatus::Idle);
|
||||
runtime.stop();
|
||||
assert_eq!(runtime.get_status(), RuntimeStatus::Failed);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_extraction_from_command_gpt() {
|
||||
let ctx = RuntimeContext {
|
||||
story_id: "1".to_string(),
|
||||
agent_name: "coder".to_string(),
|
||||
command: "gpt-4o".to_string(),
|
||||
args: vec![],
|
||||
prompt: "test".to_string(),
|
||||
cwd: "/tmp".to_string(),
|
||||
inactivity_timeout_secs: 300,
|
||||
mcp_port: 3001,
|
||||
};
|
||||
assert!(ctx.command.starts_with("gpt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_extraction_from_command_o3() {
|
||||
let ctx = RuntimeContext {
|
||||
story_id: "1".to_string(),
|
||||
agent_name: "coder".to_string(),
|
||||
command: "o3".to_string(),
|
||||
args: vec![],
|
||||
prompt: "test".to_string(),
|
||||
cwd: "/tmp".to_string(),
|
||||
inactivity_timeout_secs: 300,
|
||||
mcp_port: 3001,
|
||||
};
|
||||
assert!(ctx.command.starts_with("o"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user