#![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::{Config, RegistryProviderConfig, 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, } impl TestChannelHandle { fn new(inner: Arc) -> Self { Self { inner } } } #[async_trait] impl Channel for TestChannelHandle { fn name(&self) -> &str { "gateway" } async fn start(&self) -> Result { 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 { 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, pub db: Arc, gateway_state: Arc, agent_handle: Option>, bridge_handle: Option>, _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 = 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>>> = 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::(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::() .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::() .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::() .await .expect("invalid routines response") } pub async fn routine_by_name(&self, name: &str) -> Option { 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::() .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::() .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(); } } }