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/// done 事件的 JSON payload。
29#[derive(serde::Serialize)]
30struct DonePayload {
31    exit_code: Option<i64>,
32    oom_killed: bool,
33    timed_out: bool,
34    duration_ms: u64,
35}
36
37/// SSE handler：从 EXEC_STREAMS 取出 receiver（取出即移除，防重复连接），
38/// 包成 ReceiverStream，映射成 SSE Event 流。
39pub async fn exec_stream(
40    Query(q): Query<StreamQuery>,
41) -> Result<Sse<impl futures::Stream<Item = Result<Event, Infallible>>>, (StatusCode, String)> {
42    // 取出 rx（取出即移除）：同一 task_id 只能连一次 SSE，
43    // 防止多客户端或重连导致 receiver 被多次消费。
44    let (_, entry) = EXEC_STREAMS
45        .remove(&q.task_id)
46        .ok_or((StatusCode::NOT_FOUND, "任务不存在或已结束".to_string()))?;
47
48    let stream = ReceiverStream::new(entry.rx).map(|chunk| {
49        Ok::<_, Infallible>(match chunk {
50            OutputChunk::Stdout(s) => Event::default().event("stdout").data(s),
51            OutputChunk::Stderr(s) => Event::default().event("stderr").data(s),
52            OutputChunk::Done {
53                exit_code,
54                oom_killed,
55                timed_out,
56                duration_ms,
57            } => Event::default()
58                .event("done")
59                .json_data(DonePayload {
60                    exit_code,
61                    oom_killed,
62                    timed_out,
63                    duration_ms,
64                })
65                .unwrap_or_else(|_| Event::default().event("done").data("{}")),
66        })
67    });
68
69    Ok(Sse::new(stream).keep_alive(
70        KeepAlive::new()
71            .interval(Duration::from_secs(15))
72            .text("keep-alive"),
73    ))
74}
yggdrasil/api/code_runner/sse.rs

yggdrasil/api/code_runner/
sse.rs
