yggdrasil/api/code_runner/
sse.rs

1//! SSE 流式输出端点。`GET /api/exec/stream?task_id=X`。
2//!
3//! 鉴权 + 限流已在 [`super::execute::start_exec_stream`] 完成（校验链），此处只
4//! 校验 task 存在。task_id 是 UUID，不可枚举；只有过了校验的调用者才知道 task_id。
5//!
6//! 前端用原生 EventSource 连接本端点（`web_sys::EventSource`），按 event 类型
7//! 分发：`stdout` → 终端 writeStdout，`stderr` → writeStderr，`done` → 终态收尾。
8//! keep-alive comment 每 15s 一次，防反向代理超时关闭空闲连接。
9
10use std::convert::Infallible;
11use std::time::Duration;
12
13use axum::extract::Query;
14use axum::http::StatusCode;
15use axum::response::sse::{Event, KeepAlive, Sse};
16use futures::StreamExt;
17use tokio_stream::wrappers::ReceiverStream;
18
19use crate::api::code_runner::progress::EXEC_STREAMS;
20use crate::infra::docker::OutputChunk;
21
22/// SSE 查询参数：`?task_id=X`。
23#[derive(serde::Deserialize)]
24pub struct StreamQuery {
25    pub task_id: String,
26}
27
28#[derive(serde::Serialize)]
29struct DonePayload {
30    exit_code: Option<i64>,
31    oom_killed: bool,
32    timed_out: bool,
33    duration_ms: u64,
34    /// 系统级错误消息（镜像缺失 / daemon 不可达等）。正常执行时为 None。
35    #[serde(skip_serializing_if = "Option::is_none")]
36    error: Option<String>,
37}
38
39/// SSE handler：从 EXEC_STREAMS 取出 receiver（取出即移除，防重复连接），
40/// 包成 ReceiverStream，映射成 SSE Event 流。
41pub async fn exec_stream(
42    Query(q): Query<StreamQuery>,
43) -> Result<Sse<impl futures::Stream<Item = Result<Event, Infallible>>>, (StatusCode, String)> {
44    // 取出 rx（取出即移除）：同一 task_id 只能连一次 SSE，
45    // 防止多客户端或重连导致 receiver 被多次消费。
46    let (_, entry) = EXEC_STREAMS
47        .remove(&q.task_id)
48        .ok_or((StatusCode::NOT_FOUND, "任务不存在或已结束".to_string()))?;
49
50    let stream = ReceiverStream::new(entry.rx).map(|chunk| {
51        Ok::<_, Infallible>(match chunk {
52            OutputChunk::Stdout(s) => Event::default().event("stdout").data(s),
53            OutputChunk::Stderr(s) => Event::default().event("stderr").data(s),
54            OutputChunk::Done {
55                exit_code,
56                oom_killed,
57                timed_out,
58                duration_ms,
59                error,
60            } => Event::default()
61                .event("done")
62                .json_data(DonePayload {
63                    exit_code,
64                    oom_killed,
65                    timed_out,
66                    duration_ms,
67                    error,
68                })
69                .unwrap_or_else(|_| Event::default().event("done").data("{}")),
70        })
71    });
72
73    Ok(Sse::new(stream).keep_alive(
74        KeepAlive::new()
75            .interval(Duration::from_secs(15))
76            .text("keep-alive"),
77    ))
78}
yggdrasil/api/code_runner/sse.rs

yggdrasil/api/code_runner/
sse.rs
