crashlabs/huskies
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 37 Wiki Activity

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

Browse Source
This commit is contained in:
Huskies Agent
2026-07-21 17:29:35 +00:00
parent 1583ade9fc
commit 10dd239c92
6 changed files with 461 additions and 531 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/src/agents/runtime/openai.rs
+25 -264
View File
@@ -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") {
// 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 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<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 {
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"));
}
}