optimclaw/tests/support/gateway_workflow_harness.rs

#![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<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();
        }
    }
}