huskies: merge 1246 refactor Extract shared LLM-runtime logic duplicated between gemini and openai runtimes

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<AtomicBool>,
+}
+
+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<AtomicBool> {
+        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<AgentEvent>,
+    event_log: Arc<Mutex<Vec<AgentEvent>>>,
+    log_writer: Option<Arc<Mutex<AgentLogWriter>>>,
+) -> 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<AgentEvent>,
+    event_log: Arc<Mutex<Vec<AgentEvent>>>,
+    log_writer: Option<Arc<Mutex<AgentLogWriter>>>,
+) -> (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<RuntimeResult> {
+    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<crate::http::context::AppContext> {
+    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<String>) -> 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"));
+    }
+}

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<TokenUsage> {
 
 #[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<AppContext> {
-        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 {
+        let ctx = test_runtime_context(
-            story_id: "42_story_test".to_string(),
+            "gemini-2.5-pro",
-            agent_name: "coder-1".to_string(),
+            vec![
-            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,
-            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 {
+        let ctx = test_runtime_context("gemini-2.5-pro", vec![]);
-            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 instruction = build_system_instruction(&ctx);
         let text = instruction["parts"][0]["text"].as_str().unwrap();

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<Val
 
     let mut result = json!({
         "type": "object",
-        "properties": clean_schema_properties(properties),
+        "properties": clean_schema_properties(properties, true),
     });
 
     // Preserve required fields if present.
@@ -73,44 +74,6 @@ pub(super) fn convert_mcp_schema_to_gemini(schema: Option<&Value>) -> Option<Val
     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)
-}
-
 // ── 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());
-    }
 }

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<AtomicBool>,
+    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") {
+        let model = extract_model(&ctx, |c| c.starts_with("gemini"), "gemini-2.5-pro");
-            // 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 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 {
+        let (emit, mut total_usage, mut turn) =
-            input_tokens: 0,
+            start_conversation_loop(&ctx, tx, event_log, log_writer);
-            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) {
+            if let Some(result) = check_loop_guard(&ctx, &cancelled, &mut turn, &total_usage, &emit)
-                emit(AgentEvent::Error {
+            {
-                    story_id: ctx.story_id.clone(),
+                return Ok(result);
-                    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,
-                });
             }
 
             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 {
+                return Ok(done_result(&ctx, &emit, total_usage));
-                    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,
-                });
             }
 
             // Add the model's response to the conversation.
@@ -333,32 +261,15 @@ impl AgentRuntime for GeminiRuntime {
             }
         }
 
-        emit(AgentEvent::Done {
+        Ok(done_result(&ctx, &emit, total_usage))
-            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,
-        })
     }
 
     fn stop(&self) {
-        self.cancelled.store(true, Ordering::Relaxed);
+        self.cancelled.stop();
     }
 
     fn get_status(&self) -> RuntimeStatus {
-        if self.cancelled.load(Ordering::Relaxed) {
+        self.cancelled.status()
-            RuntimeStatus::Failed
-        } else {
-            RuntimeStatus::Idle
-        }
     }
 }
 
@@ -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<AppContext> {
-        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 {
+        let ctx = test_runtime_context("gemini-2.5-pro", vec![]);
-            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,
-        };
 
         // The model extraction logic is inside start(), but we test the
         // condition here.

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;
+    use api_common::test_app_ctx;
-
-    fn test_app_ctx() -> Arc<AppContext> {
-        let tmp = tempfile::tempdir().unwrap();
-        Arc::new(AppContext::new_test(tmp.path().to_path_buf()))
-    }
 
     #[test]
     fn runtime_context_fields() {

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<AtomicBool>,
+    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") {
+        let model = extract_model(
-            // The pool puts the model into `command` for non-CLI runtimes.
+            &ctx,
-            ctx.command.clone()
+            |c| c.starts_with("gpt") || c.starts_with("o"),
-        } else {
+            "gpt-4o",
-            // 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 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 {
+        let (emit, mut total_usage, mut turn) =
-            input_tokens: 0,
+            start_conversation_loop(&ctx, tx, event_log, log_writer);
-            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) {
+            if let Some(result) = check_loop_guard(&ctx, &cancelled, &mut turn, &total_usage, &emit)
-                emit(AgentEvent::Error {
+            {
-                    story_id: ctx.story_id.clone(),
+                return Ok(result);
-                    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,
-                });
             }
 
             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 {
+                return Ok(done_result(&ctx, &emit, total_usage));
-                    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,
-                });
             }
 
             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) {
+        self.cancelled.status()
-            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
-            )
-        })
-}
-
 /// Load MCP tool definitions directly and convert to OpenAI function-calling format.
 fn convert_mcp_tools_to_openai() -> Vec<Value> {
     let tools = list_tools();
@@ -377,7 +289,7 @@ fn convert_mcp_schema_to_openai(schema: Option<&Value>) -> Option<Value> {
     });
 
     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<Value> {
     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<TokenUsage> {
     let usage = response.get("usage")?;
@@ -449,13 +327,8 @@ fn parse_usage(response: &Value) -> Option<TokenUsage> {
 
 #[cfg(test)]
 mod tests {
+    use super::super::api_common::test_runtime_context;
     use super::*;
-    use crate::http::context::AppContext;
-
-    fn test_app_ctx() -> Arc<AppContext> {
-        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 {
+        let ctx = test_runtime_context("gpt-4o", vec![]);
-            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,
-        };
         assert!(ctx.command.starts_with("gpt"));
     }
 
     #[test]
     fn model_extraction_from_command_o3() {
-        let ctx = RuntimeContext {
+        let ctx = test_runtime_context("o3", vec![]);
-            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,
-        };
         assert!(ctx.command.starts_with("o"));
     }
 }