mirror of
https://github.com/outbackdingo/optimclaw.git
synced 2026-08-25 14:53:34 +00:00
Add reusable gateway workflow test harness and mock OpenAI server
This commit is contained in:
Illia Polosukhin
@@ -0,0 +1,130 @@
|
||||
//! Live-ish gateway workflow integration using an in-process mock OpenAI server.
|
||||
//! This exercises the same path as manual validation:
|
||||
//! - chat send through gateway
|
||||
//! - routine creation via tool call
|
||||
//! - system-event emission via tool call
|
||||
//! - webhook ingestion via gateway
|
||||
//! - status/runs checks via routines API
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
mod support;
|
||||
|
||||
#[cfg(feature = "libsql")]
|
||||
mod tests {
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::support::gateway_workflow_harness::GatewayWorkflowHarness;
|
||||
use crate::support::mock_openai_server::{
|
||||
MockOpenAiResponse, MockOpenAiRule, MockOpenAiServerBuilder, MockToolCall,
|
||||
};
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_workflow_harness_chat_and_webhook() {
|
||||
let mock = MockOpenAiServerBuilder::new()
|
||||
.with_rule(MockOpenAiRule::on_user_contains(
|
||||
"create workflow routine",
|
||||
MockOpenAiResponse::ToolCalls(vec![MockToolCall::new(
|
||||
"call_create_1",
|
||||
"routine_create",
|
||||
serde_json::json!({
|
||||
"name": "wf-ci-webhook-demo",
|
||||
"description": "CI webhook workflow demo",
|
||||
"trigger_type": "system_event",
|
||||
"trigger": {
|
||||
"source": "github",
|
||||
"event_type": "issue.opened",
|
||||
"filters": {"repository": "nearai/ironclaw"}
|
||||
},
|
||||
"action_type": "lightweight",
|
||||
"action": {
|
||||
"prompt": "Summarize webhook and report issue number"
|
||||
}
|
||||
}),
|
||||
)]),
|
||||
))
|
||||
.with_rule(MockOpenAiRule::on_user_contains(
|
||||
"emit webhook event",
|
||||
MockOpenAiResponse::ToolCalls(vec![MockToolCall::new(
|
||||
"call_emit_1",
|
||||
"event_emit",
|
||||
serde_json::json!({
|
||||
"source": "github",
|
||||
"event_type": "issue.opened",
|
||||
"payload": {
|
||||
"repository": "nearai/ironclaw",
|
||||
"issue": {"number": 777, "title": "Infra test"}
|
||||
}
|
||||
}),
|
||||
)]),
|
||||
))
|
||||
.with_default_response(MockOpenAiResponse::Text("ack".to_string()))
|
||||
.start()
|
||||
.await;
|
||||
|
||||
let harness =
|
||||
GatewayWorkflowHarness::start_openai_compatible(&mock.openai_base_url(), "mock-model")
|
||||
.await;
|
||||
|
||||
let thread_id = harness.create_thread().await;
|
||||
harness
|
||||
.send_chat(&thread_id, "create workflow routine")
|
||||
.await;
|
||||
harness.send_chat(&thread_id, "emit webhook event").await;
|
||||
|
||||
let history = harness
|
||||
.wait_for_turns(&thread_id, 2, Duration::from_secs(10))
|
||||
.await;
|
||||
let turns = history["turns"].as_array().expect("turns array missing");
|
||||
assert!(turns.len() >= 2, "expected at least 2 turns");
|
||||
|
||||
let routine = harness
|
||||
.routine_by_name("wf-ci-webhook-demo")
|
||||
.await
|
||||
.expect("routine not created");
|
||||
let routine_id = routine["id"].as_str().expect("routine id missing");
|
||||
|
||||
let runs_before = harness.routine_runs(routine_id).await;
|
||||
let before_count = runs_before["runs"]
|
||||
.as_array()
|
||||
.map(|a| a.len())
|
||||
.unwrap_or_default();
|
||||
|
||||
let hook = harness
|
||||
.github_webhook(
|
||||
"issues",
|
||||
serde_json::json!({
|
||||
"action": "opened",
|
||||
"repository": {"full_name": "nearai/ironclaw"},
|
||||
"issue": {"number": 778, "title": "Webhook endpoint test"}
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(hook["status"], "accepted");
|
||||
assert_eq!(hook["event_type"], "issue.opened");
|
||||
assert!(
|
||||
hook["fired_routines"].as_u64().unwrap_or(0) >= 1,
|
||||
"expected webhook to fire at least one routine"
|
||||
);
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
let runs_after = harness.routine_runs(routine_id).await;
|
||||
let after_count = runs_after["runs"]
|
||||
.as_array()
|
||||
.map(|a| a.len())
|
||||
.unwrap_or_default();
|
||||
assert!(
|
||||
after_count > before_count,
|
||||
"expected routine runs to increase after webhook; before={before_count}, after={after_count}"
|
||||
);
|
||||
|
||||
let requests = mock.requests().await;
|
||||
assert!(
|
||||
requests.len() >= 2,
|
||||
"expected mock LLM server to receive requests"
|
||||
);
|
||||
|
||||
harness.shutdown().await;
|
||||
mock.shutdown().await;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,467 @@
|
||||
#![allow(dead_code)]
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use async_trait::async_trait;
|
||||
use secrecy::SecretString;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
use ironclaw::agent::routine_engine::RoutineEngine;
|
||||
use ironclaw::agent::{Agent, AgentDeps, SessionManager as AgentSessionManager};
|
||||
use ironclaw::app::{AppBuilder, AppBuilderFlags};
|
||||
use ironclaw::channels::web::log_layer::LogBroadcaster;
|
||||
use ironclaw::channels::web::server::{GatewayState, RateLimiter, start_server};
|
||||
use ironclaw::channels::web::sse::SseManager;
|
||||
use ironclaw::channels::web::ws::WsConnectionTracker;
|
||||
use ironclaw::channels::{Channel, IncomingMessage, MessageStream};
|
||||
use ironclaw::config::llm::RegistryProviderConfig;
|
||||
use ironclaw::config::{Config, RoutineConfig};
|
||||
use ironclaw::context::ContextManager;
|
||||
use ironclaw::db::Database;
|
||||
use ironclaw::db::libsql::LibSqlBackend;
|
||||
use ironclaw::error::ChannelError;
|
||||
use ironclaw::llm::registry::ProviderProtocol;
|
||||
use ironclaw::llm::{
|
||||
SessionConfig as LlmSessionConfig, SessionManager as LlmSessionManager, create_llm_provider,
|
||||
};
|
||||
|
||||
use crate::support::test_channel::TestChannel;
|
||||
|
||||
struct TestChannelHandle {
|
||||
inner: Arc<TestChannel>,
|
||||
}
|
||||
|
||||
impl TestChannelHandle {
|
||||
fn new(inner: Arc<TestChannel>) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl Channel for TestChannelHandle {
|
||||
fn name(&self) -> &str {
|
||||
"gateway"
|
||||
}
|
||||
|
||||
async fn start(&self) -> Result<MessageStream, ChannelError> {
|
||||
self.inner.start().await
|
||||
}
|
||||
|
||||
async fn respond(
|
||||
&self,
|
||||
msg: &IncomingMessage,
|
||||
response: ironclaw::channels::OutgoingResponse,
|
||||
) -> Result<(), ChannelError> {
|
||||
self.inner.respond(msg, response).await
|
||||
}
|
||||
|
||||
async fn send_status(
|
||||
&self,
|
||||
status: ironclaw::channels::StatusUpdate,
|
||||
metadata: &serde_json::Value,
|
||||
) -> Result<(), ChannelError> {
|
||||
self.inner.send_status(status, metadata).await
|
||||
}
|
||||
|
||||
async fn broadcast(
|
||||
&self,
|
||||
user_id: &str,
|
||||
response: ironclaw::channels::OutgoingResponse,
|
||||
) -> Result<(), ChannelError> {
|
||||
self.inner.broadcast(user_id, response).await
|
||||
}
|
||||
|
||||
async fn health_check(&self) -> Result<(), ChannelError> {
|
||||
self.inner.health_check().await
|
||||
}
|
||||
|
||||
fn conversation_context(&self, metadata: &serde_json::Value) -> HashMap<String, String> {
|
||||
self.inner.conversation_context(metadata)
|
||||
}
|
||||
|
||||
async fn shutdown(&self) -> Result<(), ChannelError> {
|
||||
self.inner.shutdown().await
|
||||
}
|
||||
}
|
||||
|
||||
pub struct GatewayWorkflowHarness {
|
||||
pub addr: SocketAddr,
|
||||
pub auth_token: String,
|
||||
pub client: reqwest::Client,
|
||||
pub user_id: String,
|
||||
pub test_channel: Arc<TestChannel>,
|
||||
pub db: Arc<dyn Database>,
|
||||
gateway_state: Arc<GatewayState>,
|
||||
agent_handle: Option<tokio::task::JoinHandle<()>>,
|
||||
bridge_handle: Option<tokio::task::JoinHandle<()>>,
|
||||
_temp_dir: tempfile::TempDir,
|
||||
}
|
||||
|
||||
impl GatewayWorkflowHarness {
|
||||
pub async fn start_openai_compatible(base_url: &str, model: &str) -> Self {
|
||||
let temp_dir = tempfile::tempdir().expect("failed to create temp dir");
|
||||
let db_path = temp_dir.path().join("gateway_workflow_harness.db");
|
||||
let backend = LibSqlBackend::new_local(&db_path)
|
||||
.await
|
||||
.expect("failed to create test db");
|
||||
backend
|
||||
.run_migrations()
|
||||
.await
|
||||
.expect("failed to run migrations");
|
||||
let db: Arc<dyn Database> = Arc::new(backend);
|
||||
|
||||
let skills_dir = temp_dir.path().join("skills");
|
||||
let installed_skills_dir = temp_dir.path().join("installed_skills");
|
||||
let _ = std::fs::create_dir_all(&skills_dir);
|
||||
let _ = std::fs::create_dir_all(&installed_skills_dir);
|
||||
let mut config = Config::for_testing(db_path, skills_dir, installed_skills_dir);
|
||||
config.agent.auto_approve_tools = true;
|
||||
config.routines.enabled = true;
|
||||
config.routines.max_concurrent_routines = 4;
|
||||
config.llm.backend = "openai_compatible".to_string();
|
||||
config.llm.provider = Some(RegistryProviderConfig {
|
||||
protocol: ProviderProtocol::OpenAiCompletions,
|
||||
provider_id: "openai_compatible".to_string(),
|
||||
api_key: Some(SecretString::from("dummy".to_string())),
|
||||
base_url: base_url.to_string(),
|
||||
model: model.to_string(),
|
||||
extra_headers: Vec::new(),
|
||||
oauth_token: None,
|
||||
});
|
||||
|
||||
let llm_session = Arc::new(LlmSessionManager::new(LlmSessionConfig::default()));
|
||||
let llm = create_llm_provider(&config.llm, Arc::clone(&llm_session))
|
||||
.expect("failed to create openai-compatible provider");
|
||||
|
||||
let log_broadcaster = Arc::new(LogBroadcaster::new());
|
||||
let mut app_builder = AppBuilder::new(
|
||||
config,
|
||||
AppBuilderFlags::default(),
|
||||
None,
|
||||
Arc::clone(&llm_session),
|
||||
log_broadcaster,
|
||||
);
|
||||
app_builder.with_database(Arc::clone(&db));
|
||||
app_builder.with_llm(llm);
|
||||
|
||||
let mut components = app_builder
|
||||
.build_all()
|
||||
.await
|
||||
.expect("failed to build app components");
|
||||
|
||||
let ctx_mgr = Arc::new(ContextManager::new(
|
||||
components.config.agent.max_parallel_jobs,
|
||||
));
|
||||
components.tools.register_job_tools(
|
||||
ctx_mgr,
|
||||
None,
|
||||
None,
|
||||
components.db.clone(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
let routine_slot: Arc<tokio::sync::RwLock<Option<Arc<RoutineEngine>>>> =
|
||||
Arc::new(tokio::sync::RwLock::new(None));
|
||||
|
||||
if let (Some(db_arc), Some(ws)) = (&components.db, &components.workspace) {
|
||||
let (notify_tx, _notify_rx) = tokio::sync::mpsc::channel(16);
|
||||
let engine = Arc::new(RoutineEngine::new(
|
||||
RoutineConfig::default(),
|
||||
Arc::clone(db_arc),
|
||||
components.llm.clone(),
|
||||
Arc::clone(ws),
|
||||
notify_tx,
|
||||
None,
|
||||
));
|
||||
components
|
||||
.tools
|
||||
.register_routine_tools(Arc::clone(db_arc), engine);
|
||||
}
|
||||
|
||||
let test_channel = Arc::new(TestChannel::new());
|
||||
let handle = TestChannelHandle::new(Arc::clone(&test_channel));
|
||||
let channel_manager = ironclaw::channels::ChannelManager::new();
|
||||
channel_manager.add(Box::new(handle)).await;
|
||||
let channels = Arc::new(channel_manager);
|
||||
|
||||
let user_id = "gateway-test-user".to_string();
|
||||
let (gw_tx, mut gw_rx) = mpsc::channel::<IncomingMessage>(256);
|
||||
let forward_channel = Arc::clone(&test_channel);
|
||||
let bridge_handle = tokio::spawn(async move {
|
||||
while let Some(msg) = gw_rx.recv().await {
|
||||
forward_channel.send_incoming(msg).await;
|
||||
}
|
||||
});
|
||||
|
||||
let scheduler_slot: ironclaw::tools::builtin::SchedulerSlot =
|
||||
Arc::new(tokio::sync::RwLock::new(None));
|
||||
let agent_session_manager = Arc::new(AgentSessionManager::new());
|
||||
|
||||
let gateway_state = Arc::new(GatewayState {
|
||||
msg_tx: tokio::sync::RwLock::new(Some(gw_tx)),
|
||||
sse: SseManager::new(),
|
||||
workspace: components.workspace.clone(),
|
||||
session_manager: Some(Arc::clone(&agent_session_manager)),
|
||||
log_broadcaster: None,
|
||||
log_level_handle: None,
|
||||
extension_manager: components.extension_manager.clone(),
|
||||
tool_registry: Some(Arc::clone(&components.tools)),
|
||||
store: components.db.clone(),
|
||||
job_manager: None,
|
||||
prompt_queue: None,
|
||||
scheduler: Some(scheduler_slot.clone()),
|
||||
user_id: user_id.clone(),
|
||||
shutdown_tx: tokio::sync::RwLock::new(None),
|
||||
ws_tracker: Some(Arc::new(WsConnectionTracker::new())),
|
||||
llm_provider: Some(Arc::clone(&components.llm)),
|
||||
skill_registry: components.skill_registry.clone(),
|
||||
skill_catalog: components.skill_catalog.clone(),
|
||||
chat_rate_limiter: RateLimiter::new(120, 60),
|
||||
registry_entries: Vec::new(),
|
||||
cost_guard: Some(Arc::clone(&components.cost_guard)),
|
||||
routine_engine: Arc::clone(&routine_slot),
|
||||
startup_time: Instant::now(),
|
||||
});
|
||||
|
||||
let mut agent = Agent::new(
|
||||
components.config.agent.clone(),
|
||||
AgentDeps {
|
||||
store: components.db,
|
||||
llm: components.llm,
|
||||
cheap_llm: components.cheap_llm,
|
||||
safety: components.safety,
|
||||
tools: components.tools,
|
||||
workspace: components.workspace,
|
||||
extension_manager: components.extension_manager,
|
||||
skill_registry: components.skill_registry,
|
||||
skill_catalog: components.skill_catalog,
|
||||
skills_config: components.config.skills.clone(),
|
||||
hooks: components.hooks,
|
||||
cost_guard: components.cost_guard,
|
||||
sse_tx: Some(gateway_state.sse.sender()),
|
||||
http_interceptor: None,
|
||||
transcription: None,
|
||||
document_extraction: None,
|
||||
},
|
||||
channels,
|
||||
None,
|
||||
None,
|
||||
Some(RoutineConfig {
|
||||
enabled: true,
|
||||
cron_check_interval_secs: 60,
|
||||
max_concurrent_routines: 4,
|
||||
default_cooldown_secs: 300,
|
||||
max_lightweight_tokens: 4096,
|
||||
}),
|
||||
Some(Arc::clone(&components.context_manager)),
|
||||
Some(Arc::clone(&agent_session_manager)),
|
||||
);
|
||||
agent.set_routine_engine_slot(Arc::clone(&routine_slot));
|
||||
*scheduler_slot.write().await = Some(agent.scheduler());
|
||||
|
||||
let agent_handle = tokio::spawn(async move {
|
||||
let _ = agent.run().await;
|
||||
});
|
||||
|
||||
if let Some(rx) = test_channel.take_ready_rx().await {
|
||||
let _ = tokio::time::timeout(Duration::from_secs(5), rx).await;
|
||||
}
|
||||
|
||||
let auth_token = "gateway-test-token".to_string();
|
||||
let addr = start_server(
|
||||
"127.0.0.1:0".parse().expect("valid localhost addr"),
|
||||
Arc::clone(&gateway_state),
|
||||
auth_token.clone(),
|
||||
)
|
||||
.await
|
||||
.expect("failed to start gateway server");
|
||||
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(10))
|
||||
.build()
|
||||
.expect("failed to build reqwest client");
|
||||
|
||||
Self {
|
||||
addr,
|
||||
auth_token,
|
||||
client,
|
||||
user_id,
|
||||
test_channel,
|
||||
db,
|
||||
gateway_state,
|
||||
agent_handle: Some(agent_handle),
|
||||
bridge_handle: Some(bridge_handle),
|
||||
_temp_dir: temp_dir,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn base_url(&self) -> String {
|
||||
format!("http://{}", self.addr)
|
||||
}
|
||||
|
||||
pub async fn create_thread(&self) -> String {
|
||||
let resp = self
|
||||
.client
|
||||
.post(format!("{}/api/chat/thread/new", self.base_url()))
|
||||
.bearer_auth(&self.auth_token)
|
||||
.send()
|
||||
.await
|
||||
.expect("create thread request failed")
|
||||
.error_for_status()
|
||||
.expect("create thread non-2xx")
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("invalid thread response");
|
||||
resp.get("id")
|
||||
.and_then(|v| v.as_str())
|
||||
.expect("thread id missing")
|
||||
.to_string()
|
||||
}
|
||||
|
||||
pub async fn send_chat(&self, thread_id: &str, content: &str) {
|
||||
let _ = self
|
||||
.client
|
||||
.post(format!("{}/api/chat/send", self.base_url()))
|
||||
.bearer_auth(&self.auth_token)
|
||||
.json(&serde_json::json!({"thread_id": thread_id, "content": content}))
|
||||
.send()
|
||||
.await
|
||||
.expect("chat send failed")
|
||||
.error_for_status()
|
||||
.expect("chat send non-2xx");
|
||||
}
|
||||
|
||||
pub async fn history(&self, thread_id: &str) -> serde_json::Value {
|
||||
self.client
|
||||
.get(format!(
|
||||
"{}/api/chat/history?thread_id={thread_id}",
|
||||
self.base_url()
|
||||
))
|
||||
.bearer_auth(&self.auth_token)
|
||||
.send()
|
||||
.await
|
||||
.expect("history request failed")
|
||||
.error_for_status()
|
||||
.expect("history non-2xx")
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("invalid history response")
|
||||
}
|
||||
|
||||
pub async fn wait_for_turns(
|
||||
&self,
|
||||
thread_id: &str,
|
||||
min_turns: usize,
|
||||
timeout: Duration,
|
||||
) -> serde_json::Value {
|
||||
let deadline = Instant::now() + timeout;
|
||||
loop {
|
||||
let history = self.history(thread_id).await;
|
||||
let turns = history
|
||||
.get("turns")
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|v| v.len())
|
||||
.unwrap_or_default();
|
||||
if turns >= min_turns {
|
||||
return history;
|
||||
}
|
||||
assert!(Instant::now() < deadline, "timed out waiting for turns");
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_routines(&self) -> serde_json::Value {
|
||||
self.client
|
||||
.get(format!("{}/api/routines", self.base_url()))
|
||||
.bearer_auth(&self.auth_token)
|
||||
.send()
|
||||
.await
|
||||
.expect("routines request failed")
|
||||
.error_for_status()
|
||||
.expect("routines non-2xx")
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("invalid routines response")
|
||||
}
|
||||
|
||||
pub async fn routine_by_name(&self, name: &str) -> Option<serde_json::Value> {
|
||||
let routines = self.list_routines().await;
|
||||
routines
|
||||
.get("routines")
|
||||
.and_then(|v| v.as_array())
|
||||
.and_then(|arr| {
|
||||
arr.iter()
|
||||
.find(|r| r.get("name").and_then(|v| v.as_str()) == Some(name))
|
||||
.cloned()
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn routine_runs(&self, routine_id: &str) -> serde_json::Value {
|
||||
self.client
|
||||
.get(format!(
|
||||
"{}/api/routines/{routine_id}/runs",
|
||||
self.base_url()
|
||||
))
|
||||
.bearer_auth(&self.auth_token)
|
||||
.send()
|
||||
.await
|
||||
.expect("routine runs request failed")
|
||||
.error_for_status()
|
||||
.expect("routine runs non-2xx")
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("invalid routine runs response")
|
||||
}
|
||||
|
||||
pub async fn github_webhook(
|
||||
&self,
|
||||
event: &str,
|
||||
payload: serde_json::Value,
|
||||
) -> serde_json::Value {
|
||||
self.client
|
||||
.post(format!("{}/api/webhooks/github", self.base_url()))
|
||||
.header("x-github-event", event)
|
||||
.json(&payload)
|
||||
.send()
|
||||
.await
|
||||
.expect("webhook request failed")
|
||||
.error_for_status()
|
||||
.expect("webhook non-2xx")
|
||||
.json::<serde_json::Value>()
|
||||
.await
|
||||
.expect("invalid webhook response")
|
||||
}
|
||||
|
||||
pub async fn shutdown(mut self) {
|
||||
self.test_channel.signal_shutdown();
|
||||
|
||||
if let Some(tx) = self.gateway_state.shutdown_tx.write().await.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
|
||||
if let Some(handle) = self.bridge_handle.take() {
|
||||
handle.abort();
|
||||
}
|
||||
if let Some(handle) = self.agent_handle.take() {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for GatewayWorkflowHarness {
|
||||
fn drop(&mut self) {
|
||||
self.test_channel.signal_shutdown();
|
||||
if let Some(handle) = self.bridge_handle.take() {
|
||||
handle.abort();
|
||||
}
|
||||
if let Some(handle) = self.agent_handle.take() {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,277 @@
|
||||
#![allow(dead_code)]
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use axum::extract::State;
|
||||
use axum::http::StatusCode;
|
||||
use axum::routing::{get, post};
|
||||
use axum::{Json, Router};
|
||||
use serde_json::{Value, json};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::{Mutex, oneshot};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MockOpenAiRule {
|
||||
contains: String,
|
||||
response: MockOpenAiResponse,
|
||||
}
|
||||
|
||||
impl MockOpenAiRule {
|
||||
pub fn on_user_contains(contains: impl Into<String>, response: MockOpenAiResponse) -> Self {
|
||||
Self {
|
||||
contains: contains.into(),
|
||||
response,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub enum MockOpenAiResponse {
|
||||
Text(String),
|
||||
ToolCalls(Vec<MockToolCall>),
|
||||
Raw(Value),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct MockToolCall {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub arguments: Value,
|
||||
}
|
||||
|
||||
impl MockToolCall {
|
||||
pub fn new(id: impl Into<String>, name: impl Into<String>, arguments: Value) -> Self {
|
||||
Self {
|
||||
id: id.into(),
|
||||
name: name.into(),
|
||||
arguments,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct MockOpenAiServerBuilder {
|
||||
models: Vec<String>,
|
||||
rules: Vec<MockOpenAiRule>,
|
||||
default_response: Option<MockOpenAiResponse>,
|
||||
}
|
||||
|
||||
impl MockOpenAiServerBuilder {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
models: vec!["mock-model".to_string()],
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_models(mut self, models: Vec<String>) -> Self {
|
||||
self.models = models;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_rule(mut self, rule: MockOpenAiRule) -> Self {
|
||||
self.rules.push(rule);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_default_response(mut self, response: MockOpenAiResponse) -> Self {
|
||||
self.default_response = Some(response);
|
||||
self
|
||||
}
|
||||
|
||||
pub async fn start(self) -> MockOpenAiServer {
|
||||
let state = Arc::new(MockOpenAiState {
|
||||
models: self.models,
|
||||
rules: self.rules,
|
||||
default_response: self
|
||||
.default_response
|
||||
.unwrap_or_else(|| MockOpenAiResponse::Text("OK".to_string())),
|
||||
requests: Mutex::new(Vec::new()),
|
||||
response_counter: AtomicU64::new(1),
|
||||
});
|
||||
|
||||
let app = Router::new()
|
||||
.route("/v1/models", get(models_handler))
|
||||
.route("/v1/chat/completions", post(chat_completions_handler))
|
||||
.with_state(Arc::clone(&state));
|
||||
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("failed to bind mock openai server");
|
||||
let addr = listener.local_addr().expect("failed to read bound addr");
|
||||
|
||||
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
|
||||
let handle = tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, app)
|
||||
.with_graceful_shutdown(async {
|
||||
let _ = shutdown_rx.await;
|
||||
})
|
||||
.await;
|
||||
});
|
||||
|
||||
MockOpenAiServer {
|
||||
addr,
|
||||
state,
|
||||
shutdown_tx: Some(shutdown_tx),
|
||||
server_task: Some(handle),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct MockOpenAiServer {
|
||||
addr: SocketAddr,
|
||||
state: Arc<MockOpenAiState>,
|
||||
shutdown_tx: Option<oneshot::Sender<()>>,
|
||||
server_task: Option<tokio::task::JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl MockOpenAiServer {
|
||||
pub fn base_url(&self) -> String {
|
||||
format!("http://{}", self.addr)
|
||||
}
|
||||
|
||||
pub fn openai_base_url(&self) -> String {
|
||||
format!("{}/v1", self.base_url())
|
||||
}
|
||||
|
||||
pub async fn requests(&self) -> Vec<Value> {
|
||||
self.state.requests.lock().await.clone()
|
||||
}
|
||||
|
||||
pub async fn shutdown(mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(handle) = self.server_task.take() {
|
||||
let _ = handle.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for MockOpenAiServer {
|
||||
fn drop(&mut self) {
|
||||
if let Some(tx) = self.shutdown_tx.take() {
|
||||
let _ = tx.send(());
|
||||
}
|
||||
if let Some(handle) = self.server_task.take() {
|
||||
handle.abort();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct MockOpenAiState {
|
||||
models: Vec<String>,
|
||||
rules: Vec<MockOpenAiRule>,
|
||||
default_response: MockOpenAiResponse,
|
||||
requests: Mutex<Vec<Value>>,
|
||||
response_counter: AtomicU64,
|
||||
}
|
||||
|
||||
async fn models_handler(State(state): State<Arc<MockOpenAiState>>) -> Json<Value> {
|
||||
Json(json!({
|
||||
"object": "list",
|
||||
"data": state
|
||||
.models
|
||||
.iter()
|
||||
.map(|id| json!({"id": id, "object": "model"}))
|
||||
.collect::<Vec<_>>()
|
||||
}))
|
||||
}
|
||||
|
||||
async fn chat_completions_handler(
|
||||
State(state): State<Arc<MockOpenAiState>>,
|
||||
Json(body): Json<Value>,
|
||||
) -> Result<Json<Value>, (StatusCode, String)> {
|
||||
state.requests.lock().await.push(body.clone());
|
||||
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.unwrap_or("mock-model");
|
||||
let last_role = body
|
||||
.pointer("/messages")
|
||||
.and_then(|m| m.as_array())
|
||||
.and_then(|arr| arr.last())
|
||||
.and_then(|v| v.get("role"))
|
||||
.and_then(|r| r.as_str())
|
||||
.unwrap_or_default();
|
||||
|
||||
let latest_user = body
|
||||
.pointer("/messages")
|
||||
.and_then(|m| m.as_array())
|
||||
.and_then(|arr| {
|
||||
arr.iter().rev().find_map(|msg| {
|
||||
if msg.get("role").and_then(|r| r.as_str()) == Some("user") {
|
||||
msg.get("content")
|
||||
.and_then(|c| c.as_str())
|
||||
.map(ToOwned::to_owned)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
let selected = if last_role == "user" {
|
||||
state
|
||||
.rules
|
||||
.iter()
|
||||
.find(|r| latest_user.contains(&r.contains))
|
||||
.map(|r| r.response.clone())
|
||||
.unwrap_or_else(|| state.default_response.clone())
|
||||
} else {
|
||||
state.default_response.clone()
|
||||
};
|
||||
|
||||
let n = state.response_counter.fetch_add(1, Ordering::Relaxed);
|
||||
let response = match selected {
|
||||
MockOpenAiResponse::Text(content) => json!({
|
||||
"id": format!("chatcmpl-mock-{n}"),
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": content},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
|
||||
}),
|
||||
MockOpenAiResponse::ToolCalls(tool_calls) => {
|
||||
let calls = tool_calls
|
||||
.iter()
|
||||
.map(|tc| {
|
||||
json!({
|
||||
"id": tc.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.name,
|
||||
"arguments": tc.arguments.to_string()
|
||||
}
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
json!({
|
||||
"id": format!("chatcmpl-mock-{n}"),
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": serde_json::Value::Null,
|
||||
"tool_calls": calls
|
||||
},
|
||||
"finish_reason": "tool_calls"
|
||||
}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
|
||||
})
|
||||
}
|
||||
MockOpenAiResponse::Raw(v) => v,
|
||||
};
|
||||
|
||||
Ok(Json(response))
|
||||
}
|
||||
@@ -1,7 +1,9 @@
|
||||
pub mod assertions;
|
||||
pub mod cleanup;
|
||||
pub mod gateway_workflow_harness;
|
||||
pub mod instrumented_llm;
|
||||
pub mod metrics;
|
||||
pub mod mock_openai_server;
|
||||
pub mod test_channel;
|
||||
pub mod test_rig;
|
||||
pub mod trace_llm;
|
||||
|
||||
Reference in New Issue
Block a user