From 10dd239c9257e7ed618ec6dbd7d5ce588d7f3c39 Mon Sep 17 00:00:00 2001 From: Huskies Agent Date: Tue, 21 Jul 2026 17:23:42 +0000 Subject: [PATCH] huskies: merge 1246 refactor Extract shared LLM-runtime logic duplicated between gemini and openai runtimes --- server/src/agents/runtime/api_common.rs | 404 ++++++++++++++++++++++++ server/src/agents/runtime/gemini/api.rs | 62 +--- server/src/agents/runtime/gemini/mcp.rs | 82 +---- server/src/agents/runtime/gemini/mod.rs | 147 ++------- server/src/agents/runtime/mod.rs | 8 +- server/src/agents/runtime/openai.rs | 289 ++--------------- 6 files changed, 461 insertions(+), 531 deletions(-) create mode 100644 server/src/agents/runtime/api_common.rs diff --git a/server/src/agents/runtime/api_common.rs b/server/src/agents/runtime/api_common.rs new file mode 100644 index 00000000..24b367ba --- /dev/null +++ b/server/src/agents/runtime/api_common.rs @@ -0,0 +1,404 @@ +//! Shared helpers for the API-based agent runtimes (Gemini, OpenAI), which +//! talk directly to a provider's REST API rather than spawning a CLI over a +//! PTY. Both runtimes drive an almost-identical turn loop against different +//! wire formats; this module holds the logic that doesn't vary between them. +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; + +use serde_json::Value; +use tokio::sync::broadcast; + +use crate::agent_log::AgentLogWriter; + +use super::super::{AgentEvent, TokenUsage}; +use super::{RuntimeContext, RuntimeResult, RuntimeStatus}; + +/// Cancellation flag shared by the API-based runtimes' `stop()`/`get_status()`. +pub(super) struct CancellationFlag { + cancelled: Arc, +} + +impl CancellationFlag { + /// Create a fresh, un-cancelled flag. + pub(super) fn new() -> Self { + Self { + cancelled: Arc::new(AtomicBool::new(false)), + } + } + + /// Clone of the underlying flag, for the conversation loop to poll. + pub(super) fn handle(&self) -> Arc { + Arc::clone(&self.cancelled) + } + + /// Request a stop. + pub(super) fn stop(&self) { + self.cancelled.store(true, Ordering::Relaxed); + } + + /// Report `Failed` once stopped, `Idle` otherwise. + pub(super) fn status(&self) -> RuntimeStatus { + if self.cancelled.load(Ordering::Relaxed) { + RuntimeStatus::Failed + } else { + RuntimeStatus::Idle + } + } +} + +/// Build the event-emitting closure shared by API-based runtimes: forwards +/// events to the broadcast channel, the in-memory event log, and (optionally) +/// the on-disk log writer. +pub(super) fn make_emit( + tx: broadcast::Sender, + event_log: Arc>>, + log_writer: Option>>, +) -> impl Fn(AgentEvent) { + move |event: AgentEvent| { + super::super::pty::emit_event( + event, + &tx, + &event_log, + log_writer.as_ref().map(|w| w.as_ref()), + ); + } +} + +/// Zeroed token-usage accumulator, seeded before an API runtime's +/// conversation loop starts accumulating per-turn usage. +pub(super) fn zero_usage() -> TokenUsage { + TokenUsage { + input_tokens: 0, + output_tokens: 0, + cache_creation_input_tokens: 0, + cache_read_input_tokens: 0, + total_cost_usd: 0.0, + } +} + +/// Emit the initial "running" status event. +pub(super) fn emit_running(ctx: &RuntimeContext, emit: &impl Fn(AgentEvent)) { + emit(AgentEvent::Status { + story_id: ctx.story_id.clone(), + agent_name: ctx.agent_name.clone(), + status: "running".to_string(), + }); +} + +/// Set up an API-based runtime's conversation loop: builds the event +/// emitter, emits the initial "running" status, and returns it alongside a +/// zeroed usage accumulator and the turn counter (starting at 0). +pub(super) fn start_conversation_loop( + ctx: &RuntimeContext, + tx: broadcast::Sender, + event_log: Arc>>, + log_writer: Option>>, +) -> (impl Fn(AgentEvent), TokenUsage, u32) { + let emit = make_emit(tx, event_log, log_writer); + emit_running(ctx, &emit); + (emit, zero_usage(), 0u32) +} + +/// Build a successful `RuntimeResult` carrying the given token usage. All +/// API-based runtimes report `exit_ok: true` and leave the CLI-only fields +/// (`aborted_signal`, `rate_limit_exit`, `rate_limit_reset_at`) at their +/// defaults, since those concepts don't apply outside the PTY runtime. +pub(super) fn api_runtime_result(total_usage: TokenUsage) -> RuntimeResult { + RuntimeResult { + session_id: None, + token_usage: Some(total_usage), + exit_ok: true, + aborted_signal: false, + rate_limit_exit: false, + rate_limit_reset_at: None, + } +} + +/// Emit the `Done` event and build the final successful result once the +/// model produces a response with no further tool/function calls. +pub(super) fn done_result( + ctx: &RuntimeContext, + emit: &impl Fn(AgentEvent), + total_usage: TokenUsage, +) -> RuntimeResult { + emit(AgentEvent::Done { + story_id: ctx.story_id.clone(), + agent_name: ctx.agent_name.clone(), + session_id: None, + }); + api_runtime_result(total_usage) +} + +/// Safety limit on conversation turns for API-based runtimes, shared so +/// both the guard check and its error message stay in sync. +const MAX_TURNS: u32 = 200; + +/// Check the per-turn cancellation/max-turns guard at the top of an API +/// runtime's conversation loop. Returns `Some(result)` when the loop should +/// stop immediately (either the user requested a stop, or the safety turn +/// limit was exceeded); otherwise increments `*turn` and returns `None`. +pub(super) fn check_loop_guard( + ctx: &RuntimeContext, + cancelled: &AtomicBool, + turn: &mut u32, + total_usage: &TokenUsage, + emit: &impl Fn(AgentEvent), +) -> Option { + 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 Some(api_runtime_result(total_usage.clone())); + } + + *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 Some(api_runtime_result(total_usage.clone())); + } + + None +} + +/// Extract the model name for an API-based runtime: the agent pool stashes +/// the model directly in `ctx.command` for non-CLI runtimes, detected here +/// via `is_command_a_model`; otherwise fall back to a `--model` arg, and +/// finally `default_model`. +pub(super) fn extract_model( + ctx: &RuntimeContext, + is_command_a_model: impl Fn(&str) -> bool, + default_model: &str, +) -> String { + if is_command_a_model(&ctx.command) { + ctx.command.clone() + } else { + ctx.args + .iter() + .position(|a| a == "--model") + .and_then(|i| ctx.args.get(i + 1)) + .cloned() + .unwrap_or_else(|| default_model.to_string()) + } +} + +/// Build the default system-prompt text shared by both API-based runtimes: +/// prefers an explicit `--append-system-prompt` arg (set by the agent pool), +/// else falls back to a generic tool-calling preamble. +pub(super) 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 + ) + }) +} + +/// Recursively clean an MCP JSON-Schema `properties` object into a provider +/// function-calling schema: strips `$schema` (always) and +/// `additionalProperties` (when the provider doesn't support it, e.g. +/// Gemini) from the top level and from nested `properties`/`items`. +pub(super) fn clean_schema_properties( + properties: &Value, + strip_additional_properties: bool, +) -> 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"); + if strip_additional_properties { + p.remove("additionalProperties"); + } + + if let Some(nested_props) = p.get("properties").cloned() { + p.insert( + "properties".to_string(), + clean_schema_properties(&nested_props, strip_additional_properties), + ); + } + + 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"); + if strip_additional_properties { + cleaned_items.remove("additionalProperties"); + } + p.insert("items".to_string(), Value::Object(cleaned_items)); + } + } + cleaned.insert(key.clone(), prop); + } + Value::Object(cleaned) +} + +// ── Test helpers ───────────────────────────────────────────────────── + +/// Build a throwaway `AppContext` backed by a temp directory, for tests +/// that need a `RuntimeContext.app_ctx` but don't exercise it. +#[cfg(test)] +pub(super) fn test_app_ctx() -> Arc { + let tmp = tempfile::tempdir().unwrap(); + Arc::new(crate::http::context::AppContext::new_test( + tmp.path().to_path_buf(), + )) +} + +/// Build a `RuntimeContext` with sensible test defaults, overriding only +/// `command` and `args` (the fields the API-runtime tests vary). +#[cfg(test)] +pub(super) fn test_runtime_context(command: &str, args: Vec) -> RuntimeContext { + RuntimeContext { + story_id: "42_story_test".to_string(), + agent_name: "coder-1".to_string(), + command: command.to_string(), + args, + prompt: "Do the thing".to_string(), + cwd: "/tmp/wt".to_string(), + inactivity_timeout_secs: 300, + app_ctx: Some(test_app_ctx()), + session_id_to_resume: None, + fresh_prompt: None, + project_root: std::path::PathBuf::from("/tmp/project"), + model: None, + } +} + +// ── Tests ──────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn clean_schema_strips_dollar_schema_always() { + let schema = json!({ + "name": { "type": "string", "$schema": "http://json-schema.org/draft-07/schema#" } + }); + let result = clean_schema_properties(&schema, false); + assert!(result["name"].get("$schema").is_none()); + } + + #[test] + fn clean_schema_strips_additional_properties_when_requested() { + let schema = json!({ + "name": { "type": "string", "additionalProperties": false } + }); + let result = clean_schema_properties(&schema, true); + assert!(result["name"].get("additionalProperties").is_none()); + } + + #[test] + fn clean_schema_keeps_additional_properties_when_not_requested() { + let schema = json!({ + "name": { "type": "object", "additionalProperties": false } + }); + let result = clean_schema_properties(&schema, false); + assert!(result["name"].get("additionalProperties").is_some()); + } + + #[test] + fn clean_schema_recurses_into_nested_object_properties() { + let schema = json!({ + "config": { + "type": "object", + "properties": { + "key": { "type": "string", "$schema": "x" } + } + } + }); + let result = clean_schema_properties(&schema, false); + assert!(result["config"]["properties"]["key"].is_object()); + assert!( + result["config"]["properties"]["key"] + .get("$schema") + .is_none() + ); + } + + #[test] + fn clean_schema_recurses_into_array_items() { + let schema = json!({ + "items": { + "type": "array", + "items": { + "type": "object", + "properties": { "name": { "type": "string" } }, + "additionalProperties": false, + "$schema": "x" + } + } + }); + let result = clean_schema_properties(&schema, true); + let items_schema = &result["items"]["items"]; + assert!(items_schema.get("additionalProperties").is_none()); + assert!(items_schema.get("$schema").is_none()); + } + + #[test] + fn extract_model_uses_command_when_it_matches() { + let ctx = test_runtime_context("gpt-4o", vec![]); + assert_eq!( + extract_model(&ctx, |c| c.starts_with("gpt"), "fallback"), + "gpt-4o" + ); + } + + #[test] + fn extract_model_falls_back_to_args() { + let ctx = test_runtime_context("claude", vec!["--model".to_string(), "custom".to_string()]); + assert_eq!( + extract_model(&ctx, |c| c.starts_with("gpt"), "fallback"), + "custom" + ); + } + + #[test] + fn extract_model_falls_back_to_default() { + let ctx = test_runtime_context("claude", vec![]); + assert_eq!( + extract_model(&ctx, |c| c.starts_with("gpt"), "fallback"), + "fallback" + ); + } + + #[test] + fn build_system_text_uses_args() { + let ctx = test_runtime_context( + "gpt-4o", + vec![ + "--append-system-prompt".to_string(), + "Custom system prompt".to_string(), + ], + ); + assert_eq!(build_system_text(&ctx), "Custom system prompt"); + } + + #[test] + fn build_system_text_default() { + let ctx = test_runtime_context("gpt-4o", vec![]); + let text = build_system_text(&ctx); + assert!(text.contains("42_story_test")); + assert!(text.contains("/tmp/wt")); + } +} diff --git a/server/src/agents/runtime/gemini/api.rs b/server/src/agents/runtime/gemini/api.rs index ac9ef3f0..c2139e69 100644 --- a/server/src/agents/runtime/gemini/api.rs +++ b/server/src/agents/runtime/gemini/api.rs @@ -4,6 +4,7 @@ use serde_json::{Value, json}; use super::super::super::TokenUsage; use super::super::RuntimeContext; +use super::super::api_common::build_system_text; // ── Gemini API types ───────────────────────────────────────────────── @@ -19,26 +20,8 @@ pub(super) struct GeminiFunctionDeclaration { /// Build the system instruction content from the RuntimeContext. pub(super) 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 }] + "parts": [{ "text": build_system_text(ctx) }] }) } @@ -92,34 +75,18 @@ pub(super) fn parse_usage_metadata(response: &Value) -> Option { #[cfg(test)] mod tests { + use super::super::super::api_common::test_runtime_context; use super::*; - use crate::http::context::AppContext; - use std::sync::Arc; - - fn test_app_ctx() -> Arc { - let tmp = tempfile::tempdir().unwrap(); - Arc::new(AppContext::new_test(tmp.path().to_path_buf())) - } #[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![ + let ctx = test_runtime_context( + "gemini-2.5-pro", + 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, - app_ctx: Some(test_app_ctx()), - session_id_to_resume: None, - fresh_prompt: None, - project_root: std::path::PathBuf::from("/tmp/project"), - model: None, - }; + ); let instruction = build_system_instruction(&ctx); assert_eq!(instruction["parts"][0]["text"], "Custom system prompt"); @@ -127,20 +94,7 @@ mod tests { #[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, - app_ctx: Some(test_app_ctx()), - session_id_to_resume: None, - fresh_prompt: None, - project_root: std::path::PathBuf::from("/tmp/project"), - model: None, - }; + let ctx = test_runtime_context("gemini-2.5-pro", vec![]); let instruction = build_system_instruction(&ctx); let text = instruction["parts"][0]["text"].as_str().unwrap(); diff --git a/server/src/agents/runtime/gemini/mcp.rs b/server/src/agents/runtime/gemini/mcp.rs index 5b119f56..0888ef1b 100644 --- a/server/src/agents/runtime/gemini/mcp.rs +++ b/server/src/agents/runtime/gemini/mcp.rs @@ -8,6 +8,7 @@ use crate::slog; use crate::http::mcp::tools_list::list_tools; +use super::super::api_common::clean_schema_properties; use super::api::GeminiFunctionDeclaration; // ── MCP tool loading ──────────────────────────────────────────────── @@ -62,7 +63,7 @@ pub(super) fn convert_mcp_schema_to_gemini(schema: Option<&Value>) -> Option) -> Option 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) -} - // ── Tests ──────────────────────────────────────────────────────────── #[cfg(test)] @@ -170,45 +133,4 @@ mod tests { 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()); - } } diff --git a/server/src/agents/runtime/gemini/mod.rs b/server/src/agents/runtime/gemini/mod.rs index 88e7cf62..9349b1c9 100644 --- a/server/src/agents/runtime/gemini/mod.rs +++ b/server/src/agents/runtime/gemini/mod.rs @@ -1,5 +1,5 @@ //! Gemini runtime — drives Google Gemini API sessions as agent backends. -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::Ordering; use std::sync::{Arc, Mutex}; use reqwest::Client; @@ -10,7 +10,10 @@ use crate::agent_log::AgentLogWriter; use crate::http::mcp::dispatch::dispatch_tool_call; use crate::slog; -use super::super::{AgentEvent, TokenUsage}; +use super::super::AgentEvent; +use super::api_common::{ + CancellationFlag, check_loop_guard, done_result, extract_model, start_conversation_loop, +}; use super::{AgentRuntime, RuntimeContext, RuntimeResult, RuntimeStatus}; mod api; @@ -40,14 +43,14 @@ struct GeminiFunctionCall { /// 6. Tracks token usage from the API response metadata. pub struct GeminiRuntime { /// Whether a stop has been requested. - cancelled: Arc, + cancelled: CancellationFlag, } impl GeminiRuntime { /// Create a new Gemini runtime instance. pub fn new() -> Self { Self { - cancelled: Arc::new(AtomicBool::new(false)), + cancelled: CancellationFlag::new(), } } } @@ -66,19 +69,7 @@ impl AgentRuntime for GeminiRuntime { .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 - 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 model = extract_model(&ctx, |c| c.starts_with("gemini"), "gemini-2.5-pro"); let app_ctx = ctx .app_ctx @@ -86,7 +77,7 @@ impl AgentRuntime for GeminiRuntime { .ok_or_else(|| "Gemini runtime requires app_ctx to be set".to_string())?; let client = Client::new(); - let cancelled = Arc::clone(&self.cancelled); + let cancelled = self.cancelled.handle(); // Step 1: Load MCP tool definitions and convert to Gemini format. let gemini_tools = convert_mcp_tools_to_gemini(); @@ -98,65 +89,14 @@ impl AgentRuntime for GeminiRuntime { "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(), - }); + let (emit, mut total_usage, mut turn) = + start_conversation_loop(&ctx, tx, event_log, log_writer); // 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), - exit_ok: true, - aborted_signal: false, - rate_limit_exit: false, - rate_limit_reset_at: None, - }); - } - - 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), - exit_ok: true, - aborted_signal: false, - rate_limit_exit: false, - rate_limit_reset_at: None, - }); + if let Some(result) = check_loop_guard(&ctx, &cancelled, &mut turn, &total_usage, &emit) + { + return Ok(result); } slog!( @@ -248,19 +188,7 @@ impl AgentRuntime for GeminiRuntime { // 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), - exit_ok: true, - aborted_signal: false, - rate_limit_exit: false, - rate_limit_reset_at: None, - }); + return Ok(done_result(&ctx, &emit, total_usage)); } // Add the model's response to the conversation. @@ -333,32 +261,15 @@ impl AgentRuntime for GeminiRuntime { } } - 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), - exit_ok: true, - aborted_signal: false, - rate_limit_exit: false, - rate_limit_reset_at: None, - }) + Ok(done_result(&ctx, &emit, total_usage)) } fn stop(&self) { - self.cancelled.store(true, Ordering::Relaxed); + self.cancelled.stop(); } fn get_status(&self) -> RuntimeStatus { - if self.cancelled.load(Ordering::Relaxed) { - RuntimeStatus::Failed - } else { - RuntimeStatus::Idle - } + self.cancelled.status() } } @@ -366,13 +277,8 @@ impl AgentRuntime for GeminiRuntime { #[cfg(test)] mod tests { + use super::super::api_common::test_runtime_context; use super::*; - use crate::http::context::AppContext; - - fn test_app_ctx() -> Arc { - let tmp = tempfile::tempdir().unwrap(); - Arc::new(AppContext::new_test(tmp.path().to_path_buf())) - } #[test] fn gemini_runtime_stop_sets_cancelled() { @@ -385,20 +291,7 @@ mod tests { #[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, - app_ctx: Some(test_app_ctx()), - session_id_to_resume: None, - fresh_prompt: None, - project_root: std::path::PathBuf::from("/tmp/project"), - model: None, - }; + let ctx = test_runtime_context("gemini-2.5-pro", vec![]); // The model extraction logic is inside start(), but we test the // condition here. diff --git a/server/src/agents/runtime/mod.rs b/server/src/agents/runtime/mod.rs index 6e9d89fb..a7cabf89 100644 --- a/server/src/agents/runtime/mod.rs +++ b/server/src/agents/runtime/mod.rs @@ -1,4 +1,5 @@ //! Agent runtimes — pluggable backends (Claude Code, Gemini, OpenAI) for running agents. +mod api_common; mod claude_code; mod gemini; mod openai; @@ -132,12 +133,7 @@ pub trait AgentRuntime: Send + Sync { #[cfg(test)] mod tests { use super::*; - use crate::http::context::AppContext; - - fn test_app_ctx() -> Arc { - let tmp = tempfile::tempdir().unwrap(); - Arc::new(AppContext::new_test(tmp.path().to_path_buf())) - } + use api_common::test_app_ctx; #[test] fn runtime_context_fields() { diff --git a/server/src/agents/runtime/openai.rs b/server/src/agents/runtime/openai.rs index f251b26b..f1d8225e 100644 --- a/server/src/agents/runtime/openai.rs +++ b/server/src/agents/runtime/openai.rs @@ -1,5 +1,5 @@ //! OpenAI Codex runtime — drives OpenAI API sessions as agent backends. -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::Ordering; use std::sync::{Arc, Mutex}; use reqwest::Client; @@ -12,6 +12,10 @@ use crate::http::mcp::tools_list::list_tools; use crate::slog; use super::super::{AgentEvent, TokenUsage}; +use super::api_common::{ + CancellationFlag, build_system_text, check_loop_guard, clean_schema_properties, done_result, + extract_model, start_conversation_loop, +}; use super::{AgentRuntime, RuntimeContext, RuntimeResult, RuntimeStatus}; // ── Public runtime struct ──────────────────────────────────────────── @@ -28,14 +32,14 @@ use super::{AgentRuntime, RuntimeContext, RuntimeResult, RuntimeStatus}; /// 6. Tracks token usage from the API response. pub struct OpenAiRuntime { /// Whether a stop has been requested. - cancelled: Arc, + cancelled: CancellationFlag, } impl OpenAiRuntime { /// Create a new OpenAI runtime instance. pub fn new() -> Self { Self { - cancelled: Arc::new(AtomicBool::new(false)), + cancelled: CancellationFlag::new(), } } } @@ -54,18 +58,11 @@ impl AgentRuntime for OpenAiRuntime { .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 - ctx.args - .iter() - .position(|a| a == "--model") - .and_then(|i| ctx.args.get(i + 1)) - .cloned() - .unwrap_or_else(|| "gpt-4o".to_string()) - }; + let model = extract_model( + &ctx, + |c| c.starts_with("gpt") || c.starts_with("o"), + "gpt-4o", + ); let app_ctx = ctx .app_ctx @@ -73,7 +70,7 @@ impl AgentRuntime for OpenAiRuntime { .ok_or_else(|| "OpenAI runtime requires app_ctx to be set".to_string())?; let client = Client::new(); - let cancelled = Arc::clone(&self.cancelled); + let cancelled = self.cancelled.handle(); // Step 1: Fetch MCP tool definitions and convert to OpenAI format. let openai_tools = convert_mcp_tools_to_openai(); @@ -85,65 +82,14 @@ impl AgentRuntime for OpenAiRuntime { 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(), - }); + let (emit, mut total_usage, mut turn) = + start_conversation_loop(&ctx, tx, event_log, log_writer); // 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), - exit_ok: true, - aborted_signal: false, - rate_limit_exit: false, - rate_limit_reset_at: None, - }); - } - - 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), - exit_ok: true, - aborted_signal: false, - rate_limit_exit: false, - rate_limit_reset_at: None, - }); + if let Some(result) = check_loop_guard(&ctx, &cancelled, &mut turn, &total_usage, &emit) + { + return Ok(result); } slog!( @@ -218,19 +164,7 @@ impl AgentRuntime for OpenAiRuntime { 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), - exit_ok: true, - aborted_signal: false, - rate_limit_exit: false, - rate_limit_reset_at: None, - }); + return Ok(done_result(&ctx, &emit, total_usage)); } let tool_calls = tool_calls.unwrap(); @@ -297,38 +231,16 @@ impl AgentRuntime for OpenAiRuntime { } fn stop(&self) { - self.cancelled.store(true, Ordering::Relaxed); + self.cancelled.stop(); } fn get_status(&self) -> RuntimeStatus { - if self.cancelled.load(Ordering::Relaxed) { - RuntimeStatus::Failed - } else { - RuntimeStatus::Idle - } + self.cancelled.status() } } // ── 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 - ) - }) -} - /// Load MCP tool definitions directly and convert to OpenAI function-calling format. fn convert_mcp_tools_to_openai() -> Vec { let tools = list_tools(); @@ -377,7 +289,7 @@ fn convert_mcp_schema_to_openai(schema: Option<&Value>) -> Option { }); if let Some(properties) = schema.get("properties") { - result["properties"] = clean_schema_properties(properties); + result["properties"] = clean_schema_properties(properties, false); } else { result["properties"] = json!({}); } @@ -392,40 +304,6 @@ fn convert_mcp_schema_to_openai(schema: Option<&Value>) -> Option { 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) -} - /// Parse token usage from an OpenAI API response. fn parse_usage(response: &Value) -> Option { let usage = response.get("usage")?; @@ -449,13 +327,8 @@ fn parse_usage(response: &Value) -> Option { #[cfg(test)] mod tests { + use super::super::api_common::test_runtime_context; use super::*; - use crate::http::context::AppContext; - - fn test_app_ctx() -> Arc { - let tmp = tempfile::tempdir().unwrap(); - Arc::new(AppContext::new_test(tmp.path().to_path_buf())) - } #[test] fn convert_mcp_schema_simple_object() { @@ -512,92 +385,6 @@ mod tests { 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, - app_ctx: Some(test_app_ctx()), - session_id_to_resume: None, - fresh_prompt: None, - project_root: std::path::PathBuf::from("/tmp/project"), - model: None, - }; - - 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, - app_ctx: Some(test_app_ctx()), - session_id_to_resume: None, - fresh_prompt: None, - project_root: std::path::PathBuf::from("/tmp/project"), - model: None, - }; - - 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!({ @@ -631,39 +418,13 @@ mod tests { #[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, - app_ctx: Some(test_app_ctx()), - session_id_to_resume: None, - fresh_prompt: None, - project_root: std::path::PathBuf::from("/tmp/project"), - model: None, - }; + let ctx = test_runtime_context("gpt-4o", vec![]); 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, - app_ctx: Some(test_app_ctx()), - session_id_to_resume: None, - fresh_prompt: None, - project_root: std::path::PathBuf::from("/tmp/project"), - model: None, - }; + let ctx = test_runtime_context("o3", vec![]); assert!(ctx.command.starts_with("o")); } }