clawmates/crates/cm-api/tests/gateway_load.rs

//! Gateway under load (plan P6): many concurrent SSE streams against one
//! server — every run completes, every journal is strictly monotonic, and
//! every replay is byte-equal to its live stream. The §13 contract must
//! not degrade under concurrency.

use std::sync::Arc;

use cm_api::AppState;
use cm_auth::AuthService;
use cm_domain::{Role, User, UserId, Workspace, WorkspaceId};
use cm_llm::ScriptedProvider;
use cm_runtime::{Runtime, RuntimeConfig};
use eventsource_stream::Eventsource;
use futures::StreamExt;
use serde_json::{json, Value};

const STREAMS: usize = 40;

const SCENARIOS: &str = r#"
[[scenario]]
marker = "[[scenario:tool-time]]"

[[scenario.turns]]
events = [
    { type = "text", text = "Checking. " },
    { type = "tool_use", name = "clock.now", input = {} },
]

[[scenario.turns]]
events = [
    { type = "text", text = "I checked the current time for you." },
]
"#;

async fn collect_sse(response: reqwest::Response) -> Vec<(String, String, Value)> {
    let mut events = Vec::new();
    let mut stream = response.bytes_stream().eventsource();
    while let Some(event) = stream.next().await {
        let event = event.unwrap();
        let data: Value = serde_json::from_str(&event.data).unwrap();
        let name = event.event.clone();
        let done = name == "run_completed" || name == "error";
        events.push((event.id, name, data));
        if done {
            break;
        }
    }
    events
}

#[tokio::test]
async fn forty_concurrent_streams_complete_and_replay_identically() {
    let pool = cm_testkit::test_pool().await;
    let runtime = Runtime::new(
        pool.clone(),
        Arc::new(ScriptedProvider::from_toml(SCENARIOS).unwrap()),
        RuntimeConfig::basic("scripted", 1024),
    );
    let app = cm_api::router(AppState::new(pool.clone(), runtime));
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let base = format!("http://{}", listener.local_addr().unwrap());
    tokio::spawn(async move {
        axum::serve(listener, app).await.unwrap();
    });

    let ws = Workspace {
        id: WorkspaceId::new(),
        name: "Acme".into(),
        plan: "team".into(),
    };
    cm_db::repo::workspaces::insert(&pool, &ws).await.unwrap();
    let owner = User {
        id: UserId::new(),
        workspace_id: ws.id,
        email: format!("{}@acme.test", UserId::new()),
        role: Role::Owner,
        display_name: "Owner".into(),
        created_at: time::OffsetDateTime::UNIX_EPOCH,
    };
    cm_db::repo::users::insert(&pool, &owner).await.unwrap();
    AuthService::new(pool.clone())
        .set_password(owner.id, "pw")
        .await
        .unwrap();
    let client = reqwest::Client::new();
    let token = client
        .post(format!("{base}/api/auth/login"))
        .json(&json!({"email": owner.email, "password": "pw"}))
        .send()
        .await
        .unwrap()
        .json::<Value>()
        .await
        .unwrap()["token"]
        .as_str()
        .unwrap()
        .to_owned();
    let claw: Value = client
        .post(format!("{base}/api/claws"))
        .bearer_auth(&token)
        .json(&json!({"name": "Scout", "job_title": "Analyst"}))
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    let claw_id = claw["id"].as_str().unwrap().to_owned();

    // Fire all streams concurrently, one session each.
    let mut tasks = Vec::new();
    for i in 0..STREAMS {
        let client = client.clone();
        let base = base.clone();
        let token = token.clone();
        let claw_id = claw_id.clone();
        tasks.push(tokio::spawn(async move {
            let session: Value = client
                .post(format!("{base}/api/sessions"))
                .bearer_auth(&token)
                .json(&json!({"clawId": claw_id, "title": format!("Load {i}")}))
                .send()
                .await
                .unwrap()
                .json()
                .await
                .unwrap();
            let session_key = session["sessionKey"].as_str().unwrap().to_owned();
            let response = client
                .post(format!("{base}/api/gateway?clawId={claw_id}"))
                .bearer_auth(&token)
                .json(&json!({
                    "sessionKey": session_key,
                    "message": format!("time check {i} [[scenario:tool-time]]"),
                }))
                .send()
                .await
                .unwrap();
            assert_eq!(
                response.status(),
                200,
                "stream {i}: {}",
                response.text().await.unwrap_or_default()
            );
            let live = collect_sse(response).await;
            (session_key, live)
        }));
    }

    let mut runs = Vec::new();
    for task in tasks {
        runs.push(task.await.unwrap());
    }
    assert_eq!(runs.len(), STREAMS);

    for (session_key, live) in &runs {
        // Completed, with the full §13 event vocabulary in order.
        let names: Vec<&str> = live.iter().map(|(_, name, _)| name.as_str()).collect();
        assert_eq!(*names.first().unwrap(), "run_started", "{session_key}");
        assert_eq!(*names.last().unwrap(), "run_completed", "{session_key}");
        assert!(names.contains(&"step_started"));
        assert!(names.contains(&"step_finished"));

        // Strictly monotonic journal ids.
        let ids: Vec<i64> = live.iter().map(|(id, _, _)| id.parse().unwrap()).collect();
        assert!(
            ids.windows(2).all(|pair| pair[1] > pair[0]),
            "monotonic ids: {ids:?}"
        );

        // Replay equals live, byte for byte.
        let replay = collect_sse(
            client
                .post(format!("{base}/api/gateway?clawId={claw_id}"))
                .bearer_auth(&token)
                .json(&json!({"sessionKey": session_key, "resumeFrom": 0}))
                .send()
                .await
                .unwrap(),
        )
        .await;
        assert_eq!(&replay, live, "replay must equal live for {session_key}");
    }
}