431 lines
14 KiB
Rust
431 lines
14 KiB
Rust
//! OpenAI Codex runtime — drives OpenAI API sessions as agent backends.
|
|
use std::sync::atomic::Ordering;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
use reqwest::Client;
|
|
use serde_json::{Value, json};
|
|
use tokio::sync::broadcast;
|
|
|
|
use crate::agent_log::AgentLogWriter;
|
|
use crate::http::mcp::dispatch::dispatch_tool_call;
|
|
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 ────────────────────────────────────────────
|
|
|
|
/// 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 huskies' 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: CancellationFlag,
|
|
}
|
|
|
|
impl OpenAiRuntime {
|
|
/// Create a new OpenAI runtime instance.
|
|
pub fn new() -> Self {
|
|
Self {
|
|
cancelled: CancellationFlag::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
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 = extract_model(
|
|
&ctx,
|
|
|c| c.starts_with("gpt") || c.starts_with("o"),
|
|
"gpt-4o",
|
|
);
|
|
|
|
let app_ctx = ctx
|
|
.app_ctx
|
|
.clone()
|
|
.ok_or_else(|| "OpenAI runtime requires app_ctx to be set".to_string())?;
|
|
|
|
let client = Client::new();
|
|
let cancelled = self.cancelled.handle();
|
|
|
|
// Step 1: Fetch MCP tool definitions and convert to OpenAI format.
|
|
let openai_tools = convert_mcp_tools_to_openai();
|
|
|
|
// 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 (emit, mut total_usage, mut turn) =
|
|
start_conversation_loop(&ctx, tx, event_log, log_writer);
|
|
|
|
// Step 3: Conversation loop.
|
|
loop {
|
|
if let Some(result) = check_loop_guard(&ctx, &cancelled, &mut turn, &total_usage, &emit)
|
|
{
|
|
return Ok(result);
|
|
}
|
|
|
|
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.
|
|
return Ok(done_result(&ctx, &emit, 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 = dispatch_tool_call(tool_name, args.clone(), &app_ctx).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.stop();
|
|
}
|
|
|
|
fn get_status(&self) -> RuntimeStatus {
|
|
self.cancelled.status()
|
|
}
|
|
}
|
|
|
|
// ── Helper functions ─────────────────────────────────────────────────
|
|
|
|
/// Load MCP tool definitions directly and convert to OpenAI function-calling format.
|
|
fn convert_mcp_tools_to_openai() -> Vec<Value> {
|
|
let tools = list_tools();
|
|
|
|
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()
|
|
);
|
|
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, false);
|
|
} 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)
|
|
}
|
|
|
|
/// 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::super::api_common::test_runtime_context;
|
|
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 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 = test_runtime_context("gpt-4o", vec![]);
|
|
assert!(ctx.command.starts_with("gpt"));
|
|
}
|
|
|
|
#[test]
|
|
fn model_extraction_from_command_o3() {
|
|
let ctx = test_runtime_context("o3", vec![]);
|
|
assert!(ctx.command.starts_with("o"));
|
|
}
|
|
}
|