clawmates/crates/cm-api/src/routes/gateway.rs

use std::convert::Infallible;
use std::time::Duration;

use axum::extract::{Query, State};
use axum::response::sse::{Event, KeepAlive, Sse};
use axum::Json;
use cm_domain::AgentId;
use cm_runtime::{RunEventBody, RunEventEnvelope};
use futures::stream::BoxStream;
use serde::Deserialize;
use tokio::sync::broadcast;

use crate::routes::claws::workspace_agent;
use crate::routes::sessions::scoped_session;
use crate::{ApiError, AppState, Authed};

#[derive(Deserialize)]
pub struct GatewayQuery {
    #[serde(rename = "clawId")]
    claw_id: AgentId,
}

#[derive(Deserialize)]
pub struct GatewayRequest {
    #[serde(rename = "sessionKey")]
    session_key: String,
    /// New user message — starts a run.
    message: Option<String>,
    /// Reconnect offset — replays the journal after this seq, then goes
    /// live if the run is still streaming.
    #[serde(rename = "resumeFrom")]
    resume_from: Option<i64>,
}

fn envelope_to_sse(envelope: &RunEventEnvelope) -> Event {
    let payload = serde_json::to_value(&envelope.event).expect("event serializes");
    Event::default()
        .id(envelope.seq.to_string())
        .event(payload["type"].as_str().expect("tagged event"))
        .data(payload.to_string())
}

fn is_terminal(event: &RunEventBody) -> bool {
    matches!(
        event,
        RunEventBody::RunCompleted { .. } | RunEventBody::Error { .. }
    )
}

type SseStream = BoxStream<'static, Result<Event, Infallible>>;

fn live_stream(mut rx: broadcast::Receiver<RunEventEnvelope>, after: i64) -> SseStream {
    Box::pin(async_stream::stream! {
        loop {
            match rx.recv().await {
                Ok(envelope) => {
                    if envelope.seq <= after {
                        continue;
                    }
                    let done = is_terminal(&envelope.event);
                    yield Ok(envelope_to_sse(&envelope));
                    if done {
                        break;
                    }
                }
                Err(broadcast::error::RecvError::Lagged(_)) => continue,
                Err(broadcast::error::RecvError::Closed) => break,
            }
        }
    })
}

/// POST /api/gateway?clawId= — the single audited streaming channel (§15).
/// Every emitted event was journaled first, so a reconnect with
/// `resumeFrom` replays exactly what live observers saw.
pub async fn gateway(
    State(state): State<AppState>,
    Authed(user): Authed,
    Query(query): Query<GatewayQuery>,
    Json(body): Json<GatewayRequest>,
) -> Result<impl axum::response::IntoResponse, ApiError> {
    let agent = workspace_agent(&state, &user, query.claw_id).await?;
    let session = scoped_session(&state, &user, &body.session_key).await?;
    if session.agent_id != agent.id {
        return Err(ApiError::NotFound);
    }

    let stream: SseStream = match body.message {
        Some(text) => {
            let started = state
                .runtime
                .send_message(session.id, &text)
                .await
                .map_err(|e| match e {
                    cm_runtime::RuntimeError::Db(cm_db::DbError::NotFound) => ApiError::NotFound,
                    _ => ApiError::Internal,
                })?;
            live_stream(started.events, 0)
        }
        None => {
            let resume_from = body.resume_from.unwrap_or(0);
            let run = cm_db::repo::runs::latest_for_session(&state.pool, session.id)
                .await?
                .ok_or(ApiError::NotFound)?;
            // Subscribe before reading the journal so no event falls in the
            // gap; the live tail then skips anything the replay covered.
            let live = state.runtime.subscribe(run.id).await;
            let journal =
                cm_db::repo::run_events::list_after(&state.pool, run.id, resume_from).await?;
            let last_replayed = journal.last().map(|e| e.seq).unwrap_or(resume_from);
            let replay_done = journal
                .last()
                .map(|e| e.event_type == "run_completed" || e.event_type == "error")
                .unwrap_or(false);

            Box::pin(async_stream::stream! {
                for entry in journal {
                    yield Ok(Event::default()
                        .id(entry.seq.to_string())
                        .event(entry.event_type.clone())
                        .data(entry.payload.to_string()));
                }
                if !replay_done {
                    if let Some(rx) = live {
                        let mut tail = live_stream(rx, last_replayed);
                        while let Some(event) = futures::StreamExt::next(&mut tail).await {
                            yield event;
                        }
                    }
                }
            })
        }
    };

    Ok(Sse::new(stream).keep_alive(
        KeepAlive::new()
            .interval(Duration::from_secs(15))
            .text("hb"),
    ))
}