rsclaw-runtime 2026.6.26

//! SSE handlers for A2A v1.0 streaming methods.
//!
//! Two methods land here:
//!   - SendStreamingMessage : spawn a new agent task, stream its events back
//!   - SubscribeToTask      : tap into an existing task's event broadcast
//!
//! Each event is wrapped in a JSON-RPC frame `{"jsonrpc","id","result":<wire>}`
//! and emitted as one SSE `data:` line.

use std::{convert::Infallible, sync::atomic::Ordering};

use axum::response::{
    IntoResponse, Response,
    sse::{Event, KeepAlive, Sse},
};
use futures::StreamExt;
use serde_json::{Value, json};
use tokio::sync::oneshot;
use tokio_stream::wrappers::{BroadcastStream, errors::BroadcastStreamRecvError};
use tracing::{info, warn};
use uuid::Uuid;

use crate::{
    a2a::{
        event::AgentEvent,
        relay::{audit_relay, relay_target_from_params},
        types::{
            A2aArtifact, A2aMessage, A2aTask, A2aTaskStatus, JsonRpcRequest, SendMessageParams,
            TaskState,
        },
    },
    server::AppState,
};
use rsclaw_agent::{AgentMessage, AgentReply};

/// Entry point. Called by the gateway dispatcher when the JSON-RPC method is
/// `SendStreamingMessage` or `SubscribeToTask`.
pub async fn handle_streaming_rpc(
    state: AppState,
    caller: Option<crate::a2a::auth::A2aIdentity>,
    req: JsonRpcRequest,
) -> Response {
    let req_id = req.id.clone();

    // Relay forwarding: if the request targets a spoke node, route through the
    // relay hub instead of creating a local streaming task. Targets may be
    // explicit (metadata.agentId) or implicit via a known task_id route.
    let relay_target = relay_target_from_params(&req.params).or_else(|| {
        crate::a2a::relay::task_id_from_params(&req.params)
            .and_then(|tid| state.relay_hub.route_for_task(tid))
    });
    if let Some(ref target) = relay_target {
        if state.relay_hub.route_for(target).is_some() {
            let principal = caller
                .as_ref()
                .map(|id| id.id.as_str())
                .unwrap_or("anonymous-dev");
            if !crate::a2a::relay::can_invoke(caller.as_ref(), target) {
                state
                    .relay_hub
                    .metrics
                    .acl_denials
                    .fetch_add(1, Ordering::Relaxed);
                let relay_id = state.config.gateway.a2a_relay.relay_id.as_str();
                let target_node = target.split('/').next().unwrap_or("");
                audit_relay(
                    "deny",
                    principal,
                    "invoke",
                    &format!("agent:{target}"),
                    relay_id,
                    target_node,
                    None,
                    Some("a2a:invoke scope missing"),
                );
                return sse_jsonrpc_error(
                    Some(req_id),
                    -32003,
                    format!("not authorized to invoke {target}"),
                );
            }
            let mut params = req.params.clone();
            crate::a2a::relay::rewrite_target_agent_for_spoke(&mut params, target);
            match state
                .relay_hub
                .invoke_streaming(target, &req.method, params, principal)
                .await
            {
                Ok((request_id, node_id, event_rx)) => {
                    let guard = crate::a2a::relay::RelayStreamGuard::new(
                        state.relay_hub.clone(),
                        node_id,
                        request_id,
                    );
                    let stream = tokio_stream::wrappers::BroadcastStream::new(event_rx).filter_map(
                        move |result| {
                            // Force capture so guard.Drop fires when the SSE
                            // stream is dropped (consumer disconnect).
                            let _ = &guard;
                            let req_id = req_id.clone();
                            async move {
                                match result {
                                    Ok(value) => {
                                        let payload = json!({
                                            "jsonrpc": "2.0",
                                            "id": req_id,
                                            "result": value,
                                        });
                                        Some(Ok::<_, Infallible>(
                                            Event::default().json_data(payload).unwrap_or_default(),
                                        ))
                                    }
                                    Err(BroadcastStreamRecvError::Lagged(n)) => {
                                        warn!(lagged = n, "SSE relay consumer lagged");
                                        None
                                    }
                                }
                            }
                        },
                    );
                    return Sse::new(stream)
                        .keep_alive(KeepAlive::new())
                        .into_response();
                }
                Err(e) => {
                    warn!(error = %e, "relay streaming failed, falling back");
                }
            }
        }
    }

    let (task_id, rx) = match req.method.as_str() {
        "SendStreamingMessage" => spawn_streaming_task(state.clone(), caller, req.params).await,
        "SubscribeToTask" => {
            let tid = req
                .params
                .get("id")
                .and_then(|v| v.as_str())
                .map(str::to_owned)
                .unwrap_or_default();
            // Empty id, unknown task, OR a task the caller doesn't own: emit a
            // synthetic Failed and close. §7.5 — a non-owner gets the SAME
            // "task not found" outcome as a genuinely missing task, so task
            // existence never leaks. Existing in-flight tasks owned by the
            // caller have a live sender, so subscribing returns events.
            if tid.is_empty()
                || state.task_store.get(&tid).ok().flatten().is_none()
                || !crate::a2a::server::caller_owns(&state.task_store, &caller, &tid)
            {
                let rx = state.task_event_bus.subscribe(&tid);
                state.task_event_bus.publish(AgentEvent::Status {
                    task_id: tid.clone(),
                    context_id: String::new(),
                    state: TaskState::Failed,
                    message: Some(rsclaw_a2a_types::event::text_message(if tid.is_empty() {
                        "SubscribeToTask: missing task id"
                    } else {
                        "SubscribeToTask: task not found"
                    })),
                    final_: true,
                });
                state.task_event_bus.close(&tid);
                (tid, rx)
            } else {
                let rx = state.task_event_bus.subscribe(&tid);
                (tid, rx)
            }
        }
        other => {
            // Unknown method on the SSE entry — emit Failed + close so the
            // client doesn't hang on an empty broadcast channel.
            let tid = Uuid::new_v4().to_string();
            let rx = state.task_event_bus.subscribe(&tid);
            state.task_event_bus.publish(AgentEvent::Status {
                task_id: tid.clone(),
                context_id: String::new(),
                state: TaskState::Failed,
                message: Some(rsclaw_a2a_types::event::text_message(&format!(
                    "unsupported streaming method: {other}"
                ))),
                final_: true,
            });
            state.task_event_bus.close(&tid);
            (tid, rx)
        }
    };

    drop(task_id);
    let stream = BroadcastStream::new(rx).filter_map(move |result| {
        let req_id = req_id.clone();
        async move {
            match result {
                Ok(ev) => {
                    let payload = json!({
                        "jsonrpc": "2.0",
                        "id": req_id,
                        "result": ev.to_wire_event(),
                    });
                    Some(Ok::<_, Infallible>(
                        Event::default().json_data(payload).unwrap_or_default(),
                    ))
                }
                Err(BroadcastStreamRecvError::Lagged(n)) => {
                    warn!(lagged = n, "SSE consumer lagged");
                    None
                }
            }
        }
    });

    Sse::new(stream)
        .keep_alive(KeepAlive::new())
        .into_response()
}

fn sse_jsonrpc_error(id: Option<Value>, code: i64, message: String) -> Response {
    let payload = json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": {
            "code": code,
            "message": message,
        },
    });
    let stream = tokio_stream::once(Ok::<_, Infallible>(
        Event::default().json_data(payload).unwrap_or_default(),
    ));
    Sse::new(stream).into_response()
}

/// Spawn an agent task and return `(task_id, subscriber)`. The subscriber is
/// taken BEFORE any events are published to the bus so the SSE consumer
/// observes the full Submitted → Working → Completed sequence.
pub(crate) async fn spawn_streaming_task(
    state: AppState,
    caller: Option<crate::a2a::auth::A2aIdentity>,
    params: Value,
) -> (String, tokio::sync::broadcast::Receiver<AgentEvent>) {
    let params: SendMessageParams = match serde_json::from_value(params) {
        Ok(p) => p,
        Err(e) => {
            warn!(err = %e, "SendStreamingMessage: invalid params");
            // Without a terminal Failed + bus close the SSE stream stays
            // open with no events and the client hangs until its read
            // timeout. Emit a synthetic Failed carrying the parse error
            // and close the bus so BroadcastStream sees Closed and ends.
            let tid = Uuid::new_v4().to_string();
            let rx = state.task_event_bus.subscribe(&tid);
            state.task_event_bus.publish(AgentEvent::Status {
                task_id: tid.clone(),
                context_id: String::new(),
                state: TaskState::Failed,
                message: Some(rsclaw_a2a_types::event::text_message(&format!(
                    "invalid params: {e}"
                ))),
                final_: true,
            });
            state.task_event_bus.close(&tid);
            return (tid, rx);
        }
    };

    let task_id = params
        .message
        .task_id
        .clone()
        .unwrap_or_else(|| Uuid::new_v4().to_string());

    // §7.5: stamp the creating principal so only it (or the operator token)
    // can later read / cancel / subscribe to this streaming task.
    if let Some(c) = caller.as_ref() {
        let _ = state.task_store.put_owner(&task_id, &c.id);
    }

    // CRITICAL: subscribe BEFORE any publish so the SSE consumer sees
    // Submitted/Working frames. broadcast channels don't replay history.
    let early_rx = state.task_event_bus.subscribe(&task_id);

    let session_key = params
        .message
        .context_id
        .clone()
        .unwrap_or_else(|| format!("a2a:{}", Uuid::new_v4()));

    let agent_id = params
        .metadata
        .as_ref()
        .and_then(|m| m.get("agentId").and_then(|v| v.as_str()).map(str::to_owned));

    // Ingest non-text parts (Raw/Url/Data) before the runtime sees the
    // message, so `@a2a_<kind>_...` references appear in the text and the
    // runtime's existing `resolve_file_refs` pipeline loads them as image
    // attachments / file references on the agent loop entry.
    let workspace =
        crate::a2a::server::resolve_agent_workspace_pub(&state, agent_id.as_deref()).await;
    let ingested = crate::a2a::files::ingest_message_parts(&workspace, &params.message.parts).await;
    let text = ingested.text;

    let handle = match agent_id {
        Some(aid) => state.agents.get(&aid),
        None => state.agents.default_agent(),
    };
    let Ok(handle) = handle else {
        warn!("SendStreamingMessage: no agent available");
        // Publish Failed AND close the bus so the SSE stream actually
        // terminates — without the close the subscriber sees Failed but
        // BroadcastStream stays subscribed forever waiting for the next
        // event that will never come.
        state.task_event_bus.publish(AgentEvent::Status {
            task_id: task_id.clone(),
            context_id: session_key.clone(),
            state: TaskState::Failed,
            message: Some(rsclaw_a2a_types::event::text_message(
                "no agent available for this task",
            )),
            final_: true,
        });
        state.task_event_bus.close(&task_id);
        return (task_id, early_rx);
    };

    // Reply oneshot — we await this in a spawned task so we can publish
    // Completed/Failed and the final Artifact when the agent's turn finishes.
    let (reply_tx, reply_rx) = oneshot::channel::<AgentReply>();

    // mpsc channel from runtime → bus bridge. Does NOT close the bus on
    // drop — the bus stays open until the reply-watcher (below) publishes
    // the terminal status event. Otherwise the bus would close as soon as
    // the agent destructured AgentMessage (dropping event_tx) and any later
    // Completed publish would land on a fresh channel with no subscribers.
    let (event_tx, mut event_rx) = tokio::sync::mpsc::channel::<AgentEvent>(64);
    let bus_for_bridge = state.task_event_bus.clone();
    tokio::spawn(async move {
        while let Some(ev) = event_rx.recv().await {
            bus_for_bridge.publish(ev);
        }
    });

    let cancel_token = tokio_util::sync::CancellationToken::new();
    state
        .task_cancels
        .insert(task_id.clone(), cancel_token.clone());

    // Persist initial task state (Submitted).
    let initial_history = A2aMessage {
        message_id: params.message.message_id.clone(),
        role: params.message.role.clone(),
        parts: params.message.parts.clone(),
        context_id: Some(session_key.clone()),
        task_id: Some(task_id.clone()),
        metadata: params.message.metadata.clone(),
    };
    let initial_task = A2aTask {
        id: task_id.clone(),
        context_id: Some(session_key.clone()),
        status: A2aTaskStatus {
            state: TaskState::Submitted,
            message: None,
            timestamp: Some(chrono::Utc::now().to_rfc3339()),
        },
        history: vec![initial_history],
        artifacts: vec![],
        metadata: None,
    };
    if let Err(e) = state.task_store.put(&initial_task) {
        warn!(err = %e, "failed to persist initial streaming task");
    }

    // Mirror events into the persistent store as they arrive.
    let persist_store = state.task_store.clone();
    let persist_task_id = task_id.clone();
    let mut persist_rx = state.task_event_bus.subscribe(&task_id);
    tokio::spawn(async move {
        while let Ok(ev) = persist_rx.recv().await {
            match ev {
                AgentEvent::Artifact {
                    artifact_id, parts, ..
                } => {
                    let _ = persist_store.append_artifact(
                        &persist_task_id,
                        A2aArtifact {
                            artifact_id,
                            parts,
                            name: None,
                            description: None,
                            metadata: None,
                        },
                    );
                }
                AgentEvent::Status { state, final_, .. } => {
                    let _ = persist_store.set_status(&persist_task_id, state);
                    if final_ {
                        // Push config GC: terminal state means no more
                        // webhooks will fire for this task — clear its
                        // notification configs so they don't linger in
                        // the store. (No-op if there were no configs.)
                        let _ = persist_store.delete_push_configs_for_task(&persist_task_id);
                        break;
                    }
                }
                AgentEvent::InputRequired { .. } => {
                    let _ = persist_store.set_status(&persist_task_id, TaskState::InputRequired);
                }
                AgentEvent::AuthRequired { .. } => {
                    let _ = persist_store.set_status(&persist_task_id, TaskState::AuthRequired);
                }
            }
        }
    });

    // Push notification fan-out.
    state.push_dispatcher.clone().watch(task_id.clone());

    // Publish Submitted → Working status events so SSE subscribers see progress.
    state.task_event_bus.publish(AgentEvent::Status {
        task_id: task_id.clone(),
        context_id: session_key.clone(),
        state: TaskState::Submitted,
        message: None,
        final_: false,
    });
    state.task_event_bus.publish(AgentEvent::Status {
        task_id: task_id.clone(),
        context_id: session_key.clone(),
        state: TaskState::Working,
        message: None,
        final_: false,
    });

    // Spawn a watcher that, when the agent's reply arrives, publishes the
    // final artifact + Completed status (or Failed if the channel dropped).
    let bus_for_reply = state.task_event_bus.clone();
    let task_id_for_reply = task_id.clone();
    let ctx_id_for_reply = session_key.clone();
    let cancels_for_reply = state.task_cancels.clone();
    tokio::spawn(async move {
        match reply_rx.await {
            Ok(reply) => match reply.outcome {
                rsclaw_agent::registry::ReplyOutcome::Ok => {
                    let artifact_id = uuid::Uuid::new_v4().to_string();
                    bus_for_reply.publish(AgentEvent::Artifact {
                        task_id: task_id_for_reply.clone(),
                        context_id: ctx_id_for_reply.clone(),
                        artifact_id,
                        parts: crate::a2a::files::emit_reply_parts(
                            &reply.text,
                            &reply.images,
                            &reply.files,
                        ),
                        append: false,
                        last_chunk: true,
                    });
                    bus_for_reply.publish(AgentEvent::Status {
                        task_id: task_id_for_reply.clone(),
                        context_id: ctx_id_for_reply,
                        state: TaskState::Completed,
                        message: None,
                        final_: true,
                    });
                }
                rsclaw_agent::registry::ReplyOutcome::Error => {
                    // Surface the error text in the terminal Failed message so
                    // SSE/push subscribers see the cause; no Artifact (Artifact
                    // implies usable output).
                    bus_for_reply.publish(AgentEvent::Status {
                        task_id: task_id_for_reply.clone(),
                        context_id: ctx_id_for_reply,
                        state: TaskState::Failed,
                        message: Some(rsclaw_a2a_types::event::text_message(&reply.text)),
                        final_: true,
                    });
                }
                rsclaw_agent::registry::ReplyOutcome::Canceled => {
                    // CancelTask dispatcher already published the terminal
                    // Canceled event and closed the bus. Don't republish.
                }
            },
            Err(_) => {
                bus_for_reply.publish(AgentEvent::Status {
                    task_id: task_id_for_reply.clone(),
                    context_id: ctx_id_for_reply,
                    state: TaskState::Failed,
                    message: None,
                    final_: true,
                });
            }
        }
        cancels_for_reply.remove(&task_id_for_reply);
        // Now safe to drop the broadcast channel — terminal event delivered.
        bus_for_reply.close(&task_id_for_reply);
    });

    // Wire the INPUT_REQUIRED resume channel — with the same timeout
    // semantics as the sync path. Helper handles registration + timeout
    // cleanup so a client that never sends the resume SendMessage
    // doesn't leak a suspended entry.
    let (ireq_tx, ireq_rx) = tokio::sync::mpsc::channel::<tokio::sync::oneshot::Sender<String>>(4);
    crate::a2a::server::spawn_input_request_listener(
        state.clone(),
        task_id.clone(),
        session_key.clone(),
        ireq_rx,
    );

    let msg = AgentMessage {
        session_key: session_key.clone(),
        text,
        channel: "a2a".to_owned(),
        // A2A 已鉴权身份作为可信发送方(不再丢弃成常量)。竞猜联赛中枢据此做
        // "一身份一预测"的 nodeId,避免参赛者在消息体里自报 nodeId 被冒充。
        // 无 caller(鉴权关闭/dev)时回退旧常量,行为不变。
        peer_id: caller
            .as_ref()
            .map(|c| c.id.clone())
            .unwrap_or_else(|| "a2a-client".to_owned()),
        chat_id: String::new(),
        reply_tx,
        task_id: Some(task_id.clone()),
        context_id: Some(session_key),
        event_tx: Some(event_tx),
        cancel_token: Some(cancel_token),
        input_request_tx: Some(ireq_tx),
        extra_tools: vec![],
        images: vec![],
        files: vec![],
        account: None,
    };

    if let Err(e) = handle.tx.send(msg).await {
        warn!(err = ?e, "agent inbox closed");
    } else {
        info!(task_id = %task_id, "A2A SendStreamingMessage spawned");
    }
    (task_id, early_rx)
}