diff --git a/Cargo.lock b/Cargo.lock index 3075352b..9ac40e37 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2974,6 +2974,7 @@ dependencies = [ name = "tw-store" version = "0.57.1" dependencies = [ + "blake3", "bytes", "rusqlite", "serde", diff --git a/bin/twcore/src/main.rs b/bin/twcore/src/main.rs index 815ee3fc..ebdc9882 100644 --- a/bin/twcore/src/main.rs +++ b/bin/twcore/src/main.rs @@ -782,7 +782,7 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> // // body 的通道在这里建：**它是唯一同时看得见网关和存储的地方**， // 而两边各有各的同形结构，是为了不让「观测」挂到「转发」下面。 - let (body_tx, body_rx) = tokio::sync::mpsc::channel(tw_gateway::bodies::CHANNEL_CAP); + let (body_tx, body_rx) = tw_gateway::bodies::channel(); let store = build_store(&dir, state.bus.clone(), state.pricing.clone(), body_rx); if store.is_some() { state.set_body_sink(body_tx); @@ -975,25 +975,45 @@ fn build_store( "could not read the last request id, so this run may overwrite the oldest records: {e}" ), } - // 两边的 body 结构在这里对接。**一次移动，不复制** —— `Bytes` 的 - // 克隆是引用计数。 - let (tx, rx) = tokio::sync::mpsc::channel(tw_gateway::bodies::CHANNEL_CAP); + /* + 两边的 body 结构在这里对接，**落盘的不是原文**：脱敏规则认得出的值在这里换掉、 + 打码（`BodyRecord::for_disk`，见 `tw_gateway::bodies`）。放在阻塞线程上 —— + 一份 4 MB 的正文要扫好几遍，占着异步线程的话，同一个线程上的转发都得等它。 + + 交给存储层的这一头**只留一个空位**。等着落盘的正文由网关那一头按字节记账 + （`bodies::QUEUED_MAX`），一份正文的额度要等存储层收下它才还回去；这里再开一个 + 大口子的话，攒在这里的那些就没人管了。 + */ + let (tx, rx) = tokio::sync::mpsc::channel(1); let mut bodies = bodies; tokio::spawn(async move { while let Some(b) = bodies.recv().await { + let Ok(disk) = tokio::task::spawn_blocking(move || b.for_disk()).await else { + continue; + }; + let tw_gateway::bodies::ForDisk { + id, + at_ms, + kind, + body, + original_len, + held, + } = disk; let mapped = tw_store::StoredBody { - id: b.id, - at_ms: b.at_ms, - which: match b.kind { + id, + at_ms, + which: match kind { tw_gateway::bodies::BodyKind::Request => tw_store::Which::Request, tw_gateway::bodies::BodyKind::Response => tw_store::Which::Response, }, - body: b.body, - original_len: b.original_len, + body, + original_len, }; if tx.send(mapped).await.is_err() { return; } + // 存储层收下了：额度还回去 + drop(held); } }); Some(tw_store::task::spawn( diff --git a/crates/tw-api/src/ep.rs b/crates/tw-api/src/ep.rs index 9802c143..cf15a809 100644 --- a/crates/tw-api/src/ep.rs +++ b/crates/tw-api/src/ep.rs @@ -60,6 +60,8 @@ endpoints! { Fixture: GET "/request/{id}/fixture" [id], () => String, text; Sessions: GET "/sessions", api::ListQuery => Vec; SessionDetail: GET "/sessions/{id}" [id], () => api::SessionDetail; + /// 一次会话读成一段对话：每一轮新说的话、回答、工具调用和结果（已脱敏） + SessionTranscript: GET "/sessions/{id}/transcript" [id], () => api::Transcript; // ─────────────────────────────────────────────── 测速、回放、试路由 SpeedQuote: POST "/speed/quote", api::SpeedRunRequest => api::SpeedQuote; diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index acf6bff9..0f54f7f7 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -672,9 +672,13 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// `cn-resident-id` 和 `bank-card`（[`CardNetwork`]、[`CardPrefix`]），[`SecretKind`] 多了 /// `personal`。照 30 写的界面说不出这两条规则按什么认。 /// -/// **32 起有脚本插件**：事件多了 [`Event::PluginFailed`]（插件在请求上出错，或者文件变了、 -/// 加载不了而停用）。照 31 写的界面不认这个事件。 -pub const CONTROL_API_VERSION: u32 = 32; +/// **32 起会话能读成一段对话**：新端点 `GET /sessions/{id}/transcript`（[`Transcript`]） +/// 从存下来的正文里读出每一轮新说的话、回答、推理、工具调用和结果，读不到的地方逐轮说出来 +/// （[`TranscriptGap`]）。照 31 写的界面只有每一轮的用量和金额。 +/// +/// **33 起有脚本插件**：事件多了 [`Event::PluginFailed`]（插件在请求上出错，或者文件变了、 +/// 加载不了而停用）。照 32 写的界面不认这个事件。 +pub const CONTROL_API_VERSION: u32 = 33; #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] @@ -1723,7 +1727,7 @@ pub struct RetentionView { /// 正文总共最多占多少字节 pub body_max_bytes: u64, /// 正文现在实际占了多少。**不是配置，是现状** —— 没有它， - /// 「2 GB 上限」是个用户无从判断松紧的数字 + /// 「5 GB 上限」是个用户无从判断松紧的数字 pub body_bytes_now: u64, } @@ -3693,14 +3697,23 @@ pub struct RequestDetail { pub in_flight: bool, } +/// 一份正文最多存多少字节：请求和回答一样，4 MiB。更长的只存开头，[`BodyView::truncated`] +/// 说出来。 +/// +/// **存储层按它截，网关攒回答也按它攒**（`tw_store::blobs::MAX_ONE`、 +/// `tw_gateway::bodies::RESPONSE_TAP_MAX`）。两个数放在一处：各写各的话，改了一个， +/// 另一个还停在原地 —— 以前回答只攒 256 KB，比存储层肯收的少十几倍。 +pub const BODY_MAX: usize = 4 * 1024 * 1024; + /// 一份存下来的 body。 #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] pub struct BodyView { - /// **已脱敏**。这段文字会被复制到 issue 里 + /// **已脱敏**。这段文字会被复制到 issue 里。落盘的那一份就是换过、打过码的（脱敏规则 + /// 认得出的值不会原样写进磁盘），读出来再打一遍 pub text: String, /// 原本多长。**截断了要能说出来** —— 不说的话用户会以为请求本身 - /// 就长这样 + /// 就长这样。没截断的就是存下来的这一份的长度：换掉、打码的那几处和原文差几个字节 pub original_len: usize, pub truncated: bool, } @@ -4047,6 +4060,123 @@ pub struct SessionDetail { pub turns: Vec, } +/// 一次会话读成一段对话（`GET /sessions/{id}/transcript`）：每一轮新说的话、回答、推理、 +/// 工具调用和工具结果。 +/// +/// **从存下来的正文里读出来**，不是另记的一份：正文只留几天（`retention.body_days`）， +/// 太大的只留开头，没存下来的也有。读不到的地方，那一轮的 `gaps` 说出来。 +/// +/// **已脱敏**，和请求详情里的正文同一套打码。图片只说类型和大小，从不带数据。 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct Transcript { + pub session: String, + /// 第一个读得懂的请求里的系统提示：Anthropic 的 `system`、Responses 的 `instructions`、 + /// Gemini 的 `systemInstruction`，Chat 和 Responses 还有开头连着的 system、developer + /// 消息，几段之间空一行。没有是 null + pub system: Option, + /// 和 [`SessionDetail::turns`] 同样的请求，同样的顺序 + pub turns: Vec, +} + +/// 对话里的一轮，就是会话里的一个请求。 +/// +/// 客户端每一轮都把整段历史发上来：请求 i 的消息 = 请求 i-1 的消息 + 上一轮的回答 + 新的 +/// 用户消息或工具结果。`input` 只放新的那几条；上一轮的回答已经在上一轮的 `output` 里。 +/// +/// **不生成回答的调用**（数 token、Responses 的压缩）也在这里占一轮，`input`、`output` +/// 都是空的，也不和前后的请求比对：它们问的是这段对话，不是对话里的一句。 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct TranscriptTurn { + /// 请求号，写成十进制的字符串。和 [`TurnView::id`] 是同一条请求 + pub id: String, + /// 这个请求带的历史没有接着上一个读得懂的请求：压缩过、改过历史，或者它是一串读不懂 + /// 的请求之后第一个读得懂的。这时 `input` 是它的整段历史 + pub restart: bool, + /// 系统提示和上一个读得懂的请求不一样了：新的那一份（去掉了的是空串）。没变是 null + pub system_changed: Option, + /// 这个请求里新的消息。上一轮的回答没有完整读出来时（那一轮的 `gaps` 里有 `response_*`）， + /// 客户端记下的那条助手消息也在这里：它是那一轮说过什么的记录 + pub input: Vec, + /// 回答，从存下来的响应里读出来的。失败的请求（上游回了错误）没有回答，也不算缺 + pub output: Vec, + /// 这一轮哪些地方读不出来 + pub gaps: Vec, +} + +/// 请求里的一条消息。 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct TranscriptMessage { + pub role: TranscriptRole, + pub parts: Vec, +} + +slug_enum! { + /// 一条消息是谁说的。 + pub enum TranscriptRole { + User = "user", + Assistant = "assistant", + /// 只装着工具结果的消息：Anthropic 全是 `tool_result` 的用户消息、Chat 的 `tool` + /// 消息、Responses 的 `function_call_output` + Tool = "tool", + /// 对话中途的 system、developer 消息 + System = "system", + } +} + +/// 消息或回答里的一块。 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum TranscriptPart { + Text { + text: String, + }, + /// 推理。只有签名、或者被打码的推理，`text` 是空串 + Thinking { + text: String, + }, + /// `input` 是参数的 JSON 文本；自由格式的工具（Codex 的 `apply_patch`）是它的原文 + ToolCall { + id: String, + name: String, + input: String, + }, + /// `call_id` 和它回应的那个 `tool_call` 的 `id` 是同一个 + ToolResult { + call_id: String, + text: String, + is_error: bool, + }, + /// 图片。**只有类型和大小，从不带数据**；给的是地址的不知道大小 + Image { + media_type: Option, + bytes: Option, + }, + /// 别的块：文件、音频、服务端工具的调用和结果……`label` 是它的类型名，原样 + Other { + label: String, + }, +} + +slug_enum! { + /// 一轮里读不出来的地方。 + pub enum TranscriptGap { + /// 请求体没有存下来，或者已经清掉了 + RequestMissing = "request_missing", + /// 请求体只存了开头，或者解析不了 + RequestTruncated = "request_truncated", + /// 回答没有存下来，或者已经清掉了 + ResponseMissing = "response_missing", + /// 回答只存了开头：`output` 是读得出来的那一段 + ResponseTruncated = "response_truncated", + /// 回答存下来了，但读不懂 + ResponseUnreadable = "response_unreadable", + } +} + // ---------------------------------------------------------------- 请求重放 /// 把存下来的那条请求，原样发给另一个上游。 diff --git a/crates/tw-api/src/ts.rs b/crates/tw-api/src/ts.rs index 0fd9ad1b..f2f28142 100644 --- a/crates/tw-api/src/ts.rs +++ b/crates/tw-api/src/ts.rs @@ -282,6 +282,58 @@ mod tests { assert!(event.contains("answered_model?: string"), "{event}"); } + /// 对话记录：请求号是字符串；可以为空的是 null（必有的字段，不是省掉）；块是按 `kind` + /// 分的联合，角色和缺口是字面量 + #[test] + fn a_transcript_is_turns_of_messages_and_parts() { + let ts = typescript(); + assert!(ts.contains(" SessionTranscript: { req: null; res: Transcript };")); + assert!(ts.contains( + " SessionTranscript: { method: \"GET\", path: \"/sessions/{id}/transcript\", params: [\"id\"], format: \"json\" }," + )); + let transcript = decl_of(&ts, "Transcript"); + for field in [ + "session: string", + "system: string | null", + "turns: Array", + ] { + assert!(transcript.contains(field), "{field}: {transcript}"); + } + let turn = decl_of(&ts, "TranscriptTurn"); + for field in [ + "id: string", + "restart: boolean", + "system_changed: string | null", + "input: Array", + "output: Array", + "gaps: Array", + ] { + assert!(turn.contains(field), "{field}: {turn}"); + } + assert_eq!( + decl_of(&ts, "TranscriptMessage"), + "export type TranscriptMessage = { role: TranscriptRole, parts: Array, }" + ); + assert_eq!( + decl_of(&ts, "TranscriptPart"), + "export type TranscriptPart = { \"kind\": \"text\", text: string, } \ + | { \"kind\": \"thinking\", text: string, } \ + | { \"kind\": \"tool_call\", id: string, name: string, input: string, } \ + | { \"kind\": \"tool_result\", call_id: string, text: string, is_error: boolean, } \ + | { \"kind\": \"image\", media_type: string | null, bytes: number | null, } \ + | { \"kind\": \"other\", label: string, }" + ); + assert_eq!( + decl_of(&ts, "TranscriptRole"), + "export type TranscriptRole = \"user\" | \"assistant\" | \"tool\" | \"system\"" + ); + assert_eq!( + decl_of(&ts, "TranscriptGap"), + "export type TranscriptGap = \"request_missing\" | \"request_truncated\" \ + | \"response_missing\" | \"response_truncated\" | \"response_unreadable\"" + ); + } + #[test] fn every_endpoint_is_in_the_table() { let ts = typescript(); diff --git a/crates/tw-config/src/lib.rs b/crates/tw-config/src/lib.rs index 70bf42c5..135e111a 100644 --- a/crates/tw-config/src/lib.rs +++ b/crates/tw-config/src/lib.rs @@ -1183,7 +1183,7 @@ mod tests { assert_eq!(back.retention.body_days, 3); // 没写的那两个仍然是默认值，不是 0 —— 0 会让 gc 把一切都删掉 assert_eq!(back.retention.row_days, 90); - assert_eq!(back.retention.body_max_bytes, 2 * 1024 * 1024 * 1024); + assert_eq!(back.retention.body_max_bytes, 5 * 1024 * 1024 * 1024); assert!( serde_yaml_ng::to_string(&back) .unwrap() diff --git a/crates/tw-config/src/retention.rs b/crates/tw-config/src/retention.rs index 0e95d99b..5afb769f 100644 --- a/crates/tw-config/src/retention.rs +++ b/crates/tw-config/src/retention.rs @@ -21,7 +21,7 @@ pub struct Retention { /// 一行记录留几天。它撑着「上个月花了多少」那类问题 #[serde(default = "d_row_days")] pub row_days: u64, - /// 正文总共最多占多少字节。超了从最旧的整天开始删 + /// 正文总共最多占多少字节。超了从最旧的整天开始删。出厂 5 GiB #[serde(default = "d_body_max_bytes")] pub body_max_bytes: u64, } @@ -33,7 +33,7 @@ fn d_row_days() -> u64 { 90 } fn d_body_max_bytes() -> u64 { - 2 * 1024 * 1024 * 1024 + 5 * 1024 * 1024 * 1024 } impl Default for Retention { diff --git a/crates/tw-config/tests/manual/schema.rs b/crates/tw-config/tests/manual/schema.rs index 5cc1142e..a05baa63 100644 --- a/crates/tw-config/tests/manual/schema.rs +++ b/crates/tw-config/tests/manual/schema.rs @@ -1165,10 +1165,10 @@ pub fn sections() -> Vec
{ row( "body_max_bytes", Kind::Int, - Def::Is("2147483648"), + Def::Is("5368709120"), t( - "Upper bound on the bytes bodies may take; beyond it the oldest days go first. The default is 2 GiB.", - "正文最多占用的字节数，超出时从最早的日期开始删除。默认 2 GiB。", + "Upper bound on the bytes bodies may take; beyond it the oldest days go first. The default is 5 GiB.", + "正文最多占用的字节数，超出时从最早的日期开始删除。默认 5 GiB。", ), ), ], diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index 72208457..93ac21a6 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -131,6 +131,7 @@ pub fn router(state: ControlState) -> Router { .at(ep::ReplayRun, replay::run) .at(ep::Sessions, sessions) .at(ep::SessionDetail, session_detail) + .at(ep::SessionTranscript, session_transcript) .at(ep::DryRun, dryrun::dry_run) // 为客户端发专用密钥。接管本身在桌面端做 .at(ep::ClientKey, clients::client_key) @@ -389,7 +390,7 @@ async fn overview(State(s): State) -> Json { body_days: cfg.retention.body_days, row_days: cfg.retention.row_days, body_max_bytes: cfg.retention.body_max_bytes, - // **现状和配置一起给。**「上限 2 GB」这个数字，用户没法 + // **现状和配置一起给。**「上限 5 GB」这个数字，用户没法 // 判断松还是紧，除非同时看得见现在占了多少 body_bytes_now, }, @@ -1501,12 +1502,7 @@ async fn session_detail( .iter() .find(|x| x.id == id) .map(session_view) - .ok_or_else(|| { - fail( - StatusCode::NOT_FOUND, - msg!("control.session_not_found", id = id.clone() => "There is no session {id}."), - ) - })?; + .ok_or_else(|| no_such_session(&id))?; let turns = g .db() .turns(&id) @@ -1517,6 +1513,38 @@ async fn session_detail( Ok(Json(tw_api::SessionDetail { session, turns })) } +fn no_such_session(id: &str) -> Fail { + fail( + StatusCode::NOT_FOUND, + msg!("control.session_not_found", id = id.to_string() => "There is no session {id}."), + ) +} + +/// 一次会话读成一段对话（见 `tw_store::transcript`）。 +/// +/// **放到阻塞线程上跑**，库只在取行的时候锁一下：几百轮的会话要读几百份正文、解析几百 MB +/// 的 JSON，和按正文找是同一个道理（见 [`history_search`]）—— 记录和正文落盘走的是同一把锁。 +async fn session_transcript( + State(s): State, + axum::extract::Path(id): axum::extract::Path, +) -> Result, Fail> { + let store = need_store(&s)?; + let (rows, blobs) = { + let g = store.lock().await; + let rows = g.db().session_requests(&id).map_err(records)?; + // 正文目录只是一个路径，读它不需要锁 + (rows, tw_store::Blobs::new(g.blobs().root().to_path_buf())) + }; + if rows.is_empty() { + return Err(no_such_session(&id)); + } + let transcript = + tokio::task::spawn_blocking(move || tw_store::transcript::build(&id, &rows, &blobs)) + .await + .map_err(internal)?; + Ok(Json(transcript)) +} + #[derive(Debug, thiserror::Error)] pub enum ControlError { #[error("the control-plane socket {path} could not be started: {source}")] diff --git a/crates/tw-control/src/replay.rs b/crates/tw-control/src/replay.rs index 13f44156..ace9d553 100644 --- a/crates/tw-control/src/replay.rs +++ b/crates/tw-control/src/replay.rs @@ -20,6 +20,17 @@ //! **三、脱敏照做。**重放走的是控制面，不经过数据面的管线，所以 //! 脱敏那一层要在这里显式调一次。少了它，一条本来会被脱敏的请求， //! 会因为「重放」这个动作把密钥原样发给中转站。 +//! +//! # 重放的是存下来的那一份 +//! +//! 存下来的请求不是原文：脱敏规则认得出的值落盘之前就换掉了（见 +//! `tw_gateway::bodies`）—— 拦截档下是发给上游的占位符（`<>`），别的档位 +//! 打了码（`sk-an…7f9c`）。所以重放发出去的就是这样一份：拦截档下它带着的正是原来那一次 +//! 发给上游的占位符，别的档位少了原值。**换不回来，也不该换回来**：真值从没进过磁盘。 +//! +//! 那些占位符这一次没有对应的原值，回答里提到它们也原样留着。这一次重新找到的值从 +//! 存下来的那些号后面接着编（[`tw_gateway::guard::ledger_for`]），不会和它们撞号 —— +//! 撞了的话，回答里的 1 号会被还原成这一次新找到的那个值，而它原本指的是另一样。 use std::time::Instant; @@ -34,10 +45,11 @@ fn fail(code: StatusCode, detail: tw_types::Msg) -> Fail { (code, axum::Json(detail)) } -/// 找到那条请求，把**原样的**请求体取出来。 +/// 找到那条请求，把存下来的请求体取出来。 /// -/// 注意不是 `request_detail` 里那份 —— 那一份是脱敏之后给人看的 -/// （它会被复制进 issue）。重放要的是原样。 +/// 注意不是 `request_detail` 里那份 —— 那一份读出来又打了一遍码，是给人看的 +/// （它会被复制进 issue）。重放要的是存下来的那一份：它落盘之前已经换过、打过码， +/// 结构和原来的一样。 fn stored_body( g: &tw_store::Recorder, id: i64, @@ -206,12 +218,15 @@ pub async fn run( })?; // **脱敏照做。**重放不经过数据面的管线，少了这一行，一条本来会被 - // 脱敏的请求会因为「重放」这个动作把密钥原样发出去 + // 脱敏的请求会因为「重放」这个动作把密钥原样发出去。新找到的值让开存下来的那份 + // 里已经写着的占位符（见模块说明） let rt = s.gateway.runtime(); + let ledger = tw_gateway::guard::ledger_for(&raw); let (body, ledger) = tw_gateway::guard::replace( rt.config.security.redact.mode, &rt.redact, bytes::Bytes::from(raw), + &ledger, ); let url = tw_gateway::forward::upstream_url(&provider.base_url, &row.path, None); diff --git a/crates/tw-control/tests/replay.rs b/crates/tw-control/tests/replay.rs index 0befd910..647e4fa0 100644 --- a/crates/tw-control/tests/replay.rs +++ b/crates/tw-control/tests/replay.rs @@ -126,6 +126,14 @@ fn row(id: i64, provider: &str) -> tw_store::db::RequestRow { /// 一个记着一条请求（连同请求体）的控制面。 fn app(config: &str) -> (tempfile::TempDir, axum::Router) { + app_with( + config, + br#"{"model":"claude-sonnet-4-5","max_tokens":8,"messages":[{"role":"user","content":"hi"}]}"#, + ) +} + +/// 同上，存着的请求体是 `body`。 +fn app_with(config: &str, body: &[u8]) -> (tempfile::TempDir, axum::Router) { let d = tempfile::tempdir().unwrap(); let p = d.path().join("config.yaml"); std::fs::write(&p, config).unwrap(); @@ -133,7 +141,6 @@ fn app(config: &str) -> (tempfile::TempDir, axum::Router) { let r = row(1, "本机"); db.insert(&r).unwrap(); let blobs = tw_store::Blobs::new(d.path().join("blobs")); - let body = br#"{"model":"claude-sonnet-4-5","max_tokens":8,"messages":[{"role":"user","content":"hi"}]}"#; assert!(blobs.put(r.at_ms, r.id, tw_store::Which::Request, body)); let rec = tw_store::Recorder::new( db, @@ -240,3 +247,56 @@ async fn a_bedrock_upstream_is_not_offered_a_replay_it_cannot_take() { assert_eq!(v["code"], "control.replay_bedrock", "{path}: {v}"); } } + +/// 一个把收到的请求体原样回给你的上游，顺手记下收到了什么。 +async fn echoing() -> (SocketAddr, Arc>) { + let seen: Arc> = Arc::default(); + let s = seen.clone(); + let app = axum::Router::new().route( + "/v1/messages", + axum::routing::post(move |body: String| { + let s = s.clone(); + async move { + *s.lock().unwrap() = body.clone(); + body + } + }), + ); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + (addr, seen) +} + +/// 存下来的请求（拦截档下存的是换过的那一份）里写着 1 号；这一次又找到一把密钥（比如 +/// 那之后加了一条规则，或者观察档下存的）。**新找到的拿 2 号**：拿 1 号的话，回答里的 +/// 1 号会被还原成这把新的，而它原本指的是另一样。 +#[tokio::test] +async fn a_replay_numbers_new_finds_after_the_placeholders_already_stored() { + const KEY: &str = "sk-ant-api03-REPLAYKEYAAAAAAAAAAAAAAAA"; + let (upstream, seen) = echoing().await; + let stored = format!( + r#"{{"model":"claude-sonnet-4-5","max_tokens":8,"messages":[{{"role":"user","content":"旧的 <>，新的 {KEY}"}}]}}"# + ); + let (_d, app) = app_with( + &format!( + "version: 1\nlisten:\n control:\n key: c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00\nclients:\n - name: 我\n key: tw-一把钥匙就够\nproviders:\n \ + - name: 本机\n base_url: http://{upstream}\n key: sk-x\n billing: free\n\ + security:\n redact:\n mode: enforce\n" + ), + stored.as_bytes(), + ); + + let v = replay(&app, "本机").await; + let sent = seen.lock().unwrap().clone(); + assert!( + sent.contains("旧的 <>，新的 <>"), + "{sent}" + ); + // 1 号没有原值，原样留着；2 号还原成新找到的那把（给人看的打了码） + let body = v["body"].as_str().unwrap(); + assert!( + body.contains("旧的 <>，新的 sk-an…AAAA"), + "{v}" + ); +} diff --git a/crates/tw-control/tests/stored_bodies.rs b/crates/tw-control/tests/stored_bodies.rs new file mode 100644 index 00000000..5932e3eb --- /dev/null +++ b/crates/tw-control/tests/stored_bodies.rs @@ -0,0 +1,387 @@ +//! 落盘的正文里没有原值。 +//! +//! 换掉、打码在落盘那一头做（`tw_gateway::bodies::BodyRecord::for_disk`）。这里把整条路 +//! 走一遍：真的网关、一个把收到的东西回显出来的假上游、存储层写进临时目录 —— 然后把 +//! 正文目录里的**每一个文件**读一遍，一个原值都不能有。拦截档、观察档、关闭各一遍， +//! 整包的回答和流式的回答各一个请求。 +//! +//! 读出这些正文的几条路也在这里走一遍：请求详情、按正文找、重放。 + +use std::net::SocketAddr; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use tower::ServiceExt; +use tw_control::{ConfigManager, ControlState}; + +/// 用户粘进对话里的三样：一把 API key、一个身份证号、连接串里的口令。三条都是出厂就开的规则 +const KEY: &str = "sk-ant-api03-USERSOWNKEYAAAAAAAAAAAAAA"; +const ID_NUMBER: &str = "11010519491231002X"; +const DB_PASSWORD: &str = "hunter2hunter2"; + +fn said() -> String { + format!("key {KEY}，身份证 {ID_NUMBER}，库 postgres://app:{DB_PASSWORD}@db.example.com/prod") +} + +/// 假上游：存下收到的 body，再把用户那句话原样回显 —— 模型确实会重复你给它的东西 +/// （拦截档下它看到的是占位符，观察档下是原值）。 +async fn echoing_upstream() -> (SocketAddr, Arc>>) { + let seen: Arc>> = Arc::default(); + let s = seen.clone(); + let app = axum::Router::new().route( + "/v1/messages", + axum::routing::post(move |body: bytes::Bytes| { + let s = s.clone(); + async move { + *s.lock().unwrap() = body.to_vec(); + let v: serde_json::Value = serde_json::from_slice(&body).unwrap_or_default(); + let text = v["messages"][0]["content"].as_str().unwrap_or("").to_string(); + let quoted = serde_json::to_string(&text).unwrap(); + if v["stream"] == true { + let sse = format!( + "event: message_start\ndata: {{\"type\":\"message_start\",\"message\":{{\"model\":\"claude-sonnet-4-5\",\"usage\":{{\"input_tokens\":10,\"output_tokens\":1}}}}}}\n\n\ + event: content_block_start\ndata: {{\"type\":\"content_block_start\",\"index\":0,\"content_block\":{{\"type\":\"text\",\"text\":\"\"}}}}\n\n\ + event: content_block_delta\ndata: {{\"type\":\"content_block_delta\",\"index\":0,\"delta\":{{\"type\":\"text_delta\",\"text\":{quoted}}}}}\n\n\ + event: content_block_stop\ndata: {{\"type\":\"content_block_stop\",\"index\":0}}\n\n\ + event: message_delta\ndata: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":\"end_turn\"}},\"usage\":{{\"output_tokens\":5}}}}\n\n\ + event: message_stop\ndata: {{\"type\":\"message_stop\"}}\n\n" + ); + axum::response::Response::builder() + .header("content-type", "text/event-stream") + .body(Body::from(sse)) + .unwrap() + } else { + axum::response::Response::builder() + .header("content-type", "application/json") + .body(Body::from(format!( + "{{\"type\":\"message\",\"model\":\"claude-sonnet-4-5\",\"content\":[{{\"type\":\"text\",\"text\":{quoted}}}],\"stop_reason\":\"end_turn\",\"usage\":{{\"input_tokens\":10,\"output_tokens\":5}}}}" + ))) + .unwrap() + } + } + }), + ); + let l = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = l.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(l, app).await.unwrap() }); + (addr, seen) +} + +struct World { + _dir: tempfile::TempDir, + blobs: std::path::PathBuf, + gw: SocketAddr, + seen: Arc>>, + store: Arc>, + app: axum::Router, +} + +/// 网关、存储层、控制面，**和 `twcore` 一样接**：网关交出原文，落盘之前换过、打过码。 +async fn world(mode: &str) -> World { + let (up, seen) = echoing_upstream().await; + let d = tempfile::tempdir().unwrap(); + let yaml = format!( + "version: 1\nlisten:\n control:\n key: c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00\n\ + clients:\n - name: 我\n key: tw-k\n\ + providers:\n - name: 中转\n base_url: http://{up}\n key: sk-upstream\n protocol: anthropic\n billing: free\n\ + security:\n redact:\n mode: {mode}\n" + ); + let p = d.path().join("config.yaml"); + std::fs::write(&p, &yaml).unwrap(); + let cfg: tw_config::Config = serde_yaml_ng::from_str(&yaml).unwrap(); + let gw = tw_gateway::AppState::new(cfg).unwrap(); + + let (sink, mut bodies) = tw_gateway::bodies::channel(); + gw.set_body_sink(sink); + let (tx, rx) = tokio::sync::mpsc::channel(1); + tokio::spawn(async move { + while let Some(b) = bodies.recv().await { + let disk = tokio::task::spawn_blocking(move || b.for_disk()) + .await + .unwrap(); + let which = match disk.kind { + tw_gateway::bodies::BodyKind::Request => tw_store::Which::Request, + tw_gateway::bodies::BodyKind::Response => tw_store::Which::Response, + }; + let stored = tw_store::StoredBody { + id: disk.id, + at_ms: disk.at_ms, + which, + body: disk.body, + original_len: disk.original_len, + }; + if tx.send(stored).await.is_err() { + return; + } + } + }); + let blobs = d.path().join("blobs"); + let store = tw_store::task::spawn( + tw_store::Recorder::new( + tw_store::Db::open(&d.path().join("data.db")).unwrap(), + tw_store::Blobs::new(blobs.clone()), + tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), + ), + gw.bus.subscribe(), + rx, + ); + let state = ControlState { + shutdown: Default::default(), + remote: Default::default(), + cfg: Arc::new(ConfigManager::new(p, gw.clone(), gw.bus.clone())), + gateway: gw.clone(), + store: Some(store.clone()), + started: std::time::Instant::now(), + price_updater: Default::default(), + chatgpt: Default::default(), + zai: Default::default(), + }; + let app = tw_control::router(state); + let addr = tw_gateway::serve(gw, ([127, 0, 0, 1], 0).into()) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + World { + _dir: d, + blobs, + gw: addr, + seen, + store, + app, + } +} + +/// 发一个请求，等它落了库、两份正文都落了盘。交回它的号和客户端拿到的回答 +async fn ask(w: &World, stream: bool) -> (i64, String) { + let before = w.store.lock().await.db().count().unwrap(); + let body = serde_json::json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": said()}], + }); + let got = reqwest::Client::builder() + .no_proxy() + .build() + .unwrap() + .post(format!("http://{}/v1/messages", w.gw)) + .header("x-api-key", "tw-k") + .header("content-type", "application/json") + .body(body.to_string()) + .send() + .await + .unwrap() + .text() + .await + .unwrap(); + let deadline = std::time::Instant::now() + Duration::from_secs(10); + loop { + let id = { + let g = w.store.lock().await; + (g.db().count().unwrap() > before).then(|| g.db().recent(None, 1).unwrap()[0].id) + }; + if let Some(id) = id + && files(&w.blobs) + .iter() + .filter(|(name, _)| name == &format!("{id}.req") || name == &format!("{id}.res")) + .count() + == 2 + { + return (id, got); + } + assert!( + std::time::Instant::now() < deadline, + "请求没有落库，或者正文没有落盘" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +/// 正文目录里的每一个文件：(文件名, 内容) +fn files(root: &std::path::Path) -> Vec<(String, String)> { + let mut out = Vec::new(); + let Ok(days) = std::fs::read_dir(root) else { + return out; + }; + for day in days.flatten() { + for f in std::fs::read_dir(day.path()) + .into_iter() + .flatten() + .flatten() + { + let text = String::from_utf8_lossy(&std::fs::read(f.path()).unwrap()).to_string(); + out.push((f.file_name().to_string_lossy().to_string(), text)); + } + } + out +} + +fn stored(w: &World, name: &str) -> String { + files(&w.blobs) + .into_iter() + .find(|(n, _)| n == name) + .unwrap_or_else(|| panic!("没有 {name}")) + .1 +} + +/// **这一条是这个文件存在的理由**：盘上没有一个文件带着原值。 +fn assert_nothing_real_on_disk(w: &World) { + let all = files(&w.blobs); + assert!(all.len() >= 4, "{all:?}"); + for (name, text) in &all { + for secret in [KEY, ID_NUMBER, DB_PASSWORD] { + assert!(!text.contains(secret), "{name} 里有原值 {secret}：{text}"); + } + } +} + +async fn call(app: &axum::Router, req: Request) -> (StatusCode, serde_json::Value) { + let r = app.clone().oneshot(req).await.unwrap(); + let status = r.status(); + let b = axum::body::to_bytes(r.into_body(), usize::MAX) + .await + .unwrap(); + (status, serde_json::from_slice(&b).unwrap_or_default()) +} + +async fn get(app: &axum::Router, path: &str) -> (StatusCode, serde_json::Value) { + call(app, Request::get(path).body(Body::empty()).unwrap()).await +} + +async fn post( + app: &axum::Router, + path: &str, + body: serde_json::Value, +) -> (StatusCode, serde_json::Value) { + call( + app, + Request::post(path) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await +} + +/// 拦截档：存下来的请求带着上游收到的那些占位符，回答里的占位符和它对得上号。这个请求 +/// 在路上没被改写过（同格式直通、没有改参数的规则），所以两份一个字节都不差 +#[tokio::test] +async fn under_enforce_the_disk_holds_what_the_upstream_got_and_nothing_real() { + let w = world("enforce").await; + let (whole, answer) = ask(&w, false).await; + // 客户端拿回来的照样是原值：落盘那一头换掉的不影响回程的还原 + assert!(answer.contains(KEY), "{answer}"); + let sent = String::from_utf8(w.seen.lock().unwrap().clone()).unwrap(); + assert_eq!(stored(&w, &format!("{whole}.req")), sent); + assert!( + sent.contains("key <>，身份证 <>，库 postgres://app:<>@db.example.com/prod"), + "{sent}" + ); + assert!( + stored(&w, &format!("{whole}.res")) + .contains("key <>，身份证 <>"), + ); + + let (streamed, answer) = ask(&w, true).await; + assert!(answer.contains(KEY), "{answer}"); + let sent = String::from_utf8(w.seen.lock().unwrap().clone()).unwrap(); + assert_eq!(stored(&w, &format!("{streamed}.req")), sent); + let res = stored(&w, &format!("{streamed}.res")); + assert!(res.contains("event: content_block_delta"), "{res}"); + assert!(res.contains("<>"), "{res}"); + + assert_nothing_real_on_disk(&w); +} + +/// 观察档：上游收到的是原值、回显的也是原值，盘上的打了码 —— 和安全日志里一种写法 +#[tokio::test] +async fn under_observe_the_upstream_sees_the_values_and_the_disk_does_not() { + let w = world("observe").await; + let (whole, _) = ask(&w, false).await; + let sent = String::from_utf8(w.seen.lock().unwrap().clone()).unwrap(); + assert!(sent.contains(KEY), "观察档动了发出去的那一份：{sent}"); + let req = stored(&w, &format!("{whole}.req")); + assert!( + req.contains( + "key sk-an…AAAA，身份证 …002X，库 postgres://app:hunte…ter2@db.example.com/prod" + ), + "{req}" + ); + assert!( + !req.contains("<(&req).expect("存下来的请求还是 JSON"); + let (streamed, _) = ask(&w, true).await; + assert!(stored(&w, &format!("{streamed}.res")).contains("sk-an…AAAA")); + + assert_nothing_real_on_disk(&w); +} + +/// 关着也一样：关掉的是出站的检测和替换，不是落盘前的打码 +#[tokio::test] +async fn with_redaction_off_the_disk_still_holds_nothing_real() { + let w = world("off").await; + ask(&w, false).await; + ask(&w, true).await; + assert_nothing_real_on_disk(&w); +} + +/// 读出这些正文的几条路：详情、按正文找、重放。都照常工作，也都碰不到原值 +#[tokio::test] +async fn the_detail_search_and_replay_work_on_what_was_stored() { + let w = world("enforce").await; + let (id, _) = ask(&w, false).await; + + // 详情：存下来的就是给人看的样子，读的时候再打一遍码什么都不变 + let (st, v) = get(&w.app, &format!("/request/{id}")).await; + assert_eq!(st, StatusCode::OK, "{v}"); + let req = v["request_body"]["text"].as_str().unwrap(); + assert_eq!(req, stored(&w, &format!("{id}.req"))); + assert_eq!(v["request_body"]["truncated"], false, "{v}"); + assert_eq!(v["request_body"]["original_len"], req.len(), "{v}"); + assert!(!v.to_string().contains(KEY), "{v}"); + + // 按正文找：找得到那句话，摘录里是占位符；只在原值里出现的字找不到 + let (st, v) = post( + &w.app, + "/history/search", + serde_json::json!({"q": "身份证", "content": true}), + ) + .await; + assert_eq!(st, StatusCode::OK, "{v}"); + let hits = v["hits"].as_array().unwrap(); + assert!(!hits.is_empty(), "{v}"); + assert!( + hits[0]["after"] + .as_str() + .unwrap() + .contains("<>"), + "{v}" + ); + let (_, v) = post( + &w.app, + "/history/search", + serde_json::json!({"q": "USERSOWNKEY", "content": true}), + ) + .await; + assert!(v["hits"].as_array().unwrap().is_empty(), "{v}"); + + // 重放：发出去的是存下来的那一份，占位符原样发，回显里也原样留着 —— 没有原值可还原 + let (st, v) = post( + &w.app, + "/replay/run", + serde_json::json!({"id": id, "provider": "中转"}), + ) + .await; + assert_eq!(st, StatusCode::OK, "{v}"); + let replayed = String::from_utf8(w.seen.lock().unwrap().clone()).unwrap(); + assert_eq!(replayed, stored(&w, &format!("{id}.req"))); + assert!(!replayed.contains(KEY), "{replayed}"); + assert!( + v["body"].as_str().unwrap().contains("key <>"), + "{v}" + ); +} diff --git a/crates/tw-control/tests/transcript.rs b/crates/tw-control/tests/transcript.rs new file mode 100644 index 00000000..8a11dde4 --- /dev/null +++ b/crates/tw-control/tests/transcript.rs @@ -0,0 +1,222 @@ +//! 对话记录端点（`GET /sessions/{id}/transcript`）的形状。 +//! +//! 四种格式怎么读、前后两个请求怎么比在 tw-store 里测；这里盯的是**界面拿到的那份 +//! JSON**，一个字段一个字段地钉住：桌面端照着它写，改了哪个字段名那边就悄悄读不到了。 + +use std::sync::Arc; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use serde_json::{Value, json}; +use tower::ServiceExt; +use tw_control::{ConfigManager, ControlState}; +use tw_store::Which; + +const BASE: &str = "version: 1\nlisten:\n control:\n key: c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00\nclients:\n - name: 我\n key: tw-一把钥匙就够\nproviders:\n - name: 官方\n base_url: https://api.anthropic.com\n key: sk-x\n"; + +/// Claude Code 的一轮 +fn turn(id: i64) -> tw_store::db::RequestRow { + tw_store::db::RequestRow { + session_log_bytes: None, + key_masked: None, + peer: None, + id, + at_ms: 1_790_000_000_000 + id, + client: "claude-code".into(), + client_hint: None, + session: Some("s1".into()), + provider: "官方".into(), + model: "claude-sonnet-4-5".into(), + sent_model: "claude-sonnet-4-5".into(), + answered_model: None, + path: "/v1/messages".into(), + status: Some(200), + ttfb_ms: Some(100), + ttft_ms: None, + duration_ms: Some(200), + tokens_per_sec: None, + bytes: Some(10), + input_tokens: Some(50), + output_tokens: Some(20), + cache_read_tokens: None, + cache_write_tokens: None, + input_estimate: None, + cost_micros: Some(1_000), + cost_estimated: false, + error: None, + local: false, + cancelled: false, + routing: None, + billing: tw_api::Billing::PerToken, + cache_saved_micros: None, + price_source: None, + translated: None, + } +} + +/// 一个带着请求库和正文的控制面。`bodies`：哪一条、请求体、回答（没存下的是 None） +fn app( + rows: &[tw_store::db::RequestRow], + bodies: &[(i64, Option, Option)], +) -> (tempfile::TempDir, axum::Router) { + let d = tempfile::tempdir().unwrap(); + let p = d.path().join("config.yaml"); + std::fs::write(&p, BASE).unwrap(); + let db = tw_store::Db::open(&d.path().join("data.db")).unwrap(); + for r in rows { + db.insert(r).unwrap(); + } + let blobs = tw_store::Blobs::new(d.path().join("blobs")); + for (id, request, response) in bodies { + let at = rows.iter().find(|r| r.id == *id).unwrap().at_ms; + for (which, body) in [(Which::Request, request), (Which::Response, response)] { + if let Some(b) = body { + assert!(blobs.put(at, *id, which, b.to_string().as_bytes())); + } + } + } + let rec = tw_store::Recorder::new( + db, + blobs, + tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), + ); + let cfg: tw_config::Config = serde_yaml_ng::from_str(BASE).unwrap(); + let gw = tw_gateway::AppState::new(cfg).unwrap(); + let bus = gw.bus.clone(); + let state = ControlState { + shutdown: Default::default(), + remote: Default::default(), + cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), + gateway: gw, + store: Some(Arc::new(tokio::sync::Mutex::new(rec))), + started: std::time::Instant::now(), + price_updater: Default::default(), + chatgpt: Default::default(), + zai: Default::default(), + }; + (d, tw_control::router(state)) +} + +async fn get(app: &axum::Router, path: &str) -> (StatusCode, Value) { + let r = app + .clone() + .oneshot(Request::builder().uri(path).body(Body::empty()).unwrap()) + .await + .unwrap(); + let st = r.status(); + let b = axum::body::to_bytes(r.into_body(), 1 << 20).await.unwrap(); + (st, serde_json::from_slice(&b).unwrap()) +} + +/// 每一种块、每一个字段都在：请求号是字符串，可以为空的字段写成 null 而不是省掉， +/// 块和缺口是界面认的那几个词。工具结果里的密钥打了码，图片只有类型和大小 +#[tokio::test] +async fn the_transcript_is_the_shape_the_desktop_app_reads() { + let png = "QUJD".repeat(2); + let first = json!({"model": "claude-sonnet-4-5", "system": "sys", "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "hi"}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}} + ]} + ]}); + let answer = json!({"type": "message", "content": [ + {"type": "redacted_thinking", "data": "abc"}, + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}, + {"type": "text", "text": "hello"}, + {"type": "tool_use", "id": "t1", "name": "Read", "input": {"p": 1}} + ]}); + let second = json!({"model": "claude-sonnet-4-5", "system": "sys", "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "hi"}, + {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}} + ]}, + {"role": "assistant", "content": answer["content"]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", "content": [ + {"type": "text", "text": "key sk-ant-api03-SECRETSECRETSECRET"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": png}} + ]}]} + ]}); + let (_d, app) = app( + &[turn(1), turn(2)], + &[(1, Some(first), Some(answer)), (2, Some(second), None)], + ); + + let (st, t) = get(&app, "/sessions/s1/transcript").await; + assert_eq!(st, StatusCode::OK, "{t}"); + assert_eq!( + t, + json!({ + "session": "s1", + "system": "sys", + "turns": [ + { + "id": "1", + "restart": false, + "system_changed": null, + "input": [{"role": "user", "parts": [ + {"kind": "text", "text": "hi"}, + {"kind": "image", "media_type": null, "bytes": null} + ]}], + "output": [ + {"kind": "thinking", "text": ""}, + {"kind": "other", "label": "server_tool_use"}, + {"kind": "text", "text": "hello"}, + {"kind": "tool_call", "id": "t1", "name": "Read", "input": "{\"p\":1}"} + ], + "gaps": [] + }, + { + "id": "2", + "restart": false, + "system_changed": null, + "input": [{"role": "tool", "parts": [ + {"kind": "tool_result", "call_id": "t1", "text": "key sk-an…CRET", "is_error": false}, + {"kind": "image", "media_type": "image/png", "bytes": 6} + ]}], + "output": [], + "gaps": ["response_missing"] + } + ] + }) + ); + + // 和会话详情是同样的请求、同样的顺序 + let (_, detail) = get(&app, "/sessions/s1").await; + let ids: Vec = detail["turns"] + .as_array() + .unwrap() + .iter() + .map(|x| x["id"].to_string()) + .collect(); + assert_eq!(ids, ["1", "2"]); +} + +/// 没有这次会话：404，和会话详情同一个码 +#[tokio::test] +async fn an_unknown_session_is_not_found() { + let (_d, app) = app(&[turn(1)], &[]); + let (st, body) = get(&app, "/sessions/nope/transcript").await; + assert_eq!(st, StatusCode::NOT_FOUND); + assert_eq!(body["code"], "control.session_not_found"); + assert_eq!(body["args"]["id"], "nope"); + let (st, body) = get(&app, "/sessions/nope").await; + assert_eq!(st, StatusCode::NOT_FOUND); + assert_eq!(body["code"], "control.session_not_found"); +} + +/// 正文都清掉了的会话照样有它的每一轮，每一轮说出缺了什么 +#[tokio::test] +async fn a_session_whose_bodies_are_gone_still_lists_its_turns() { + let (_d, app) = app(&[turn(1), turn(2)], &[]); + let (st, t) = get(&app, "/sessions/s1/transcript").await; + assert_eq!(st, StatusCode::OK); + assert_eq!(t["system"], Value::Null); + for x in t["turns"].as_array().unwrap() { + assert_eq!( + x["gaps"], + json!(["request_missing", "response_missing"]), + "{x}" + ); + assert_eq!(x["input"], json!([])); + } +} diff --git a/crates/tw-gateway/src/bodies.rs b/crates/tw-gateway/src/bodies.rs index d1f62bf1..82610968 100644 --- a/crates/tw-gateway/src/bodies.rs +++ b/crates/tw-gateway/src/bodies.rs @@ -4,15 +4,52 @@ //! 转发那条路上 —— 一次慢磁盘写就会变成一次慢请求，而观测永远不该有这 //! 个权力。通道满了就丢：丢的是一条观测记录，而等它是在惩罚 //! 真实用户。 +//! +//! # 落盘的不是原文 +//! +//! 交出去的是原文，落盘之前由收的那一头换掉、打码（[`BodyRecord::for_disk`]）： +//! 那要把整份正文扫好几遍，不该在转发那条路上做。**脱敏规则认得出的值不会原样进磁盘**， +//! 哪一档都一样，关着也一样： +//! +//! - 拦截档下，请求里的值换成**发给上游的那个占位符**：存下来的就是上游收到的那一份， +//! 回答里出现的占位符和它对得上号（见 [`crate::guard::look`]） +//! - 别的一律打码，和安全日志里报的是同一种写法 +//! - 最后整段再按形状打一遍码，和读的时候是同一个函数（[`tw_secret::mask_body`]）， +//! 兜住规则没认出来的 +//! +//! 以前存的是客户端发来的原文，密钥只在读出来的时候才打码：磁盘上躺着的一直是真值。 + +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use bytes::Bytes; +use tw_guard::redact::replace::{Ledger, Scheme}; +use tw_guard::redact::rules::{Kind, Rule, RuleSet}; -/// 响应体最多留多少。 +/// 打码要多看的那一截。 /// -/// **和落盘的上限（4 MB）是两个数。**这个是「在内存里攒着等着交出去」的 -/// 量，同时在飞的请求越多它乘得越狠 —— 256 KB × 32 并发是 8 MB，可以； -/// 4 MB × 32 是 128 MB，不行。而详情抽屉要看的东西，开头这些字节里全有。 -pub const RESPONSE_TAP_MAX: usize = 256 * 1024; +/// 一把密钥正好跨在截断处的话，只看前半截认不出它（私钥要见到结尾那一行才算），存下来 +/// 的开头里就留着它的前半截。多看一截，跨在截断处的那一个整个认得出、整个换掉，然后再截。 +/// 64 KB 比任何一种认得出的凭据都长得多：一把 8192 位的 RSA 私钥不到 7 KB。 +pub const MARGIN: usize = 64 * 1024; + +/// 交去落盘的一份最多带多少字节：存下来的那 [`tw_api::BODY_MAX`]，加上打码要多看的 +/// [`MARGIN`]。更长的请求体只交开头这么多，响应体也只攒这么多。 +pub const WINDOW: usize = tw_api::BODY_MAX + MARGIN; + +/// 响应体最多攒多少（[`WINDOW`]）。 +/// +/// **存下来的和请求体一样长**，都是 [`tw_api::BODY_MAX`]（4 MB），多攒的那一截只给打码 +/// 看。以前这里是 256 KB，比存储层肯收的少十几倍：一个长回答的后半截 —— 最后那几个工具 +/// 调用、停止的原因 —— 总在被扔掉的那一段里。 +/// +/// **内存**：每个在飞的回答一份，**按实际长度长**（见 [`ResponseTap::feed`]），攒到上限 +/// 为止。绝大多数回答几 KB 到几百 KB；攒得满 4 MB 的是很长的流 —— SSE 里每几个字就 +/// 包着一帧，几万 token 的回答就有几 MB。32 个同时在流、个个都过了 4 MB，是 130 MB 上下， +/// 流一结束就还回去。观测层没起来时一个字节都不攒（`crate::ending::Ending::feed`）。交出去 +/// 之后由 [`QUEUED_MAX`] 管着。 +pub const RESPONSE_TAP_MAX: usize = WINDOW; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum BodyKind { @@ -20,34 +57,225 @@ pub enum BodyKind { Response, } +/// 落盘之前怎么处理一份正文。 +/// +/// 带着做这件事要的东西：这个请求开始时生效的那套规则（换了配置，已经在路上的照旧）， +/// 和拦截档下这个请求的账本（见 [`crate::guard::look`]）。 +#[derive(Clone)] +pub struct Redaction { + pub rules: Arc, + /// 原值 → 发给上游的占位符。拦截档下才有东西 + pub ledger: Ledger, +} + +/// **不打印账本**：里面是原值。 +impl std::fmt::Debug for Redaction { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Redaction") + .field("replaced", &self.ledger.len()) + .finish_non_exhaustive() + } +} + +impl Default for Redaction { + /// 没人交代的时候：出厂的那套规则，不换、只打码。**宁可多打。** + fn default() -> Self { + Self { + rules: Arc::new(RuleSet::defaults()), + ledger: Ledger::new(Scheme::SECRET), + } + } +} + +impl Redaction { + /// 换过、打过码的样子。 + /// + /// 两道：脱敏规则认出的每一处，账本里有的换成那个占位符（上游收到的就是它），没有的 + /// 打码；然后整段再按形状打一遍（[`tw_secret::mask_body`]）兜住规则没认出来的。 + /// + /// **读的时候还会再打一遍**，打第二遍不再改动什么（`mask_body` 认得自己打过的码）。 + pub fn apply(&self, text: &str) -> String { + // 我们自己的占位符不算（见 `crate::guard::hits`）：回答里回显的、重放过来的，原样存 + let hits = crate::guard::hits(text, &self.rules); + let sent: HashMap<&str, &str> = self.ledger.replacements().collect(); + let mut out = String::with_capacity(text.len()); + let mut at = 0; + for h in &hits { + out.push_str(&text[at..h.bytes.start]); + let value = &text[h.bytes.clone()]; + match sent.get(value) { + Some(placeholder) => out.push_str(placeholder), + None => out.push_str(&masked(&h.rule, value)), + } + at = h.bytes.end; + } + out.push_str(&text[at..]); + tw_secret::mask_body(&out) + } +} + +/// 一个认出来、没换成占位符的值存成什么样。 +/// +/// 和安全日志里报的是同一种写法（[`tw_guard::redact::rules::masked`]），只有内网地址和 +/// 内部域名不同：日志里它们原样报（打了码就说不出是哪台机器），正文里一样打掉 —— 打开 +/// 这两条规则的人，就是不想让它们留在别处的人。 +/// +/// **反斜杠去掉**：一个值里带着转义时，打码可能正好留下半个，存下来的那份就不再是 JSON。 +fn masked(rule: &Rule, value: &str) -> String { + let m = match rule.kind() { + Kind::Internal => tw_secret::mask_secret(value), + _ => tw_guard::redact::rules::masked(rule, value), + }; + m.replace('\\', "") +} + /// 一份要存起来的 body。 #[derive(Debug)] pub struct BodyRecord { pub id: u64, pub at_ms: i64, pub kind: BodyKind, + /// 原文。交进通道的最多 [`WINDOW`] 字节 pub body: Bytes, /// 原始长度。**截断了要能说出来** —— 不说的话，用户会以为请求本身 /// 就长这样 pub original_len: usize, + /// 落盘之前怎么换、怎么打码 + pub redaction: Redaction, + /// 占着的那份额度（见 [`QUEUED_MAX`]） + held: Option, +} + +impl BodyRecord { + pub fn new( + id: u64, + at_ms: i64, + kind: BodyKind, + body: Bytes, + original_len: usize, + redaction: Redaction, + ) -> Self { + Self { + id, + at_ms, + kind, + body, + original_len, + redaction, + held: None, + } + } + + /// 落盘的那一份：换过、打过码（[`Redaction::apply`]），和要记下的原本长度。 + /// + /// **在阻塞线程上调**：一份 4 MB 的正文要扫好几遍。 + /// + /// 不是 UTF-8 的照读的时候的办法转成文字（坏字节换成 U+FFFD）：读的人看到的本来就是 + /// 那样，而二进制的正文里没有能看的东西。 + pub fn for_disk(self) -> ForDisk { + let window = &self.body[..self.body.len().min(WINDOW)]; + let body = self.redaction.apply(&String::from_utf8_lossy(window)); + // 截过的（交来的只是开头）报原本的长度。没截过的就是换过、打过码的这一份的长度： + // 比存储层的上限还长的，由存储层截、由它记下（`tw_store::Blobs::put_with_len`） + let whole = self.original_len.max(self.body.len()); + let original_len = if whole > window.len() { + whole + } else { + body.len() + }; + ForDisk { + id: self.id, + at_ms: self.at_ms, + kind: self.kind, + body: Bytes::from(body), + original_len, + held: self.held, + } + } +} + +/// 落盘的那一份（[`BodyRecord::for_disk`]）。 +#[derive(Debug)] +pub struct ForDisk { + pub id: u64, + pub at_ms: i64, + pub kind: BodyKind, + /// 换过、打过码的那一份。可能比 [`tw_api::BODY_MAX`] 长一点，由存储层截 + pub body: Bytes, + /// 原始长度，交给 `tw_store::Blobs::put_with_len` + pub original_len: usize, + /// 还占着的额度。**拿着它，直到存储层收下这一份**：交接途中的那一份也是攒在内存里的 + pub held: Option, } -/// 往哪儿交。`None` 表示观测层没起来 —— 那时什么都不做，转发照旧。 -pub type BodySender = tokio::sync::mpsc::Sender; +/// 一份正文占着的额度（见 [`QUEUED_MAX`]）。**被丢掉时还回去**：写完了、通道满了、 +/// 收的那一头不在了，都一样。 +#[derive(Debug)] +pub struct Held { + n: usize, + queued: Arc, +} + +impl Drop for Held { + fn drop(&mut self) { + self.queued.fetch_sub(self.n, Ordering::Relaxed); + } +} + +/// 往哪儿交，和交出去还没写完的有多少字节。`None` 表示观测层没起来 —— 那时什么都不做， +/// 转发照旧。 +#[derive(Debug, Clone)] +pub struct BodySink { + tx: tokio::sync::mpsc::Sender, + queued: Arc, +} + +/// 建一条通道：发的那一头给网关（`AppState::set_body_sink`），收的那一头交给落盘。 +pub fn channel() -> (BodySink, tokio::sync::mpsc::Receiver) { + let (tx, rx) = tokio::sync::mpsc::channel(CHANNEL_CAP); + let sink = BodySink { + tx, + queued: Arc::default(), + }; + (sink, rx) +} -/// 通道容量。 +/// 通道最多攒几份。**按字节的上限在 [`QUEUED_MAX`]**，这个数管住的是一大堆很小的正文。 /// -/// 攒不下就丢。**这个数字要小**：它乘上单条 256 KB 就是内存上限， -/// 而攒着一堆等着写盘的 body 本身就说明磁盘跟不上，那时留着它们也没用。 +/// 攒不下就丢：攒着一堆等着写盘的 body 本身就说明磁盘跟不上，那时留着它们也没用。 pub const CHANNEL_CAP: usize = 64; +/// 交出去、还没写完的正文合计最多多少字节。超了就丢新来的。 +/// +/// **按字节记，不按份数。**一份正文最多 4 MB 出头，按份数定上限的话，64 份就是 256 MB +/// 攒在内存里等磁盘；而平时攒着的几乎都是几 KB 的小正文，把份数压小又会在忙的时候白白 +/// 丢掉它们。单独一份总放得下：[`WINDOW`] 比它小得多。 +pub const QUEUED_MAX: usize = 32 * 1024 * 1024; + /// 交一份出去。**满了就丢，绝不等待。** -pub fn offer(tx: &Option, rec: BodyRecord) { - let Some(tx) = tx else { return }; - if tx.try_send(rec).is_err() { - // 不记日志：这条路上每个请求都会走一次，而写盘跟不上的时候 - // 日志会跟着刷屏 —— 那才是真的把事情变糟。 +/// +/// 比 [`WINDOW`] 长的只交开头 —— **拷出来**，不切片：切片拽着整个请求体（最大 256 MB） +/// 一起等在通道里。 +pub fn offer(sink: &Option, mut rec: BodyRecord) { + let Some(s) = sink else { return }; + if rec.body.len() > WINDOW { + rec.body = Bytes::copy_from_slice(&rec.body[..WINDOW]); } + let n = rec.body.len(); + // 先占额度，占不下就丢 + if s.queued.fetch_add(n, Ordering::Relaxed) + n > QUEUED_MAX { + s.queued.fetch_sub(n, Ordering::Relaxed); + return; + } + rec.held = Some(Held { + n, + queued: s.queued.clone(), + }); + // 交不出去的那一份（满了、收的那一头不在了）就地丢掉，额度跟着还回去。 + // + // 不记日志：这条路上每个请求都会走一次，而写盘跟不上的时候 + // 日志会跟着刷屏 —— 那才是真的把事情变糟。 + let _ = s.tx.try_send(rec); } /// 一边流一边攒响应体，**攒到上限就停**。 @@ -68,7 +296,15 @@ impl ResponseTap { if room == 0 { return; } - self.buf.extend_from_slice(&chunk[..chunk.len().min(room)]); + let take = chunk.len().min(room); + // **按实际长度长，到上限为止。**交给 Vec 自己翻倍的话，攒到 4 MB 出头时一下就要 + // 8 MB + let need = self.buf.len() + take; + if need > self.buf.capacity() { + let want = (self.buf.capacity() * 2).max(need).min(RESPONSE_TAP_MAX); + self.buf.reserve_exact(want - self.buf.len()); + } + self.buf.extend_from_slice(&chunk[..take]); } /// 攒到的那部分，以及**原始的总长度**。 @@ -81,6 +317,8 @@ impl ResponseTap { mod tests { use super::*; + const KEY: &str = "sk-ant-api03-AAAAAAAAAAAAAAAAAAAAAAAAAA"; + #[test] fn a_short_response_is_kept_whole() { let mut t = ResponseTap::new(); @@ -96,12 +334,20 @@ mod tests { // **截断了要能说出来** —— 不说的话，用户会以为响应本身就这么长。 let mut t = ResponseTap::new(); let chunk = vec![b'x'; 100 * 1024]; - for _ in 0..10 { + for _ in 0..50 { t.feed(&chunk); } + assert!(t.buf.capacity() <= RESPONSE_TAP_MAX, "攒满时多占了内存"); let (b, total) = t.finish(); assert_eq!(b.len(), RESPONSE_TAP_MAX); - assert_eq!(total, 10 * 100 * 1024, "原始长度没记住"); + assert_eq!(total, 50 * 100 * 1024, "原始长度没记住"); + } + + #[test] + fn the_tap_keeps_as_much_as_the_store_does_and_a_little_more() { + // 存下来的回答和请求一样长；多攒的那一截只给打码看 + assert_eq!(RESPONSE_TAP_MAX, tw_api::BODY_MAX + MARGIN); + assert_eq!(tw_api::BODY_MAX, 4 * 1024 * 1024); } #[test] @@ -115,27 +361,229 @@ mod tests { assert_eq!(total, RESPONSE_TAP_MAX + 5); } + fn record(kind: BodyKind, body: &str, redaction: Redaction) -> BodyRecord { + BodyRecord::new( + 1, + 0, + kind, + Bytes::from(body.to_string()), + body.len(), + redaction, + ) + } + + fn written(rec: BodyRecord) -> String { + String::from_utf8(rec.for_disk().body.to_vec()).unwrap() + } + + /// 拦截档下这个请求的账本：和 `start` 里一样，按客户端原文编号 + fn enforced(body: &str) -> Redaction { + let rules = RuleSet::defaults(); + let (_, ledger) = + crate::guard::look(tw_config::SecurityMode::Enforce, &rules, body.as_bytes()); + Redaction { + rules: Arc::new(rules), + ledger, + } + } + + #[test] + fn under_enforce_the_stored_request_carries_what_the_upstream_got() { + let body = format!( + r#"{{"model":"m","messages":[{{"role":"user","content":"我的 key 是 {KEY}，身份证 11010519491231002X"}}]}}"# + ); + let r = enforced(&body); + let (sent, _) = crate::guard::replace( + tw_config::SecurityMode::Enforce, + &r.rules, + Bytes::from(body.clone()), + &r.ledger, + ); + let stored = written(record(BodyKind::Request, &body, r)); + assert!(!stored.contains(KEY), "{stored}"); + // 一个字节都不差：存下来的就是发给上游的那一份 + assert_eq!(stored.as_bytes(), &sent[..]); + assert!( + stored.contains("<>，身份证 <>"), + "{stored}" + ); + serde_json::from_str::(&stored).expect("存下来的还是 JSON"); + } + + #[test] + fn without_enforce_what_the_rules_recognize_is_masked_before_it_is_written() { + // 观察档、关着的时候上游收到的是原值，存下来的打码 —— 和安全日志里一种写法 + let body = format!( + r#"{{"messages":[{{"role":"user","content":"key {KEY} 库 postgres://app:hunter2hunter2@db/x 卡 6222 0212 3456 7894"}}]}}"# + ); + let stored = written(record(BodyKind::Request, &body, Redaction::default())); + for secret in [KEY, "hunter2hunter2", "6222 0212 3456 7894"] { + assert!(!stored.contains(secret), "{secret} 原样进了磁盘：{stored}"); + } + assert!(stored.contains("key sk-an…AAAA"), "{stored}"); + assert!( + stored.contains("postgres://app:hunte…ter2@db/x"), + "{stored}" + ); + assert!(stored.contains("卡 …7894"), "{stored}"); + assert!( + !stored.contains("<(&stored).expect("存下来的还是 JSON"); + // 读的时候再打一遍，什么都不变 + assert_eq!(tw_secret::mask_body(&stored), stored); + } + + #[test] + fn what_the_rules_miss_is_masked_by_its_shape() { + // 网关自己的钥匙（`tw-`）没有一条脱敏规则认：读的时候那一道兜住它，写的时候也一样 + let body = + r#"{"messages":[{"role":"user","content":"钥匙 tw-0123456789abcdef0123456789"}]}"#; + let stored = written(record(BodyKind::Request, body, Redaction::default())); + assert!(!stored.contains("0123456789abcdef"), "{stored}"); + } + + #[test] + fn a_response_keeps_its_placeholders_and_loses_any_real_value() { + // 拦截档：上游回显的是占位符，原样存。它另外说出来的一把密钥（不在账本里）打码 + let request = format!(r#"{{"messages":[{{"role":"user","content":"{KEY}"}}]}}"#); + let r = enforced(&request); + let other = "ghp_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; + let answer = format!( + "event: content_block_delta\ndata: {{\"delta\":{{\"text\":\"你给的是 <>，另一把是 {other}\"}}}}\n\n" + ); + let stored = written(record(BodyKind::Response, &answer, r.clone())); + assert!(stored.contains("你给的是 <>"), "{stored}"); + assert!(!stored.contains(other), "{stored}"); + assert!(stored.contains("另一把是 ghp_B…BBBB"), "{stored}"); + + // 观察档：上游看见过原值，回显出来的也是原值 + let echoed = format!("{{\"text\":\"你给的是 {KEY}\"}}"); + let stored = written(record(BodyKind::Response, &echoed, Redaction::default())); + assert!(!stored.contains(KEY), "{stored}"); + } + + #[test] + fn a_secret_across_the_cut_is_masked_whole_before_the_cut() { + // 一把私钥正好跨在 4 MB 处：只看截下来的开头认不出它（要见到 END 才算） + let pem = "-----BEGIN PRIVATE KEY-----\\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC7\\n-----END PRIVATE KEY-----"; + let lead = tw_api::BODY_MAX - 40; + let body = format!( + r#"{{"content":"{}{pem}"}}"#, + "x".repeat(lead - r#"{"content":""#.len()) + ); + let disk = record(BodyKind::Request, &body, Redaction::default()).for_disk(); + let text = String::from_utf8(disk.body.to_vec()).unwrap(); + assert!(!text.contains("MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC7")); + assert!( + !text.contains("-----BEGIN PRIVATE KEY-----\\nMIIE"), + "{}", + &text[lead - 10..] + ); + // 存储层截到 4 MB 时只剩打过码的那一份的开头 + assert!(disk.body.len() < tw_api::BODY_MAX + 100); + } + + #[test] + fn the_length_written_down_says_whether_it_was_cut() { + // 没截过的：换过、打过码的这一份多长就是多长，存储层不会当它截过 + let body = format!(r#"{{"content":"{KEY}"}}"#); + let disk = record(BodyKind::Request, &body, Redaction::default()).for_disk(); + assert_eq!(disk.original_len, disk.body.len()); + assert!(disk.original_len < body.len(), "打码之后变短了"); + + // 交来的只是开头（响应体攒到了上限）：报原本的长度 + let mut r = record(BodyKind::Response, "data: {}\n\n", Redaction::default()); + r.original_len = 9 * 1024 * 1024; + assert_eq!(r.for_disk().original_len, 9 * 1024 * 1024); + + // 整份交来、却比窗口长：只看窗口里的，报整份的长度 + let huge = "x".repeat(WINDOW + 10); + let disk = record(BodyKind::Request, &huge, Redaction::default()).for_disk(); + assert_eq!(disk.body.len(), WINDOW); + assert_eq!(disk.original_len, WINDOW + 10); + } + + #[test] + fn a_body_that_is_not_text_is_stored_the_way_it_reads() { + let raw = Bytes::from(vec![0xff, b'a', 0xfe]); + let disk = + BodyRecord::new(1, 0, BodyKind::Response, raw, 3, Redaction::default()).for_disk(); + assert_eq!(&disk.body[..], "\u{fffd}a\u{fffd}".as_bytes()); + } + + #[test] + fn the_ledger_never_shows_up_in_a_debug_print() { + let r = enforced(&format!("{{\"k\":\"{KEY}\"}}")); + let dump = format!("{r:?}"); + assert!(!dump.contains("AAAAAAAAAAAA"), "{dump}"); + } + #[tokio::test] async fn offering_into_a_full_channel_drops_rather_than_waits() { // **等它是在惩罚真实用户。**观测永远不该有让请求变慢的权力。 - let (tx, _rx) = tokio::sync::mpsc::channel(1); - let tx = Some(tx); - let rec = || BodyRecord { - id: 1, - at_ms: 0, - kind: BodyKind::Request, - body: Bytes::from_static(b"x"), - original_len: 1, + let (sink, _rx) = channel(); + let tx = Some(sink.clone()); + let rec = || { + BodyRecord::new( + 1, + 0, + BodyKind::Request, + Bytes::from_static(b"x"), + 1, + Redaction::default(), + ) }; - // 通道容量是 1，塞十次不该挂住 + // 塞的比通道的份数多，不该挂住 let start = std::time::Instant::now(); - for _ in 0..10 { + for _ in 0..CHANNEL_CAP * 2 { offer(&tx, rec()); } assert!( start.elapsed() < std::time::Duration::from_millis(50), "offer 挂住了" ); + // 丢掉的那些占的额度还回去了：攒着的只有收下的那些 + assert_eq!(sink.queued.load(Ordering::Relaxed), CHANNEL_CAP); + } + + #[tokio::test] + async fn what_waits_to_be_written_is_capped_in_bytes_and_given_back_once_written() { + let (tx, mut rx) = channel(); + let sink = Some(tx.clone()); + let big = || { + BodyRecord::new( + 1, + 0, + BodyKind::Response, + Bytes::from(vec![b'x'; WINDOW + 1000]), + WINDOW + 1000, + Redaction::default(), + ) + }; + for _ in 0..20 { + offer(&sink, big()); + } + // 交进去的每一份都只带窗口那么长，合计不过上限 + let fits = QUEUED_MAX / WINDOW; + let mut got = Vec::new(); + while let Ok(r) = rx.try_recv() { + assert_eq!(r.body.len(), WINDOW); + got.push(r); + } + assert_eq!(got.len(), fits); + assert_eq!(tx.queued.load(Ordering::Relaxed), fits * WINDOW); + // 写完一份（落盘的那一份被丢掉）就还一份 + let disk = got.pop().unwrap().for_disk(); + assert_eq!(tx.queued.load(Ordering::Relaxed), fits * WINDOW); + drop(disk); + assert_eq!(tx.queued.load(Ordering::Relaxed), (fits - 1) * WINDOW); + drop(got); + assert_eq!(tx.queued.load(Ordering::Relaxed), 0); + // 腾出地方就又收了 + offer(&sink, big()); + assert!(rx.try_recv().is_ok()); } #[tokio::test] @@ -143,13 +591,14 @@ mod tests { // 观测层没起来时，转发照旧。 offer( &None, - BodyRecord { - id: 1, - at_ms: 0, - kind: BodyKind::Request, - body: Bytes::new(), - original_len: 0, - }, + BodyRecord::new( + 1, + 0, + BodyKind::Request, + Bytes::new(), + 0, + Redaction::default(), + ), ); } } diff --git a/crates/tw-gateway/src/ending.rs b/crates/tw-gateway/src/ending.rs index 1464616f..efd4fa85 100644 --- a/crates/tw-gateway/src/ending.rs +++ b/crates/tw-gateway/src/ending.rs @@ -22,7 +22,7 @@ use std::time::Instant; -use crate::bodies::{BodyKind, BodyRecord, BodySender, ResponseTap}; +use crate::bodies::{BodyKind, BodyRecord, BodySink, Redaction, ResponseTap}; use tw_dialect::convert::Reader; use tw_dialect::ir; use tw_dialect::usage::{Sniffer, Usage}; @@ -47,7 +47,9 @@ pub struct Ending { started: Instant, /// 请求开始的时刻。响应体按它归档，和请求体那一份对得上 at_ms: i64, - sink: Option, + sink: Option, + /// 响应体落盘之前怎么换、怎么打码：和请求体那一份是同一套（见 [`Ending::redact_with`]） + redaction: Option, /// 上游的响应头。**没到的时候客户端就走了的，没有状态码可报** —— 那时 /// 报一个 0 或者 499，都是在编 status: Option, @@ -117,7 +119,7 @@ impl Ending { model: String, started: Instant, at_ms: i64, - sink: Option, + sink: Option, ) -> Self { Self { bus, @@ -126,6 +128,7 @@ impl Ending { started, at_ms, sink, + redaction: None, status: None, bytes: 0, sniffer: Sniffer::new(), @@ -160,6 +163,13 @@ impl Ending { }); } + /// 响应体落盘之前按什么换、打码。**和请求体用同一套**：开始时生效的规则，拦截档下 + /// 这个请求的账本 —— 上游回答里的占位符和存下来的请求对得上号。没交代的按出厂规则 + /// 打码（见 [`Redaction::default`]） + pub fn redact_with(&mut self, r: Redaction) { + self.redaction = Some(r); + } + /// 这一次由谁回答：成功走完时记下它和它读写了多少缓存。 pub fn answered_by(&mut self, ticket: crate::affinity::Ticket) { self.answer = Some(ticket); @@ -171,11 +181,15 @@ impl Ending { } /// 上游来了一块。**这里看的是上游原话**（带占位符的那一版）：usage - /// 数字不受影响，而请求详情里存的正是「发出去的和收回来的」。 + /// 数字不受影响，而请求详情里存的正是「发出去的和收回来的」。原话里要是带着 + /// 认得出的值（观察档下模型回显的密钥），落盘之前打码（见 [`crate::bodies`]）。 pub fn feed(&mut self, chunk: &[u8]) { self.bytes += chunk.len() as u64; self.sniffer.feed(chunk); - self.tap.feed(chunk); + // 没有去处（观测层没起来）就不攒：一个回答最多攒 4 MB，攒了也交不出去 + if self.sink.is_some() { + self.tap.feed(chunk); + } self.spot(chunk); self.watch_for_errors(chunk); } @@ -299,13 +313,14 @@ impl Ending { if !recorded.is_empty() { crate::bodies::offer( &self.sink, - BodyRecord { - id: self.id, - at_ms: self.at_ms, - kind: BodyKind::Response, - body: recorded, + BodyRecord::new( + self.id, + self.at_ms, + BodyKind::Response, + recorded, original_len, - }, + self.redaction.take().unwrap_or_default(), + ), ); } let sniffer = std::mem::take(&mut self.sniffer); @@ -673,7 +688,7 @@ mod tests { fn dropped_mid_stream_it_reports_a_cancellation_with_what_it_saw() { let bus = tw_observe::EventBus::new(); let mut rx = bus.subscribe(); - let (tx, mut bodies) = tokio::sync::mpsc::channel(4); + let (tx, mut bodies) = crate::bodies::channel(); let mut e = Ending::new( bus.clone(), 7, @@ -709,6 +724,20 @@ mod tests { assert_eq!(body.at_ms, 1_000); } + /// 观测层没起来（没有去处）的时候，回答一个字节都不攒 + #[test] + fn with_nowhere_to_send_it_the_answer_is_not_kept() { + let bus = tw_observe::EventBus::new(); + let mut e = Ending::new(bus, 7, MODEL.into(), Instant::now(), 1_000, None); + e.responded(200); + e.feed(MESSAGE_START); + let (kept, seen) = std::mem::take(&mut e.tap).finish(); + assert!(kept.is_empty() && seen == 0, "{seen}"); + // 数还是照数的 + assert_eq!(e.bytes, MESSAGE_START.len() as u64); + e.finished(200); + } + /// 响应头还没到，客户端就走了。**没有状态码，也没有用量** —— 两个都是 /// None，不是 0：0 会让一次真实的调用看起来是免费的，状态码 0 则是编的。 #[test] diff --git a/crates/tw-gateway/src/fixture.rs b/crates/tw-gateway/src/fixture.rs index 6aef83f0..c64f586f 100644 --- a/crates/tw-gateway/src/fixture.rs +++ b/crates/tw-gateway/src/fixture.rs @@ -209,7 +209,11 @@ pub fn record( ) -> Fixture { let clean = |r: Recorded| -> Recorded { // 两道：先按凭据规则换成占位符（结构还在，值没了），再走一遍 - // 通用打码兜住规则没认出来的 + // 通用打码兜住规则没认出来的。 + // + // 存下来的正文落盘前已经按用户的规则换过、打过码（见 `crate::bodies`），这里按全部 + // 内置规则再来一遍：用户关掉的那几条、内网地址，到这里也换掉。新发的号让开存下来 + // 的那份里已经写着的占位符，两样东西不共用一个号 let all: Vec<&str> = tw_guard::redact::rules::BUILTINS .iter() .map(|b| b.id) @@ -217,7 +221,8 @@ pub fn record( let redacted = tw_guard::redact::replace::redact( &r.body, &tw_guard::redact::rules::RuleSet::only(&all), - tw_guard::redact::replace::Ledger::new(tw_guard::redact::replace::Scheme::SECRET), + tw_guard::redact::replace::Ledger::new(tw_guard::redact::replace::Scheme::SECRET) + .avoiding(&r.body), ) .text; Recorded { diff --git a/crates/tw-gateway/src/guard.rs b/crates/tw-gateway/src/guard.rs index e8241708..2b91263c 100644 --- a/crates/tw-gateway/src/guard.rs +++ b/crates/tw-gateway/src/guard.rs @@ -12,13 +12,40 @@ //! 上游的类别替换，于是同一个请求观察时报「检测到」，切到拦截后一处不换 —— //! 用户看到的证据，和他切过去之后得到的保护，说的不是一件事。 //! -//! 两步分开：[`find`] 在尝试上游之前对客户端发来的原文看一遍，报出去的记录 +//! 两步分开：[`look`] 在尝试上游之前对客户端发来的原文看一遍，报出去的记录 //! 只有这一份；[`replace`] 在每一跳发出去之前替换 —— 那一跳的请求体可能是 //! 转换过格式的，要换的是真正发出去的那一份。 +//! +//! # 一个值一个占位符，整个请求里都一样 +//! +//! 拦截档下 [`look`] 按客户端原文里出现的先后给找到的值编好号，每一跳都接着这本账换。 +//! 以前每一跳各起一本账，按那一跳发出去的那份的先后编号：转换过格式、字段换了顺序的 +//! 那一跳，同一把密钥可能是 2 号，而上一跳、存下来的那份请求里它是 1 号 —— 请求详情里 +//! 存的请求和回答对不上号。 use tw_config::SecurityMode as Mode; use tw_guard::redact::replace::{Ledger, Scheme}; -use tw_guard::redact::rules::{Finding, RuleSet}; +use tw_guard::redact::rules::{Finding, Hit, RuleSet}; + +/// 按规则找一遍，**不算我们自己的占位符**。 +/// +/// 连接串里写着 `postgres://app:<>@db` 的那一段，在口令那条规则看来就是 +/// 一个口令 —— 可它是我们换上去的：存下来的请求拿去重放、用户把详情里看到的东西贴回 +/// 对话，都会带着它。当成凭据的话，它会被再换一次、在安全日志里报一次、落盘时被打成 +/// `<>`。压在一个占位符上的命中都不算。 +pub fn hits(text: &str, rules: &RuleSet) -> Vec { + let mut hits = tw_guard::redact::rules::scan(text, rules); + if hits.is_empty() || !text.contains(Scheme::SECRET.open) { + return hits; + } + let ours = Scheme::SECRET.find_in(text); + hits.retain(|h| { + !ours + .iter() + .any(|(at, _, _)| at.start < h.bytes.end && h.bytes.start < at.end) + }); + hits +} /// 找一遍。**观察档和拦截档都找**，关闭时不找。 /// @@ -30,26 +57,69 @@ pub fn find(mode: Mode, rules: &RuleSet, body: &[u8]) -> Vec { let Ok(text) = std::str::from_utf8(body) else { return Vec::new(); }; - let hits = tw_guard::redact::rules::scan(text, rules); - tw_guard::redact::rules::findings(text, &hits) + tw_guard::redact::rules::findings(text, &hits(text, rules)) +} + +/// 一本新账，让开 `body` 里已经写着的占位符（见 [`Ledger::avoiding`]）。 +/// +/// 存下来的请求（拦截档下存的是换过的那一份）拿去重放时，里面写着的 1 号不能再发给 +/// 新找到的值 —— 回显里的 1 号会被还原成那个新值。 +pub fn ledger_for(body: &[u8]) -> Ledger { + let fresh = Ledger::new(Scheme::SECRET); + match std::str::from_utf8(body) { + Ok(text) => fresh.avoiding(text), + Err(_) => fresh, + } +} + +/// 看一遍客户端发来的原文：报出去的记录（同 [`find`]），和这个请求的账本。 +/// +/// **拦截档下账本在这里就编好号**：原文里找到的每个值按出现的先后发号，让开原文里本来 +/// 就写着的占位符。之后每一跳都接着这本账换（[`replace`]），存下来的那份请求也照它换 +/// （[`crate::bodies::Redaction`]）。不在拦截档时账本是空的。 +pub fn look(mode: Mode, rules: &RuleSet, body: &[u8]) -> (Vec, Ledger) { + let empty = || Ledger::new(Scheme::SECRET); + if !mode.detects() || rules.is_empty() { + return (Vec::new(), empty()); + } + let Ok(text) = std::str::from_utf8(body) else { + return (Vec::new(), empty()); + }; + let hits = hits(text, rules); + let found = tw_guard::redact::rules::findings(text, &hits); + if !mode.acts() { + return (found, empty()); + } + let seed = empty().avoiding(text); + let ledger = if hits.is_empty() { + seed + } else { + tw_guard::redact::replace::apply(text, &hits, seed).ledger + }; + (found, ledger) } -/// 拦截档下换掉要发出去的这一份。返回换过的体和还原用的账本；**不在拦截档、 -/// 或者没找到东西时与进来时逐字节相同**，账本是空的。 -pub fn replace(mode: Mode, rules: &RuleSet, body: bytes::Bytes) -> (bytes::Bytes, Ledger) { +/// 拦截档下换掉要发出去的这一份，**接着 `ledger` 的账**（见 [`look`]）。返回换过的体和 +/// 还原用的账本；**不在拦截档、或者没找到东西时与进来时逐字节相同**，账本就是交进来的那本。 +pub fn replace( + mode: Mode, + rules: &RuleSet, + body: bytes::Bytes, + ledger: &Ledger, +) -> (bytes::Bytes, Ledger) { if !mode.acts() || rules.is_empty() { - return (body, Ledger::new(Scheme::SECRET)); + return (body, ledger.clone()); } // 按字节乱切一个非 UTF-8 的体，得到的是一份坏掉的请求 let Ok(text) = std::str::from_utf8(&body) else { - return (body, Ledger::new(Scheme::SECRET)); + return (body, ledger.clone()); }; - let hits = tw_guard::redact::rules::scan(text, rules); + let hits = hits(text, rules); if hits.is_empty() { // 没命中就原样返回，连一次拷贝都不做 - return (body, Ledger::new(Scheme::SECRET)); + return (body, ledger.clone()); } - let r = tw_guard::redact::replace::apply(text, &hits, Ledger::new(Scheme::SECRET)); + let r = tw_guard::redact::replace::apply(text, &hits, ledger.clone()); (bytes::Bytes::from(r.text), r.ledger) } @@ -292,9 +362,13 @@ mod tests { )) } + fn fresh() -> Ledger { + Ledger::new(Scheme::SECRET) + } + #[test] fn enforce_replaces_with_a_placeholder_and_keeps_the_body_valid_json() { - let (out, ledger) = replace(Mode::Enforce, &RuleSet::defaults(), body()); + let (out, ledger) = replace(Mode::Enforce, &RuleSet::defaults(), body(), &fresh()); let text = String::from_utf8(out.to_vec()).unwrap(); assert!(!text.contains(KEY), "{text}"); assert!(text.contains("<>"), "{text}"); @@ -316,7 +390,7 @@ mod tests { assert_eq!(seen.len(), 1); assert!(!seen[0].masked.contains("AAAAAAAAAAAA")); assert_eq!(seen, find(Mode::Enforce, &RuleSet::defaults(), &body())); - let (out, ledger) = replace(Mode::Observe, &RuleSet::defaults(), body()); + let (out, ledger) = replace(Mode::Observe, &RuleSet::defaults(), body(), &fresh()); assert_eq!(out, body()); assert!(ledger.is_empty()); } @@ -324,7 +398,7 @@ mod tests { #[test] fn off_does_not_even_look() { assert!(find(Mode::Off, &RuleSet::defaults(), &body()).is_empty()); - let (out, _) = replace(Mode::Off, &RuleSet::defaults(), body()); + let (out, _) = replace(Mode::Off, &RuleSet::defaults(), body(), &fresh()); assert_eq!(out, body()); } @@ -332,7 +406,7 @@ mod tests { fn a_binary_body_is_left_alone_instead_of_being_mangled() { // 按字节乱切一个非 UTF-8 的体，得到的是一份坏掉的请求。 let raw = bytes::Bytes::from(vec![0xff, 0xfe, 0x00, 0x01]); - let (out, l) = replace(Mode::Enforce, &RuleSet::defaults(), raw.clone()); + let (out, l) = replace(Mode::Enforce, &RuleSet::defaults(), raw.clone(), &fresh()); assert_eq!(out, raw); assert!(l.is_empty()); assert!(find(Mode::Enforce, &RuleSet::defaults(), &raw).is_empty()); @@ -341,7 +415,7 @@ mod tests { #[test] fn a_body_with_nothing_to_redact_is_returned_untouched() { let plain = bytes::Bytes::from_static(b"{\"messages\":[]}"); - let (out, l) = replace(Mode::Enforce, &RuleSet::defaults(), plain.clone()); + let (out, l) = replace(Mode::Enforce, &RuleSet::defaults(), plain.clone(), &fresh()); assert_eq!(out, plain); assert!(l.is_empty()); } @@ -354,4 +428,94 @@ mod tests { assert!(!it[0].custom); assert!(!it[0].masked.contains("AAAAAAAAAAAA"), "{}", it[0].masked); } + + /// 每一跳接着原文那本账换：同一把密钥在每一跳都是同一个号，哪怕那一跳发出去的那份 + /// 把字段换了顺序（转换过格式，或者改写参数时按键名重排过）。以前各起一本账，下面 + /// 这一跳里 `system` 排到了 `messages` 后面，两把密钥的号就对调了 + #[test] + fn every_hop_numbers_a_value_the_way_the_client_body_did() { + let other = "ghp_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; + let client = + format!(r#"{{"system":"{KEY}","messages":[{{"role":"user","content":"{other}"}}]}}"#); + let (found, l0) = look(Mode::Enforce, &RuleSet::defaults(), client.as_bytes()); + assert_eq!((found.len(), l0.len()), (2, 2)); + let hop = + format!(r#"{{"messages":[{{"role":"user","content":"{other}"}}],"system":"{KEY}"}}"#); + let (out, ledger) = replace(Mode::Enforce, &RuleSet::defaults(), hop.clone().into(), &l0); + assert_eq!( + String::from_utf8(out.to_vec()).unwrap(), + r#"{"messages":[{"role":"user","content":"<>"}],"system":"<>"}"# + ); + assert_eq!(ledger.len(), 2); + // 各起一本账的话号就对调了 —— 这条测试防的就是它 + let (alone, _) = replace(Mode::Enforce, &RuleSet::defaults(), hop.into(), &fresh()); + assert!( + String::from_utf8(alone.to_vec()) + .unwrap() + .contains(r#""system":"<>""#) + ); + } + + #[test] + fn look_numbers_only_under_enforce_and_reports_the_same_either_way() { + let (seen, l) = look(Mode::Observe, &RuleSet::defaults(), &body()); + assert_eq!(seen, find(Mode::Observe, &RuleSet::defaults(), &body())); + assert!(l.is_empty(), "观察档不该编号"); + let (acted, l) = look(Mode::Enforce, &RuleSet::defaults(), &body()); + assert_eq!(acted, seen); + assert_eq!(l.len(), 1); + let (none, l) = look(Mode::Off, &RuleSet::defaults(), &body()); + assert!(none.is_empty() && l.is_empty()); + } + + /// 连接串里的占位符长得像口令，可它不是凭据：不再换一次、不报、原样留着 + #[test] + fn a_placeholder_where_a_password_would_be_is_not_a_password() { + let t = + format!("postgres://app:<>@db/x 和 postgres://app:hunter2@db/y 和 {KEY}"); + let found: Vec = hits(&t, &RuleSet::defaults()) + .iter() + .map(|h| t[h.bytes.clone()].to_string()) + .collect(); + assert_eq!(found, vec!["hunter2".to_string(), KEY.to_string()]); + let body = format!(r#"{{"content":"{t}"}}"#); + assert_eq!( + find(Mode::Observe, &RuleSet::defaults(), body.as_bytes()).len(), + 2 + ); + let (out, _) = replace( + Mode::Enforce, + &RuleSet::defaults(), + body.clone().into(), + &ledger_for(body.as_bytes()), + ); + let out = String::from_utf8(out.to_vec()).unwrap(); + assert!(out.contains("postgres://app:<>@db/x"), "{out}"); + assert!(out.contains("postgres://app:<>@db/y"), "{out}"); + } + + #[test] + fn a_placeholder_already_in_the_client_body_is_not_handed_out_again() { + // 用户把请求详情里看到的请求贴回了对话：里面写着 1 号 + let pasted = format!( + r#"{{"messages":[{{"role":"user","content":"上次发的是 <>，这次是 {KEY}"}}]}}"# + ); + let (_, l0) = look(Mode::Enforce, &RuleSet::defaults(), pasted.as_bytes()); + let (out, ledger) = replace(Mode::Enforce, &RuleSet::defaults(), pasted.into(), &l0); + let out = String::from_utf8(out.to_vec()).unwrap(); + assert!(out.contains("这次是 <>"), "{out}"); + assert_eq!( + tw_guard::redact::replace::restore("<> / <>", &ledger), + format!("<> / {KEY}") + ); + // 重放用的那本新账也让开它 + let stored = format!("<> {KEY}"); + let (out, _) = replace( + Mode::Enforce, + &RuleSet::defaults(), + bytes::Bytes::from(stored.clone()), + &ledger_for(stored.as_bytes()), + ); + assert_eq!(&out[..], b"<> <>"); + } } diff --git a/crates/tw-gateway/src/server/pipeline.rs b/crates/tw-gateway/src/server/pipeline.rs index 67b2837a..2d3e7d83 100644 --- a/crates/tw-gateway/src/server/pipeline.rs +++ b/crates/tw-gateway/src/server/pipeline.rs @@ -52,6 +52,9 @@ struct Started { choice: Choice, /// 这是哪段对话（见 [`crate::affinity::identity`]）。认不出来是 None conversation: Option, + /// 出站脱敏的账本：拦截档下按客户端原文编好了号，每一跳接着它换（见 + /// [`crate::guard::look`]）。别的档位是空的 + ledger: tw_guard::redact::replace::Ledger, } pub(super) async fn pipeline( @@ -87,7 +90,19 @@ pub(super) async fn pipeline( // `passthrough` 按这个错误报 —— 不排队：它一个字节都不会发出去 Routed::Refused(choice, why) => { let to = ("", tw_api::Billing::PerToken); - let id = open(&state, &req, &reading, &choice, to, fp.as_deref(), ending); + // 一个字节都没发出去，也没什么可报的；存下来的请求照样按这一档换、打码 + let (_, ledger) = look(&rt, &req); + let redaction = redaction(&rt, ledger); + let id = open( + &state, + &req, + &reading, + &choice, + to, + fp.as_deref(), + ending, + redaction, + ); state.bus.emit(super::routed_nowhere(id, choice)); return Err(why); } @@ -523,7 +538,7 @@ fn route( Ok(Routed::Go(choice, decision)) } -/// 发出开始事件：熔断过滤、`RequestStarted`、结局、入站脱敏的记录、请求体留档。 +/// 发出开始事件：熔断过滤、出站脱敏看一遍、`RequestStarted`、结局、脱敏的记录、请求体留档。 /// /// **发了开始，就欠一个结局。**从这里起，管线返回的错误由调用方报成 /// 失败，这个 future 被丢掉由 Drop 报成取消（见 `passthrough`）。 @@ -559,6 +574,12 @@ fn start( .find(|p| p.name == first) .map(|p| p.billing) .unwrap_or_default(); + + // 出站脱敏：按全局的规则看一遍客户端发来的原文。**观察档和拦截档报的 + // 是同一条记录**，差别只在换没换 —— 真正的替换在每一跳发出去之前做， + // 那一跳的请求体可能是转换过格式的。拦截档下账本在这里就编好号：每一跳、 + // 存下来的那份请求都按它换，同一个值处处是同一个占位符 + let (found, ledger) = look(rt, req); let id = open( state, req, @@ -567,13 +588,9 @@ fn start( (first, billing.into()), fp, ending, + redaction(rt, ledger.clone()), ); - - // 出站脱敏：按全局的规则看一遍客户端发来的原文。**观察档和拦截档报的 - // 是同一条记录**，差别只在换没换 —— 真正的替换在每一跳发出去之前做， - // 那一跳的请求体可能是转换过格式的。 let redact_mode = rt.config.security.redact.mode; - let found = crate::guard::find(redact_mode, &rt.redact, &req.body); if !found.is_empty() { state.bus.emit(tw_api::Event::SecretsFound { id, @@ -588,15 +605,36 @@ fn start( alive, choice, conversation: crate::affinity::identity(&req.headers, fp), + ledger, + } +} + +/// 出站脱敏看一遍客户端发来的原文（见 [`crate::guard::look`]）。 +fn look( + rt: &Runtime, + req: &Inbound, +) -> ( + Vec, + tw_guard::redact::replace::Ledger, +) { + crate::guard::look(rt.config.security.redact.mode, &rt.redact, &req.body) +} + +/// 这个请求的正文落盘之前怎么换、怎么打码：此刻生效的规则，和这个请求的账本。 +fn redaction(rt: &Runtime, ledger: tw_guard::redact::replace::Ledger) -> crate::bodies::Redaction { + crate::bodies::Redaction { + rules: rt.redact.clone(), + ledger, } } /// 发 `RequestStarted`、把这个请求欠着的结局放进 `ending`、把请求体交去留档， /// 交回这个请求的号。`to` 是要发往的那一家和它怎么收钱；一家都不会去的（被规则 -/// 拒绝了）是空的名字。 +/// 拒绝了）是空的名字。`redaction` 是请求体、响应体落盘之前怎么换、打码。 /// /// **会话在这里定**（见 [`crate::session::Sessions`]）：开始事件带着它，落库的 /// 那一行记的也是它。 +#[allow(clippy::too_many_arguments)] fn open( state: &AppState, req: &Inbound, @@ -605,6 +643,7 @@ fn open( to: (&str, tw_api::Billing), fp: Option<&str>, ending: &mut Option, + redaction: crate::bodies::Redaction, ) -> u64 { let facts = &reading.facts; let id = state.bus.next_id(); @@ -636,27 +675,32 @@ fn open( at_ms, }); let sink = state.body_sink(); - *ending = Some(crate::ending::Ending::new( + let mut end = crate::ending::Ending::new( state.bus.clone(), id, facts.model.clone(), req.started, at_ms as i64, sink.clone(), - )); + ); + end.redact_with(redaction.clone()); + *ending = Some(end); // 请求体交给观测层。**这时候它已经完整在内存里了**，所以这一步 // 除了一次 `Bytes` 的引用计数之外没有别的成本（说过入站是要 - // 整个解析的，所以本来就在）。 + // 整个解析的，所以本来就在）；比存得下的还长的，只拷开头那一段（见 + // `bodies::offer`）。**交出去的是原文**：换掉、打码在落盘那一头做，不占 + // 转发这条路（见 `crate::bodies`） crate::bodies::offer( &sink, - crate::bodies::BodyRecord { + crate::bodies::BodyRecord::new( id, - at_ms: at_ms as i64, - kind: crate::bodies::BodyKind::Request, - body: req.body.clone(), - original_len: req.body.len(), - }, + at_ms as i64, + crate::bodies::BodyKind::Request, + req.body.clone(), + req.body.len(), + redaction, + ), ); id } diff --git a/crates/tw-gateway/src/server/pipeline/hop.rs b/crates/tw-gateway/src/server/pipeline/hop.rs index 2af55cb4..7d8456ba 100644 --- a/crates/tw-gateway/src/server/pipeline/hop.rs +++ b/crates/tw-gateway/src/server/pipeline/hop.rs @@ -20,8 +20,8 @@ use tw_types::msg; pub(super) struct Served<'a> { pub(super) upstream: reqwest::Response, pub(super) provider: &'a tw_config::Provider, - /// 成功那一次的脱敏账本。**必须是成功那一次的** —— 每一跳发出去的体可能 - /// 转换过格式，占位符按那一份的顺序编号 + /// 成功那一次的脱敏账本。**必须是成功那一次的** —— 每一跳都接着原文那本账换， + /// 而那一跳发出去的体（可能转换过格式）里还有原文没有的值时，号是那一跳新发的 pub(super) ledger: tw_guard::redact::replace::Ledger, /// 成功那一跳的转换。**必须是成功那一次的** —— 故障转移从 Anthropic 上游 /// 切到 OpenAI 上游时，两跳转成的格式不一样；直通时是 None @@ -200,9 +200,14 @@ pub(super) async fn try_upstreams<'a>( // 这一家在这段对话里拒过的别家封存的推理：发之前先去掉（见 `crate::seal`） let unsealed = unseal_upfront(state, req, started, provider, &out); // 出站脱敏的拦截档：换掉**这一跳真正发出去的那一份**（可能转换过 - // 格式）。规则是全局的，每一跳换掉的是同一批东西 - let (body, ledger) = - crate::guard::replace(rt.config.security.redact.mode, &rt.redact, unsealed); + // 格式）。规则是全局的，每一跳换掉的是同一批东西；**接着原文那本账换**， + // 同一个值在每一跳、在存下来的那份请求里都是同一个占位符 + let (body, ledger) = crate::guard::replace( + rt.config.security.redact.mode, + &rt.redact, + unsealed, + &started.ledger, + ); // 用这个 provider 自己的 Client —— 它带着该走的代理。**在取密钥 // 之前拿到**：OAuth 换 token 也要走这条代理。 diff --git a/crates/tw-gateway/src/state.rs b/crates/tw-gateway/src/state.rs index 230ec89f..58946d26 100644 --- a/crates/tw-gateway/src/state.rs +++ b/crates/tw-gateway/src/state.rs @@ -149,7 +149,7 @@ pub struct AppState { /// **有界通道，满了就丢。**直接调用意味着文件 I/O 跑在转发那条路上 /// —— 一次慢磁盘写就变成一次慢请求，而观测永远不该有这个权力。 /// `None` 表示观测层没起来，那时什么都不做。 - body_sink: Arc>>, + body_sink: Arc>>, /// 每个上游最近一次报的订阅额度。 /// /// **在内存里，不落库。**它是「现在还剩多少」，不是历史 —— 存一份 @@ -312,9 +312,9 @@ impl AppState { &self.relisten } - /// 接上 body 的去处。**观测层起来之后才调** —— 在那之前 body 一律 - /// 丢掉，而请求照常。 - pub fn set_body_sink(&self, tx: crate::bodies::BodySender) { + /// 接上 body 的去处（[`crate::bodies::channel`] 的发的那一头）。**观测层起来之后才调** + /// —— 在那之前 body 一律丢掉，而请求照常。 + pub fn set_body_sink(&self, tx: crate::bodies::BodySink) { if let Ok(mut g) = self.body_sink.lock() { *g = Some(tx); } @@ -368,7 +368,7 @@ impl AppState { } } - pub(crate) fn body_sink(&self) -> Option { + pub(crate) fn body_sink(&self) -> Option { self.body_sink.lock().ok().and_then(|g| g.clone()) } diff --git a/crates/tw-gateway/tests/claude_desktop.rs b/crates/tw-gateway/tests/claude_desktop.rs index 94c3db2f..5050e645 100644 --- a/crates/tw-gateway/tests/claude_desktop.rs +++ b/crates/tw-gateway/tests/claude_desktop.rs @@ -117,7 +117,7 @@ async fn gateway_pinging_for(p: Provider, ping_for: Duration) -> Gateway { state.ping_every = Duration::from_millis(100); state.ping_for = ping_for; let events = state.bus.subscribe(); - let (tx, bodies) = tokio::sync::mpsc::channel(16); + let (tx, bodies) = tw_gateway::bodies::channel(); state.set_body_sink(tx); let addr = tw_gateway::serve(state, ([127, 0, 0, 1], 0).into()) .await diff --git a/crates/tw-guard/src/redact/replace.rs b/crates/tw-guard/src/redact/replace.rs index 8e962b2d..49c245c8 100644 --- a/crates/tw-guard/src/redact/replace.rs +++ b/crates/tw-guard/src/redact/replace.rs @@ -4,6 +4,7 @@ //! 而「瞎猜」在一个正帮你调试 `.env` 的助手身上，比看不见更糟。 use std::collections::HashMap; +use std::ops::Range; use crate::redact::rules::{Hit, RuleSet}; @@ -33,6 +34,38 @@ impl Scheme { pub fn placeholder(&self, label: &str, n: usize) -> String { format!("{}{label}_{n}{}", self.open, self.close) } + + /// `text` 里写着的每一个这种写法的占位符：在哪儿（字节区间）、标签、号码。 + /// + /// **夹在别的字里的也算**（`<<>`）：还原是按子串换的。号码大到装不下的 + /// 不算 —— 那不可能是发出去的。 + pub fn find_in<'a>(&self, text: &'a str) -> Vec<(Range, &'a str, usize)> { + let mut out = Vec::new(); + let mut from = 0; + while let Some(i) = text[from..].find(self.open) { + let at = from + i; + let start = at + self.open.len(); + // 下一处从这一处的下一个字节找起：`<<>` 里的占位符从第二个 `<` 起。 + // 开头那段是 ASCII，加一还落在字的边界上 + from = at + 1; + let rest = &text[start..]; + let len = rest + .find(|c: char| !(c.is_ascii_alphanumeric() || c == '_')) + .unwrap_or(rest.len()); + if !rest[len..].starts_with(self.close) { + continue; + } + let Some((label, n)) = rest[..len].rsplit_once('_') else { + continue; + }; + if label.is_empty() || n.is_empty() || !n.bytes().all(|b| b.is_ascii_digit()) { + continue; + } + let Ok(n) = n.parse::() else { continue }; + out.push((at..start + len + self.close.len(), label, n)); + } + out + } } /// 一次脱敏留下的账本。 @@ -85,6 +118,22 @@ impl Ledger { self.seen.iter().map(|(o, p)| (o.as_str(), p.as_str())) } + /// 让开 `text` 里已经写着的占位符：新发的号接在它们后面数。 + /// + /// 一段文字里本来就可能写着占位符：存下来的请求（拦截档下存的是换过的那一份）拿去 + /// 重放、导出成用例，用户把请求详情里看到的东西贴回对话。照常从 1 数的话，新找到的 + /// 值会拿到一个已经写在那儿的号 —— 两样东西共用一个占位符，回显里的那一处被还原成 + /// 新值，而它原本指的是另一样。 + /// + /// **只让号，不记账**：写在那儿的占位符不知道原来是什么，还原时原样留着。 + pub fn avoiding(mut self, text: &str) -> Self { + for (_, label, n) in self.scheme.find_in(text) { + let taken = self.issued.entry(label.to_string()).or_insert(0); + *taken = (*taken).max(n); + } + self + } + /// 这个值的占位符，头一次见就发一个新号。 fn issue(&mut self, original: &str, label: &str) -> String { if let Some(p) = self.seen.get(original) { @@ -373,6 +422,62 @@ mod tests { assert_eq!(pairs["c@d.com"], "{{EMAIL_2}}"); } + #[test] + fn placeholders_already_in_the_text_keep_their_numbers_to_themselves() { + // 存下来的请求里写着 1 号（拦截档下换过），重放时又找到一把新的：它拿 1 号的话， + // 回显里的 1 号会被还原成这把新的，而那一处原本指的是另一样 + let t = format!("旧的 <>，身份证 <>，新的 {KEY}"); + let r = redact(&t, &all(), l().avoiding(&t)); + assert!(r.text.contains("新的 <>"), "{}", r.text); + assert!(r.text.contains("旧的 <>"), "{}", r.text); + // 让开的号不记账：回显里的 1 号原样留着，不被还原成任何东西 + assert_eq!(r.ledger.len(), 1); + assert_eq!( + restore("<> 和 <>", &r.ledger), + format!("<> 和 {KEY}") + ); + } + + #[test] + fn placeholders_are_found_where_they_are_written() { + let t = "前 <> 中 <<>> 后 <> < = Scheme::SECRET + .find_in(t) + .into_iter() + .map(|(at, label, n)| (&t[at], label, n)) + .collect(); + assert_eq!( + found, + vec![ + ("<>", "TW_SECRET", 1), + ("<>", "TW_ID_NUMBER", 12), + ] + ); + } + + #[test] + fn only_whole_placeholders_of_this_scheme_are_set_aside() { + let taken = |t: &str| { + let r = redact(KEY, &all(), l().avoiding(t)); + r.text + }; + // 夹在别的字里的照样算：还原是按子串换的 + assert_eq!(taken("<<>>"), "<>"); + // 不是占位符的不让：没收尾的、号不是数字的、别的写法的、标签是空的 + for t in [ + "<>", + "{{TW_SECRET_4}}", + "<<_4>>", + "<>", + "<>", + ] { + assert_eq!(taken(t), "<>", "{t}"); + } + // 别的标签只让它自己的号 + assert_eq!(taken("<>"), "<>"); + } + #[test] fn decoded_text_is_matched_as_written_and_restored_into_json_escaped() { // 正文上的 `password="x"`：截在引号处的话，换下来的只是 `password=`， diff --git a/crates/tw-secret/src/mask.rs b/crates/tw-secret/src/mask.rs index 606b5597..6694f1ae 100644 --- a/crates/tw-secret/src/mask.rs +++ b/crates/tw-secret/src/mask.rs @@ -543,32 +543,48 @@ mod tests { /// **按值的形状判，不按键名。**body 是 JSON，键名五花八门（`api_key`、 /// `token`、`Authorization`、某个 MCP server 自己起的名字），而凭据的 /// 形状是有限的几种。 +/// +/// **打第二遍不改动什么**：正文落盘之前打过一遍，读出来给人看时再打一遍。打过的码 +/// （`sk-an…7f9c`）里那 5 个字的开头自己又像一把密钥，不认得它的话第二遍会把它打成 +/// `………`，越打越看不出是什么。所以紧挨着 `…` 的 5 个字的开头、4 个字的结尾原样留着 —— +/// 这么短的一段本来就不带信息，[`mask_secret`] 对它也只会给 `…`。 +/// +/// 按字节走，不先摊成一个字符数组：一份正文最多 4 MB，摊开就是 16 MB。 pub fn mask_body(text: &str) -> String { - // 一个可能是凭据的 token 由这些字符组成 - fn is_tok(c: char) -> bool { - c.is_ascii_alphanumeric() || "-_.".contains(c) + // 一个可能是凭据的 token 由这些字符组成。**全是 ASCII**，所以按字节找边界 + fn is_tok(b: u8) -> bool { + b.is_ascii_alphanumeric() || matches!(b, b'-' | b'_' | b'.') } + const ELLIPSIS: &str = "…"; + let bytes = text.as_bytes(); let mut out = String::with_capacity(text.len()); - let chars: Vec = text.chars().collect(); let mut i = 0; - while i < chars.len() { - // 只在一个 token 的开头尝试匹配，否则 `xsk-abc` 里的 `sk-abc` - // 会被当成密钥 - let at_boundary = i == 0 || !is_tok(chars[i - 1]); - if at_boundary { - let mut j = i; - while j < chars.len() && is_tok(chars[j]) { - j += 1; - } - let tok: String = chars[i..j].iter().collect(); - if looks_like_credential(&tok) { - out.push_str(&mask_secret(&tok)); - i = j; - continue; - } + while i < bytes.len() { + if !is_tok(bytes[i]) { + // 一段不是 token 的字原样照抄，到下一个 token 的开头为止 + let j = bytes[i..] + .iter() + .position(|&b| is_tok(b)) + .map_or(bytes.len(), |n| i + n); + out.push_str(&text[i..j]); + i = j; + continue; } - out.push(chars[i]); - i += 1; + // 只在一个 token 的开头尝试匹配（走到这里的都是：前一个字不是 token 的），否则 + // `xsk-abc` 里的 `sk-abc` 会被当成密钥 + let j = bytes[i..] + .iter() + .position(|&b| !is_tok(b)) + .map_or(bytes.len(), |n| i + n); + let tok = &text[i..j]; + let head_of_a_mask = tok.len() == 5 && text[j..].starts_with(ELLIPSIS); + let tail_of_a_mask = tok.len() == 4 && text[..i].ends_with(ELLIPSIS); + if !head_of_a_mask && !tail_of_a_mask && looks_like_credential(tok) { + out.push_str(&mask_secret(tok)); + } else { + out.push_str(tok); + } + i = j; } out } @@ -615,6 +631,38 @@ mod body_tests { assert!(!out.contains("bbbbbbbbbbbbbbbb"), "{out}"); } + /// 正文落盘前打过一遍，读出来再打一遍。**第二遍不能再改**：`sk-an…7f9c` 的开头 + /// `sk-an` 自己又像一把密钥，不认得它的话，第二遍会把它打成 `………7f9c` + #[test] + fn masking_twice_changes_nothing_the_first_time_did_not() { + for s in [ + r#"{"content":"我的 key 是 sk-ant-api03-abcdefghijklmnopqrstuvwxyz，别外传"}"#, + "sk-ant-api03-aaaaaaaaaaaaaaaa 和 ghp_bbbbbbbbbbbbbbbbbbbb…", + "AKIA… tw-abc sk-x…AKIA 网关 tw-0123456789abcdef0123456789", + "a1b2c3d4e5f6a7b8c9d0e1f2a3b4c5d6e7f8 和 \\nsk-ant-api03-zzzzzzzzzzzzzzzzzzzz", + "已经打过的 sk-an…wxyz、ghp_b…bbbb、………、…7894", + "", + ] { + let once = mask_body(s); + assert_eq!(mask_body(&once), once, "{s}"); + } + // 已经打过的码原样留着，别的照打 + assert_eq!( + mask_body("sk-an…wxyz 和 sk-ant-api03-abcdefghijklmnop"), + "sk-an…wxyz 和 sk-an…mnop" + ); + } + + /// 改成按字节走之后，打出来的和原来按字符走的一样（多字节的字夹在中间也一样） + #[test] + fn multibyte_text_around_a_key_is_copied_through_untouched() { + assert_eq!( + mask_body("前面的中文sk-ant-api03-abcdefghijklmnop后面也是中文"), + "前面的中文sk-an…mnop后面也是中文" + ); + assert_eq!(mask_body("密钥：🔑sk-ant-x"), "密钥：🔑………"); + } + #[test] fn the_config_secrets_that_have_no_recognisable_shape_are_still_masked() { // **这三种实测都能原样穿过 `mask_body`。**它们是这个函数存在的理由 diff --git a/crates/tw-store/Cargo.toml b/crates/tw-store/Cargo.toml index 4e7784bf..ec40258c 100644 --- a/crates/tw-store/Cargo.toml +++ b/crates/tw-store/Cargo.toml @@ -18,6 +18,8 @@ tw-pricing = { workspace = true } tw-dialect = { workspace = true } # 按正文找时交出去的那一小段要和详情抽屉同一套打码 tw-secret = { workspace = true } +# 对话记录比对前后两个请求的历史：每条消息只留一个指纹，不留正文 +blake3 = { workspace = true } bytes = { workspace = true } # functions：搜索要一个 SQL 函数，按界面同一种小写规则比子串（见 `search::sql`）。 # 只是打开 rusqlite 已有的接口，不多编任何东西 diff --git a/crates/tw-store/src/blobs.rs b/crates/tw-store/src/blobs.rs index fe9f2bda..95f472f1 100644 --- a/crates/tw-store/src/blobs.rs +++ b/crates/tw-store/src/blobs.rs @@ -6,16 +6,13 @@ use std::path::{Path, PathBuf}; -/// body 默认留几天。 -pub const KEEP_DAYS: u64 = 7; -/// 总量上限。超了从最旧的天目录开始删。 -pub const MAX_BYTES: u64 = 2 * 1024 * 1024 * 1024; - -/// 单个 body 的上限。 +/// 单个 body 的上限（[`tw_api::BODY_MAX`]，4 MiB）。 /// /// 超过就只留开头。**一个 200 MB 的请求体存下来对排查没有额外帮助** —— /// 而它会把当天的目录一次撑爆，把别的请求的 body 挤掉。 -pub const MAX_ONE: usize = 4 * 1024 * 1024; +/// +/// 留几天、总共留多少不在这里定：那是配置里的 `retention`，回收按它来（见 [`Blobs::gc`]）。 +pub const MAX_ONE: usize = tw_api::BODY_MAX; pub struct Blobs { root: PathBuf, @@ -83,6 +80,9 @@ impl Blobs { /// 而它们比 config.yaml 多得多。默认 umask 通常给 0644，那意味着 /// 同一台机器上的别的用户能把它们全读走（那条「权限就是认证」 /// 的同一个道理）。 + /// + /// 交到这里的已经是换过、打过码的那一份（网关的 `bodies::BodyRecord::for_disk`）： + /// 脱敏规则认得出的值进不了磁盘。**规则认不全**，所以权限照样收紧。 pub fn put(&self, at_ms: i64, id: i64, which: Which, body: &[u8]) -> bool { let p = self.path_for(at_ms, id, which); let Some(dir) = p.parent() else { return false }; @@ -125,8 +125,12 @@ impl Blobs { len > MAX_ONE } - /// 存了多少、原本多长。**两个数一起返回** —— 详情页要靠它说出 - /// 「只存了开头 256 KB」。 + /// 写一个 body，连同它原本多长。详情页要靠它说出「只存了开头 4 MB」，重放靠它 + /// 拒绝一份截断过的请求。 + /// + /// `original_len` 比**真正存下的**长时才另记一个 `.len` —— 截断可能发生在交来之前 + /// （网关只攒了开头），也可能发生在这里（`body` 比 [`MAX_ONE`] 长）。以前只看前一种： + /// 一个 5 MB 的请求体存下 4 MB，却没有一处说它被截过，重放照样把半截 JSON 发了出去。 pub fn put_with_len( &self, at_ms: i64, @@ -138,7 +142,7 @@ impl Blobs { if !self.put(at_ms, id, which, body) { return false; } - if original_len > body.len() { + if original_len > body.len().min(MAX_ONE) { let p = self .path_for(at_ms, id, which) .with_extension(format!("{}.len", which.suffix())); @@ -360,7 +364,7 @@ mod tests { // 用户自己放的目录不算一天 std::fs::create_dir_all(b.root().join("0000-backup")).unwrap(); assert_eq!(b.oldest_ms(), Some(10 * DAY)); - b.gc(13 * DAY, 1, MAX_BYTES); + b.gc(13 * DAY, 1, u64::MAX); assert_eq!(b.oldest_ms(), Some(12 * DAY)); } @@ -392,6 +396,25 @@ mod tests { assert!(!Blobs::was_truncated(10)); } + /// 截在这里的也要留下原本多长。以前只有「交来之前就截过」的才记：一个比上限长的 + /// 请求体整份交进来、在这里被截，读回去的人看不出它少了一截 —— 重放照样把半截发出去 + #[test] + fn a_body_cut_here_records_how_long_it_was() { + let (_d, b) = setup(); + let huge = vec![b'x'; MAX_ONE + 1000]; + assert!(b.put_with_len(0, 1, Which::Request, &huge, huge.len())); + assert_eq!(b.get(0, 1, Which::Request).unwrap().len(), MAX_ONE); + assert_eq!(b.original_len(0, 1, Which::Request), Some(MAX_ONE + 1000)); + + // 截在交来之前的：交来的是开头，原本的长度另给 + assert!(b.put_with_len(0, 2, Which::Response, b"head", 9_999)); + assert_eq!(b.original_len(0, 2, Which::Response), Some(9_999)); + + // 整份都存下了的不记：读的人拿存下的长度当原本的 + assert!(b.put_with_len(0, 3, Which::Request, b"whole", 5)); + assert_eq!(b.original_len(0, 3, Which::Request), None); + } + #[test] fn gc_deletes_whole_days_older_than_the_cutoff() { let (_d, b) = setup(); @@ -400,7 +423,7 @@ mod tests { b.put(now - i * DAY, i, Which::Request, &vec![b'x'; 1000]); } assert_eq!(b.days().len(), 10); - let freed = b.gc(now, 3, MAX_BYTES); + let freed = b.gc(now, 3, u64::MAX); // 留 now、now-1、now-2、now-3 这四天（cutoff 是 now-3 那天） assert_eq!(b.days().len(), 4, "{:?}", b.days()); assert_eq!(freed, 6000); @@ -458,7 +481,7 @@ mod tests { #[test] fn gc_on_an_empty_or_missing_root_does_nothing_and_says_zero() { let (_d, b) = setup(); - assert_eq!(b.gc(0, 7, MAX_BYTES), 0); + assert_eq!(b.gc(0, 7, u64::MAX), 0); assert_eq!(b.total_bytes(), 0); } } diff --git a/crates/tw-store/src/db.rs b/crates/tw-store/src/db.rs index 5c20a108..9f247eb3 100644 --- a/crates/tw-store/src/db.rs +++ b/crates/tw-store/src/db.rs @@ -574,6 +574,16 @@ impl Db { })?; Ok(rows.collect::, _>>()?) } + + /// 一次会话里的每一个请求，整行。**和 [`Db::turns`] 同样的筛法、同样的顺序**：对话记录 + /// （[`crate::transcript`]）一轮对一轮地跟着会话详情走。 + pub fn session_requests(&self, session: &str) -> Result, DbError> { + let mut st = self.conn.prepare( + "SELECT * FROM requests WHERE session = ?1 AND local = 0 ORDER BY at_ms, id", + )?; + let rows = st.query_map([session], row_from)?; + Ok(rows.collect::, _>>()?) + } } impl Db { diff --git a/crates/tw-store/src/lib.rs b/crates/tw-store/src/lib.rs index 59c911b2..dd1dcff7 100644 --- a/crates/tw-store/src/lib.rs +++ b/crates/tw-store/src/lib.rs @@ -15,6 +15,7 @@ pub mod model_name; pub mod recorder; pub mod search; pub mod task; +pub mod transcript; pub use blobs::{Blobs, Which}; pub use db::{Db, DbError, Latency, RequestRow, SecurityEvent, Summary, TokenRate}; diff --git a/crates/tw-store/src/search/text.rs b/crates/tw-store/src/search/text.rs index 208af6d0..77579146 100644 --- a/crates/tw-store/src/search/text.rs +++ b/crates/tw-store/src/search/text.rs @@ -175,8 +175,8 @@ pub fn answer(body: &[u8], upstream: Dialect) -> Option { text_of(events) } _ => { - // **一小块一小块地喂。**拆帧器从缓冲区头上一帧一帧地取，整份 256 KB 一次喂进去 - // 的话，每取一帧都要把后面的字节往前挪一遍 + // **一小块一小块地喂。**拆帧器从缓冲区头上一帧一帧地取，整份（最多 4 MB）一次喂 + // 进去的话，每取一帧都要把后面的字节往前挪一遍 let mut r = tw_dialect::convert::Reader::new(upstream); let mut events = Vec::new(); for piece in body.chunks(16 * 1024) { diff --git a/crates/tw-store/src/task.rs b/crates/tw-store/src/task.rs index e4e33ade..d890bd5a 100644 --- a/crates/tw-store/src/task.rs +++ b/crates/tw-store/src/task.rs @@ -21,22 +21,19 @@ pub struct StoredBody { pub id: u64, pub at_ms: i64, pub which: Which, + /// 落盘的那一份：**已经换过、打过码**（网关的 `bodies::BodyRecord::for_disk`）。 + /// 比 [`crate::blobs::MAX_ONE`] 长的由这里截 pub body: bytes::Bytes, - /// 原始长度。截断了要能说出来 + /// 原始长度。截断了要能说出来（见 [`crate::Blobs::put_with_len`]） pub original_len: usize, } -/// 多久回收一次。 -/// -/// 一小时。**不是启动时跑一次就完** —— 一个开着不关的桌面应用会连续跑 -/// 好几天，而 body 的保留策略是按天算的。 -const GC_EVERY: std::time::Duration = std::time::Duration::from_secs(3600); - -/// metadata 留多少天。 -pub const METADATA_KEEP_DAYS: u64 = 90; - /// 起来。返回的 handle 给别的地方查历史用 —— **同一个 Recorder**， /// 不是第二个连接：两个连接会让「刚写进去的还查不到」变成可能。 +/// +/// **回收不在这里。**它按配置里的 `retention` 走（`twcore` 每小时读一次配置，调 +/// [`Recorder::gc`]）。这里以前另跑着一个按写死的 7 天、90 天、2 GB 回收的循环，和按 +/// 配置的那个同时跑：配得比出厂更宽的期限和上限，每小时被它削回出厂值一次。 pub fn spawn( recorder: Recorder, mut rx: tokio::sync::broadcast::Receiver, @@ -67,30 +64,5 @@ pub fn spawn( } } }); - let r = shared.clone(); - tokio::spawn(async move { - let mut tick = tokio::time::interval(GC_EVERY); - // 第一次立刻跑：上次退出之后攒下的过期数据该清了 - loop { - tick.tick().await; - let now = now_ms(); - let freed = r.lock().await.gc( - now, - crate::blobs::KEEP_DAYS, - METADATA_KEEP_DAYS, - crate::blobs::MAX_BYTES, - ); - if freed > 0 { - tracing::info!(mb = freed / 1024 / 1024, "reclaimed expired request bodies"); - } - } - }); shared } - -fn now_ms() -> i64 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as i64) - .unwrap_or(0) -} diff --git a/crates/tw-store/src/transcript/answer.rs b/crates/tw-store/src/transcript/answer.rs new file mode 100644 index 00000000..81c9d98d --- /dev/null +++ b/crates/tw-store/src/transcript/answer.rs @@ -0,0 +1,432 @@ +//! 回答：存下来的上游原话读成一块一块。 +//! +//! `body` 是回答它的那一家的原话：整包的 JSON、SSE 的流（Bedrock 的二进制帧在网关进门时 +//! 已经转成了 SSE），或者 Gemini 不带 `alt=sse` 时那个逐步写出的 JSON 数组。按开头的字节 +//! 分辨，不看请求要没要流：上游只给流、客户端要整包的时候，存下来的是流（和按正文找的 +//! `search::text::answer` 同一个分法）。 +//! +//! - **整包**和请求里的助手消息是同一个形状，按 [`super::read`] 读； +//! - **流**按 tw-dialect 各家的流解析器读成事件再拼回来 —— 网关转换格式时读的就是它们。 +//! 解析器只认得各家之间对应得上的块，服务端工具的块、Gemini 的代码执行这些它跳过， +//! 这里在同一帧上另外认出来，记成 `other`，位置不乱。 + +use std::collections::{HashMap, HashSet}; + +use serde_json::Value; +use tw_api::TranscriptPart; +use tw_dialect::frame::{self, Frame}; +use tw_dialect::ir::{Block, BlockKind, Delta, Dialect, Event, ToolInput}; + +use super::read; + +/// 读出来的回答。 +#[derive(Debug, Default)] +pub(super) struct Answer { + pub(super) parts: Vec, + /// 认出了这家格式的回答没有。存着东西却认不出来，是 `response_unreadable` + pub(super) recognized: bool, +} + +/// 读一份回答。`freeform` 是客户端定义成自由格式的工具：别的格式的上游把它们的原文包在 +/// `{"input": …}` 里，拆出来，和客户端自己记下的一样。 +pub(super) fn read(body: &[u8], upstream: Dialect, freeform: &HashSet) -> Answer { + let mut a = match body.iter().find(|b| !b.is_ascii_whitespace()) { + None => Answer::default(), + Some(b'{') => whole(body, upstream), + Some(b'[') if upstream == Dialect::Gemini => gemini_array(body), + Some(_) => stream(body, upstream), + }; + tidy(&mut a.parts, upstream, freeform); + a +} + +/// 整包的回答。JSON 都解析不了的（截断了的）什么都读不出来 +fn whole(body: &[u8], upstream: Dialect) -> Answer { + let Ok(v) = serde_json::from_slice::(body) else { + return Answer::default(); + }; + let mut pieces = Vec::new(); + // 认得出是这家的回答就算读懂了，哪怕里面什么都没有（被拦下的提示、空的回答） + let says = |key: &str, prefix: &str| { + v.get(key) + .and_then(Value::as_str) + .is_some_and(|s| s.starts_with(prefix)) + }; + let recognized = match upstream { + Dialect::Anthropic => { + if let Some(blocks) = v.get("content").and_then(Value::as_array) { + read::anthropic_blocks(blocks, &mut pieces); + } + v.get("content").is_some() || says("type", "message") + } + Dialect::Chat => { + if let Some(m) = v + .get("choices") + .and_then(|c| c.get(0)) + .and_then(|c| c.get("message")) + { + pieces = read::chat_assistant(m); + } + v.get("choices").is_some() || says("object", "chat.completion") + } + Dialect::Responses => { + for it in v + .get("output") + .and_then(Value::as_array) + .into_iter() + .flatten() + { + pieces.extend(read::responses_item(it).pieces); + } + v.get("output").is_some() || says("object", "response") + } + Dialect::Gemini => { + if let Some(parts) = v + .get("candidates") + .and_then(|c| c.get(0)) + .and_then(|c| c.get("content")) + .and_then(|c| c.get("parts")) + .and_then(Value::as_array) + { + read::gemini_parts(parts, &mut pieces); + } + ["candidates", "promptFeedback", "usageMetadata"] + .iter() + .any(|k| v.get(*k).is_some()) + } + // 客户端不说 Converse，没有它的消息读法：整包交给 tw-dialect 解 + Dialect::Bedrock => { + if v.get("output").and_then(|o| o.get("message")).is_none() { + return Answer::default(); + } + let parts = tw_dialect::bedrock::decode_response(&v) + .blocks + .into_iter() + .map(block) + .collect(); + return Answer { + parts, + recognized: true, + }; + } + }; + Answer { + parts: pieces.iter().map(super::part).collect(), + recognized, + } +} + +fn block(b: Block) -> TranscriptPart { + match b { + Block::Text(text) => TranscriptPart::Text { text }, + Block::Thinking(t) => TranscriptPart::Thinking { text: t.text }, + Block::ToolCall(c) => TranscriptPart::ToolCall { + id: c.id, + name: c.name, + input: match c.input { + ToolInput::Text(s) => s, + json => json.to_json_text(), + }, + }, + } +} + +/// 流里的事件拼回一块一块。块按第一次出现的先后排 +#[derive(Default)] +struct Assembly { + parts: Vec, + /// 块号 → 在 `parts` 里的位置 + at: HashMap, + recognized: bool, +} + +impl Assembly { + fn other(&mut self, label: &str) { + self.push(TranscriptPart::Other { + label: label.to_string(), + }); + } + + /// 解析器不认、这里另外认出来的一块。**之后的字另起一块**：解析器那边的文字块可能还 + /// 开着（Gemini 每帧都往同一块里接），接回前面那块的话，字就跑到这一块前头去了 + fn push(&mut self, part: TranscriptPart) { + self.recognized = true; + self.parts.push(part); + let parts = &self.parts; + self.at.retain(|_, &mut i| { + !matches!( + parts[i], + TranscriptPart::Text { .. } | TranscriptPart::Thinking { .. } + ) + }); + } + + fn events(&mut self, events: Vec) { + for e in events { + self.event(e); + } + } + + fn event(&mut self, e: Event) { + self.recognized = true; + match e { + Event::BlockStart { index, kind } => { + self.at.insert(index, self.parts.len()); + self.parts.push(match kind { + BlockKind::Text => TranscriptPart::Text { + text: String::new(), + }, + BlockKind::Thinking => TranscriptPart::Thinking { + text: String::new(), + }, + BlockKind::ToolCall { id, name } => TranscriptPart::ToolCall { + id, + name, + input: String::new(), + }, + }); + } + Event::Delta { index, delta } => { + if let Some(&i) = self.at.get(&index) + && append(&mut self.parts[i], &delta) + { + return; + } + // 没报开始的块（Bedrock 的文字和推理就不报）：第一段增量到了就算开了一块 + self.at.insert(index, self.parts.len()); + self.parts.push(match delta { + Delta::Text(text) => TranscriptPart::Text { text }, + Delta::Thinking(text) => TranscriptPart::Thinking { text }, + Delta::Signature(_) => TranscriptPart::Thinking { + text: String::new(), + }, + Delta::ToolInput(input) => TranscriptPart::ToolCall { + id: String::new(), + name: String::new(), + input, + }, + }); + } + // 用量、结束原因、上游在流里报的错：不是回答的内容。错误在会话详情里有 + Event::Start { .. } + | Event::BlockStop { .. } + | Event::Usage(_) + | Event::Stop(_) + | Event::Error { .. } => {} + } + } + + fn finish(self) -> Answer { + Answer { + parts: self.parts, + recognized: self.recognized, + } + } +} + +/// 一段增量接到它那一块上。块的种类对不上的接不上 +fn append(part: &mut TranscriptPart, delta: &Delta) -> bool { + match (part, delta) { + (TranscriptPart::Text { text }, Delta::Text(d)) + | (TranscriptPart::Thinking { text }, Delta::Thinking(d)) + | (TranscriptPart::ToolCall { input: text, .. }, Delta::ToolInput(d)) => { + text.push_str(d); + true + } + // 签名不交出去：只有签名的推理块是一块空的推理 + (TranscriptPart::Thinking { .. }, Delta::Signature(_)) => true, + _ => false, + } +} + +/// tw-dialect 各家的流解析器 +enum Parser { + Anthropic(tw_dialect::anthropic::stream::Parser), + Chat(tw_dialect::chat::stream::Parser), + Responses(tw_dialect::responses::stream::Parser), + Gemini(tw_dialect::gemini::stream::Parser), + Bedrock(tw_dialect::bedrock::stream::Parser), +} + +impl Parser { + fn new(upstream: Dialect) -> Parser { + match upstream { + Dialect::Anthropic => Parser::Anthropic(Default::default()), + Dialect::Chat => Parser::Chat(Default::default()), + Dialect::Responses => Parser::Responses(Default::default()), + Dialect::Gemini => Parser::Gemini(Default::default()), + Dialect::Bedrock => Parser::Bedrock(Default::default()), + } + } + + fn finish(&mut self, out: &mut Vec) { + match self { + Parser::Anthropic(_) => {} + Parser::Chat(p) => p.finish(out), + Parser::Responses(p) => p.finish(out), + Parser::Gemini(p) => p.finish(out), + Parser::Bedrock(p) => p.finish(out), + } + } +} + +/// SSE 的流。**截断了的读到哪儿算哪儿**：最后半帧解析不了，就停在它前面 +fn stream(body: &[u8], upstream: Dialect) -> Answer { + let mut asm = Assembly::default(); + let mut parser = Parser::new(upstream); + let mut rest = body; + while !rest.is_empty() { + // 整份都在手里，不必像收流时那样攒着：按帧的边界一段一段切 + let (raw, next) = match frame::frame_end(rest) { + Some((end, sep)) => (&rest[..end], &rest[end + sep..]), + None => (rest, &rest[rest.len()..]), + }; + rest = next; + let Some(f) = frame::parse(raw) else { + continue; + }; + let mut events = Vec::new(); + match &mut parser { + Parser::Anthropic(p) => { + if let Some(label) = anthropic_other(&f) { + asm.other(&label); + } + p.frame(&f, &mut events); + } + Parser::Responses(p) => { + if let Some(label) = responses_other(&f) { + asm.other(&label); + } + p.frame(&f, &mut events); + } + Parser::Gemini(p) => { + // 自己解析这一帧，解析器也用这一份：不解析两遍 + if let Ok(v) = serde_json::from_str::(&f.data) { + gemini_chunk(&v, p, &mut asm, &mut events); + } + } + Parser::Chat(p) => p.frame(&f, &mut events), + Parser::Bedrock(p) => p.frame(&f, &mut events), + } + asm.events(events); + } + let mut events = Vec::new(); + parser.finish(&mut events); + asm.events(events); + asm.finish() +} + +/// 解析器跳过的块（服务端工具的调用和结果、MCP……）：在它开始的那一帧上认出来 +fn anthropic_other(f: &Frame) -> Option { + if !f.data.contains("\"content_block_start\"") { + return None; + } + let v: Value = serde_json::from_str(&f.data).ok()?; + let kind = v.get("content_block")?.get("type")?.as_str()?; + (!matches!(kind, "text" | "thinking" | "redacted_thinking" | "tool_use")) + .then(|| kind.to_string()) +} + +/// 解析器跳过的输出项（托管工具的调用：web_search_call、image_generation_call……） +fn responses_other(f: &Frame) -> Option { + if !f.data.contains("\"response.output_item.added\"") { + return None; + } + let v: Value = serde_json::from_str(&f.data).ok()?; + let kind = v.get("item")?.get("type")?.as_str()?; + (!matches!( + kind, + "message" | "reasoning" | "function_call" | "custom_tool_call" + )) + .then(|| kind.to_string()) +} + +/// Gemini 的一帧（数组里的一个元素）。解析器只认文字、推理和函数调用；图片和代码执行 +/// 在这里认，排在这一帧的文字前面 +fn gemini_chunk( + v: &Value, + p: &mut tw_dialect::gemini::stream::Parser, + asm: &mut Assembly, + events: &mut Vec, +) { + let parts = v + .get("candidates") + .and_then(|c| c.get(0)) + .and_then(|c| c.get("content")) + .and_then(|c| c.get("parts")) + .and_then(Value::as_array) + .map(Vec::as_slice) + .unwrap_or(&[]); + let mut pieces = Vec::new(); + read::gemini_parts(parts, &mut pieces); + for piece in &pieces { + if matches!(piece, read::Piece::Image { .. } | read::Piece::Other(_)) { + asm.push(super::part(piece)); + } + } + p.chunk(v, events); +} + +/// 不带 `alt=sse` 的 Gemini 流：一个逐步写出的 JSON 数组，一个元素是一帧。截断了的 +/// 读到最后一个完整的元素 +fn gemini_array(body: &[u8]) -> Answer { + let mut asm = Assembly::default(); + let mut parser = tw_dialect::gemini::stream::Parser::default(); + let start = body + .iter() + .position(|b| *b == b'[') + .map_or(body.len(), |i| i + 1); + let mut rest = &body[start..]; + loop { + let skip = rest + .iter() + .position(|b| !(b.is_ascii_whitespace() || *b == b',')) + .unwrap_or(rest.len()); + rest = &rest[skip..]; + if rest.first().is_none_or(|b| *b == b']') { + break; + } + let mut one = serde_json::Deserializer::from_slice(rest).into_iter::(); + let Some(Ok(v)) = one.next() else { + break; + }; + let used = one.byte_offset(); + let mut events = Vec::new(); + gemini_chunk(&v, &mut parser, &mut asm, &mut events); + asm.events(events); + rest = &rest[used..]; + } + let mut events = Vec::new(); + parser.finish(&mut events); + asm.events(events); + asm.finish() +} + +/// 收拾一下：空的文字块不算；没有参数的函数调用是 `{}`；别的格式的上游包起来的自由格式 +/// 原文拆出来（和网关转换时一样，见 tw-dialect 的 `Normalizer`） +fn tidy(parts: &mut Vec, upstream: Dialect, freeform: &HashSet) { + parts.retain(|p| !matches!(p, TranscriptPart::Text { text } if text.is_empty())); + for p in parts { + let TranscriptPart::ToolCall { name, input, .. } = p else { + continue; + }; + if freeform.contains(name.as_str()) { + if upstream != Dialect::Responses { + *input = unwrap_freeform(input); + } + } else if input.trim().is_empty() { + *input = "{}".into(); + } + } +} + +fn unwrap_freeform(json_text: &str) -> String { + match serde_json::from_str::(json_text) { + Ok(Value::Object(o)) => match o.get("input") { + Some(Value::String(s)) => s.clone(), + _ => json_text.to_string(), + }, + Ok(Value::String(s)) => s, + _ => json_text.to_string(), + } +} diff --git a/crates/tw-store/src/transcript/mod.rs b/crates/tw-store/src/transcript/mod.rs new file mode 100644 index 00000000..d5a8e676 --- /dev/null +++ b/crates/tw-store/src/transcript/mod.rs @@ -0,0 +1,1932 @@ +//! 一次会话读成一段对话（`GET /sessions/{id}/transcript`）。 +//! +//! Claude Code 这类客户端每一轮都把整段历史发上来：请求 i 的消息 = 请求 i-1 的消息 + 上一轮 +//! 的回答 + 新的用户消息或工具结果。所以按会话里的顺序一个一个往下读，**每个请求只交出它 +//! 新带来的那几条**；回答从存下来的响应里读（见 [`answer`]）。 +//! +//! # 怎么认「新的那几条」 +//! +//! 上一个读得懂的请求的消息是这一个的开头（逐条比指纹），剩下的就是新的；剩下的第一条是 +//! 助手消息的话，它是上一轮的回答，已经在上一轮的 `output` 里了，去掉。对不上（压缩过、 +//! 改过历史），或者上一个请求读不懂，就交出这个请求的整段历史，标 `restart`。 +//! +//! 比的是规整过的样子（见 [`read::Piece`]）：缓存断点、推理签名、键的顺序都不算。**推理 +//! 整个不算**：有的客户端到了下一轮就把前面的推理去掉（接口本来也不读它们），算上的话， +//! 这样的客户端每一轮都对不上。 +//! +//! # 读多少、留多少 +//! +//! 几百轮的会话，每个请求几 MB。每份正文读一次、解析一次，读完这一轮就丢；从一个请求留到 +//! 下一个的只有每条消息 16 字节的指纹和系统提示。 + +mod answer; +mod read; + +use std::collections::HashSet; + +use serde_json::Value; +use tw_api::{TranscriptGap, TranscriptMessage, TranscriptPart, TranscriptRole, TranscriptTurn}; +use tw_dialect::ir::Dialect; + +use crate::blobs::{Blobs, Which}; +use crate::db::RequestRow; +use crate::search::text::{client_dialect, dialect_of}; +use read::{Input, Piece, Role}; + +/// 读成一段对话。`rows` 是这次会话的请求，和会话详情同样的顺序 +/// （[`crate::Db::session_requests`]）。 +/// +/// **已脱敏**，和请求详情里的正文同一套打码（`tw_secret::mask_body`）。 +pub fn build(session: &str, rows: &[RequestRow], blobs: &Blobs) -> tw_api::Transcript { + let mut t = tw_api::Transcript { + session: session.to_string(), + system: None, + turns: Vec::with_capacity(rows.len()), + }; + let mut chain = Chain::default(); + for row in rows { + chain.turn(row, blobs, &mut t); + } + mask(&mut t); + t +} + +/// 一条消息比对用的指纹 +type Fp = [u8; 16]; + +/// 读到哪儿了：上一个读得懂的请求留下来的东西。 +#[derive(Default)] +struct Chain { + /// 它的每条消息的指纹（只有推理的消息不比，不在里面）。还没有读得懂的请求是 None + prev: Option>, + /// 它的系统提示 + system: Option, + /// 它之后有没有读不懂的请求 + broken: bool, + /// 它是第几轮，回答完整地交出去了没有 + last: Option<(usize, bool)>, +} + +impl Chain { + fn turn(&mut self, row: &RequestRow, blobs: &Blobs, t: &mut tw_api::Transcript) { + let mut turn = TranscriptTurn { + id: row.id.to_string(), + restart: false, + system_changed: None, + input: Vec::new(), + output: Vec::new(), + gaps: Vec::new(), + }; + // 不生成回答的调用（数 token、Responses 的压缩）：问的是这段对话，不是对话里的一句。 + // 不读，也不和前后比 + let Some(client) = client_dialect(&row.path) else { + t.turns.push(turn); + return; + }; + let mut freeform = HashSet::new(); + let readable = match request(row, blobs) { + Ok(v) => { + self.read(client, &v, &mut turn, t, &mut freeform); + true + } + Err(gap) => { + turn.gaps.push(gap); + self.broken = true; + false + } + }; + let (output, gap) = output(row, blobs, upstream_of(row, client), &freeform); + turn.output = output; + turn.gaps.extend(gap); + if readable { + let shown = !turn.output.is_empty() + && !turn.gaps.iter().any(|g| { + matches!( + g, + TranscriptGap::ResponseMissing + | TranscriptGap::ResponseTruncated + | TranscriptGap::ResponseUnreadable + ) + }); + self.last = Some((t.turns.len(), shown)); + } + t.turns.push(turn); + } + + /// 读得懂的请求：系统提示，新的消息 + fn read( + &mut self, + client: Dialect, + v: &Value, + turn: &mut TranscriptTurn, + t: &mut tw_api::Transcript, + freeform: &mut HashSet, + ) { + let read::Body { + system, + items, + freeform: tools, + } = read::body(client, v); + freeform.extend(tools); + let system = system + .iter() + .map(|s| s.as_ref()) + .filter(|s| !s.is_empty()) + .collect::>() + .join("\n\n"); + match &self.system { + None => t.system = (!system.is_empty()).then(|| system.clone()), + Some(prev) if *prev != system => turn.system_changed = Some(system.clone()), + Some(_) => {} + } + + let messages = merge(items); + let fps: Vec<(usize, Fp)> = messages + .iter() + .enumerate() + .filter_map(|(i, m)| fp(m).map(|f| (i, f))) + .collect(); + let start = match (&self.prev, self.broken) { + // 前面有读不懂的：它带来的那几条没交出去过，从头交 + (_, true) => None, + (None, false) => Some(0), + (Some(prev), false) => continues(prev, &fps), + }; + turn.restart = start.is_none(); + let mut new = &messages[start.unwrap_or(0)..]; + // 剩下的第一条助手消息是上一轮的回答，已经在上一轮的 `output` 里。**那一轮的回答没有 + // 完整交出去的不去掉**（没存下、只存了开头、读不懂）：这一条就是它说过什么的记录 + if start.is_some() + && let Some((at, true)) = self.last + && let Some(said) = new.first().filter(|m| m.role == Role::Assistant) + { + same_calls(&mut t.turns[at].output, said); + new = &new[1..]; + } + turn.input = new.iter().map(message).collect(); + + self.prev = Some(fps.into_iter().map(|(_, f)| f).collect()); + self.system = Some(system); + self.broken = false; + } +} + +/// 读请求体：没存下来的、只存了开头的、解析不了的各是一种缺口。 +fn request(row: &RequestRow, blobs: &Blobs) -> Result { + let raw = blobs + .get(row.at_ms, row.id, Which::Request) + .ok_or(TranscriptGap::RequestMissing)?; + if blobs + .original_len(row.at_ms, row.id, Which::Request) + .is_some_and(|n| n > raw.len()) + { + return Err(TranscriptGap::RequestTruncated); + } + serde_json::from_slice(&raw).map_err(|_| TranscriptGap::RequestTruncated) +} + +/// 回答它的那一家说的格式：转换过的记在行上，直通的就是客户端那一种（和按正文找一样） +fn upstream_of(row: &RequestRow, client: Dialect) -> Dialect { + row.translated + .as_deref() + .and_then(|j| serde_json::from_str::(j).ok()) + .map_or(client, |t| dialect_of(t.to)) +} + +/// 这一轮的回答，和读不出来的那个缺口。 +fn output( + row: &RequestRow, + blobs: &Blobs, + upstream: Dialect, + freeform: &HashSet, +) -> (Vec, Option) { + // 上游接下了、回了 2xx 的才有回答可读 + let answered = row.status.is_some_and(|s| (200..300).contains(&s)); + let Some(body) = blobs.get(row.at_ms, row.id, Which::Response) else { + // 没走到上游的、上游回了错误的，本来就没有回答；一个字节都没收到客户端就走了的也是 + let missing = answered && row.bytes != Some(0); + return ( + Vec::new(), + missing.then_some(TranscriptGap::ResponseMissing), + ); + }; + // 上游回了错误：存下来的是错误，不是回答。失败在会话详情里已经有了，这里不重复 + if !answered { + return (Vec::new(), None); + } + let truncated = blobs + .original_len(row.at_ms, row.id, Which::Response) + .is_some_and(|n| n > body.len()); + let a = answer::read(&body, upstream, freeform); + let gap = if truncated { + Some(TranscriptGap::ResponseTruncated) + } else if !a.recognized { + Some(TranscriptGap::ResponseUnreadable) + } else { + None + }; + (a.parts, gap) +} + +/// 一条消息。同一个角色连着的几条并成一条：Responses 里推理、文字、几个调用是几个输入项， +/// 可它们是同一次回答；Chat 里几个调用的结果是几条 `tool` 消息。用户和系统消息不并 +struct Message<'a> { + role: Role, + pieces: Vec>, +} + +fn merge(items: Vec>) -> Vec> { + let mut out: Vec = Vec::with_capacity(items.len()); + for it in items { + if it.pieces.is_empty() { + continue; + } + match out.last_mut() { + Some(last) + if last.role == it.role && matches!(it.role, Role::Assistant | Role::Tool) => + { + last.pieces.extend(it.pieces) + } + _ => out.push(Message { + role: it.role, + pieces: it.pieces, + }), + } + } + out +} + +/// 一条消息比对用的指纹。推理不算（见模块的说明）；除了推理什么都没有的消息不比，是 None。 +fn fp(m: &Message) -> Option { + fn field(h: &mut blake3::Hasher, s: &str) { + h.update(&(s.len() as u64).to_le_bytes()); + h.update(s.as_bytes()); + } + let mut h = blake3::Hasher::new(); + h.update(&[m.role as u8]); + let mut any = false; + for p in &m.pieces { + match p { + Piece::Thinking(_) => continue, + Piece::Text(t) => { + h.update(&[1]); + field(&mut h, t); + } + Piece::ToolCall { id, name, input } => { + h.update(&[2]); + field(&mut h, id); + field(&mut h, name); + match input { + Input::Args(s) | Input::Raw(s) => { + h.update(&[0]); + field(&mut h, s); + } + // 键按字母排着写出来（serde_json 的对象就是这么存的）：顺序不算 + Input::Json(v) => { + h.update(&[1]); + let _ = serde_json::to_writer(&mut h, v); + } + } + } + Piece::ToolResult { + call_id, + text, + is_error, + } => { + h.update(&[3, *is_error as u8]); + field(&mut h, call_id); + field(&mut h, text); + } + Piece::Image { + media_type, + bytes, + data, + } => { + h.update(&[4]); + field(&mut h, media_type.unwrap_or_default()); + h.update(&bytes.unwrap_or(u64::MAX).to_le_bytes()); + field(&mut h, data.unwrap_or_default()); + } + Piece::Other(label) => { + h.update(&[5]); + field(&mut h, label); + } + } + any = true; + } + let mut out = [0; 16]; + out.copy_from_slice(&h.finalize().as_bytes()[..16]); + any.then_some(out) +} + +/// 上一个请求的消息是不是这一个的开头。是的话，新的从第几条开始 +fn continues(prev: &[Fp], cur: &[(usize, Fp)]) -> Option { + if cur.len() < prev.len() || cur.iter().zip(prev).any(|((_, a), b)| a != b) { + return None; + } + Some(prev.len().checked_sub(1).map_or(0, |last| cur[last].0 + 1)) +} + +/// 上一轮回答里工具调用的号，换成客户端记下的那个。 +/// +/// 转换过格式的回答里，号是上游给的（Gemini 根本不给，网关和这里各自现编一个）；下一轮的 +/// 工具结果认的是客户端拿到的那个号。名字一个一个对得上才换。 +fn same_calls(output: &mut [TranscriptPart], said: &Message) { + let theirs: Vec<(&str, &str)> = said + .pieces + .iter() + .filter_map(|p| match p { + Piece::ToolCall { id, name, .. } => Some((id.as_ref(), name.as_ref())), + _ => None, + }) + .collect(); + let mine: Vec<(&mut String, &mut String)> = output + .iter_mut() + .filter_map(|p| match p { + TranscriptPart::ToolCall { id, name, .. } => Some((id, name)), + _ => None, + }) + .collect(); + if mine.len() != theirs.len() + || mine + .iter() + .zip(&theirs) + .any(|((_, name), (_, said))| name.as_str() != *said) + { + return; + } + for ((id, _), (said, _)) in mine.into_iter().zip(theirs) { + if id != said { + *id = said.to_string(); + } + } +} + +fn message(m: &Message) -> TranscriptMessage { + TranscriptMessage { + role: match m.role { + Role::User => TranscriptRole::User, + Role::Assistant => TranscriptRole::Assistant, + Role::Tool => TranscriptRole::Tool, + Role::System => TranscriptRole::System, + }, + parts: m.pieces.iter().map(part).collect(), + } +} + +/// 一块交出去的样子。**还没打码**：打码在最后一起做 +fn part(p: &Piece) -> TranscriptPart { + match p { + Piece::Text(t) => TranscriptPart::Text { + text: t.to_string(), + }, + Piece::Thinking(t) => TranscriptPart::Thinking { + text: t.to_string(), + }, + Piece::ToolCall { id, name, input } => TranscriptPart::ToolCall { + id: id.to_string(), + name: name.to_string(), + input: input.text(), + }, + Piece::ToolResult { + call_id, + text, + is_error, + } => TranscriptPart::ToolResult { + call_id: call_id.to_string(), + text: text.to_string(), + is_error: *is_error, + }, + Piece::Image { + media_type, bytes, .. + } => TranscriptPart::Image { + media_type: media_type.map(str::to_string), + bytes: *bytes, + }, + Piece::Other(label) => TranscriptPart::Other { + label: label.to_string(), + }, + } +} + +/// 打码，和请求详情里的正文同一套。**最后一起做**：前后比对、换工具调用的号都要原文 +fn mask(t: &mut tw_api::Transcript) { + fn m(s: &mut String) { + if !s.is_empty() { + *s = tw_secret::mask_body(s); + } + } + if let Some(s) = &mut t.system { + m(s); + } + for turn in &mut t.turns { + if let Some(s) = &mut turn.system_changed { + m(s); + } + let parts = turn + .input + .iter_mut() + .flat_map(|msg| msg.parts.iter_mut()) + .chain(turn.output.iter_mut()); + for p in parts { + match p { + TranscriptPart::Text { text } | TranscriptPart::Thinking { text } => m(text), + TranscriptPart::ToolCall { id, name, input } => { + m(id); + m(name); + m(input); + } + TranscriptPart::ToolResult { call_id, text, .. } => { + m(call_id); + m(text); + } + TranscriptPart::Image { media_type, .. } => { + if let Some(s) = media_type { + m(s); + } + } + TranscriptPart::Other { label } => m(label), + } + } + } +} + +#[cfg(test)] +mod tests { + use serde_json::{Value, json}; + use tw_api::{ + Transcript, TranscriptGap as Gap, TranscriptMessage, TranscriptPart as P, + TranscriptRole as R, + }; + + use super::*; + use crate::db::Db; + use crate::db::tests::row; + + const NOW: i64 = 1_790_000_000_000; + + /// 一份测试用的盘：请求库和正文目录，记录都归在会话 `s` 里 + struct Disk { + _dir: tempfile::TempDir, + db: Db, + blobs: Blobs, + next: i64, + } + + impl Disk { + fn new() -> Disk { + let dir = tempfile::tempdir().unwrap(); + Disk { + blobs: Blobs::new(dir.path().join("blobs")), + db: Db::in_memory().unwrap(), + _dir: dir, + next: 0, + } + } + + /// 一轮：请求体和回答都存下了 + fn turn(&mut self, path: &str, request: &Value, response: &[u8]) -> i64 { + let request = request.to_string(); + self.put(path, Some(request.as_bytes()), Some(response), |_| {}) + } + + /// 一轮，存下了哪些自己定；`edit` 改这一行的记录 + fn put( + &mut self, + path: &str, + request: Option<&[u8]>, + response: Option<&[u8]>, + edit: impl FnOnce(&mut RequestRow), + ) -> i64 { + self.next += 1; + let id = self.next; + let mut r = row(id, NOW + id * 1000); + r.path = path.into(); + r.session = Some("s".into()); + edit(&mut r); + self.db.insert(&r).unwrap(); + if let Some(b) = request { + assert!(self.blobs.put(r.at_ms, id, Which::Request, b)); + } + if let Some(b) = response { + assert!(self.blobs.put(r.at_ms, id, Which::Response, b)); + } + id + } + + /// 存下来的只是开头：原本有 `len` 那么长 + fn cut(&self, id: i64, which: Which, len: usize) { + let at = NOW + id * 1000; + let body = self.blobs.get(at, id, which).unwrap(); + assert!(self.blobs.put_with_len(at, id, which, &body, len)); + } + + fn transcript(&self) -> Transcript { + let rows = self.db.session_requests("s").unwrap(); + build("s", &rows, &self.blobs) + } + } + + fn sse(frames: &[(&str, Value)]) -> Vec { + let mut s = String::new(); + for (event, data) in frames { + if !event.is_empty() { + s.push_str(&format!("event: {event}\n")); + } + s.push_str(&format!("data: {data}\n\n")); + } + s.into_bytes() + } + + /// 一段字切成两半，按字符切 + fn halves(s: &str) -> (String, String) { + let n = s.chars().count() / 2; + (s.chars().take(n).collect(), s.chars().skip(n).collect()) + } + + /// Anthropic 的流：给几块内容，写成每块的开始、两段增量、结束。服务端工具那样的块只有 + /// 开始和结束 + fn anthropic_stream(blocks: &[Value]) -> Vec { + let mut f = vec![( + "message_start", + json!({"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-5", + "usage": {"input_tokens": 10}}}), + )]; + for (i, b) in blocks.iter().enumerate() { + let delta = |d: Value| { + ( + "content_block_delta", + json!({"type": "content_block_delta", "index": i, "delta": d}), + ) + }; + let start = |block: Value| { + ( + "content_block_start", + json!({"type": "content_block_start", "index": i, "content_block": block}), + ) + }; + match b["type"].as_str().unwrap() { + "text" => { + f.push(start(json!({"type": "text", "text": ""}))); + let (a, z) = halves(b["text"].as_str().unwrap()); + f.push(delta(json!({"type": "text_delta", "text": a}))); + f.push(delta(json!({"type": "text_delta", "text": z}))); + } + "thinking" => { + f.push(start(json!({"type": "thinking", "thinking": ""}))); + f.push(delta( + json!({"type": "thinking_delta", "thinking": b["thinking"]}), + )); + f.push(delta( + json!({"type": "signature_delta", "signature": b["signature"]}), + )); + } + "tool_use" => { + f.push(start( + json!({"type": "tool_use", "id": b["id"], "name": b["name"], "input": {}}), + )); + let (a, z) = halves(&b["input"].to_string()); + f.push(delta( + json!({"type": "input_json_delta", "partial_json": a}), + )); + f.push(delta( + json!({"type": "input_json_delta", "partial_json": z}), + )); + } + _ => f.push(start(b.clone())), + } + f.push(( + "content_block_stop", + json!({"type": "content_block_stop", "index": i}), + )); + } + f.push(( + "message_delta", + json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 9}}), + )); + f.push(("message_stop", json!({"type": "message_stop"}))); + sse(&f) + } + + fn text(t: &str) -> P { + P::Text { text: t.into() } + } + + fn thinking(t: &str) -> P { + P::Thinking { text: t.into() } + } + + fn call(id: &str, name: &str, input: &str) -> P { + P::ToolCall { + id: id.into(), + name: name.into(), + input: input.into(), + } + } + + fn result(call_id: &str, t: &str) -> P { + P::ToolResult { + call_id: call_id.into(), + text: t.into(), + is_error: false, + } + } + + fn msg(role: R, parts: Vec

) -> TranscriptMessage { + TranscriptMessage { role, parts } + } + + fn user(t: &str) -> Value { + json!({"role": "user", "content": t}) + } + + fn assistant(t: &str) -> Value { + json!({"role": "assistant", "content": t}) + } + + /// Anthropic 的一轮：一段 system、给出的消息，流式回答一段话 + fn anthropic(messages: &[Value]) -> Value { + json!({"model": "claude-sonnet-4-5", "system": "你是助手", "stream": true, "messages": messages}) + } + + fn says(t: &str) -> Vec { + anthropic_stream(&[json!({"type": "text", "text": t})]) + } + + fn with_cache(mut m: Value) -> Value { + if let Some(last) = m["content"].as_array_mut().and_then(|c| c.last_mut()) { + last["cache_control"] = json!({"type": "ephemeral"}); + } + m + } + + // ───────────────────────────────────────────────── 四种客户端格式 + + /// 照 Claude Code 的样子：系统提示是几块，缓存断点每一轮挪到最后一条上，推理带着签名， + /// 工具调用和结果，结果里有一张截图。每一轮只交出新的那几条 + #[test] + fn a_claude_code_session_reads_turn_by_turn() { + let mut d = Disk::new(); + let system = json!([ + {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.0.0"}, + {"type": "text", "text": "You are Claude Code.", "cache_control": {"type": "ephemeral"}} + ]); + let tools = json!([{"name": "Read", "input_schema": {"type": "object"}}, + {"name": "Edit", "input_schema": {"type": "object"}}]); + let body = |messages: Vec| { + json!({"model": "claude-sonnet-4-5", "system": system, "tools": tools, "max_tokens": 32000, + "stream": true, "messages": messages}) + }; + let u1 = json!({"role": "user", "content": [ + {"type": "text", "text": "上下文"}, + {"type": "text", "text": "修一下 main.rs 里的 bug"} + ]}); + let a1 = |signature: &str| { + json!({"role": "assistant", "content": [ + {"type": "thinking", "thinking": "先读文件", "signature": signature}, + {"type": "text", "text": "我先看看文件。"}, + {"type": "tool_use", "id": "toolu_1", "name": "Read", "input": {"file_path": "src/main.rs"}} + ]}) + }; + let r1 = json!({"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", + "content": "fn main() {\n let key = \"sk-ant-api03-SECRETSECRETSECRET\";\n}"}]}); + let a2 = |with_thinking: bool| { + let mut blocks = vec![ + json!({"type": "text", "text": "密钥写死了，我来改。"}), + json!({"type": "tool_use", "id": "toolu_2", "name": "Edit", + "input": {"file_path": "src/main.rs", "old": "let key", "new": "let key = env()"}}), + ]; + if with_thinking { + blocks.insert( + 0, + json!({"type": "thinking", "thinking": "要改成读环境变量", "signature": "sig-2"}), + ); + } + json!({"role": "assistant", "content": blocks}) + }; + let png = "QUJD".repeat(10); + let r2 = json!({"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "toolu_2", "content": [ + {"type": "text", "text": "已修改"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": png}} + ]}, + {"type": "text", "text": "看起来不错，再跑一下测试"} + ]}); + + d.turn( + "/v1/messages", + &body(vec![with_cache(u1.clone())]), + &anthropic_stream(&[ + json!({"type": "thinking", "thinking": "先读文件", "signature": "sig-1"}), + json!({"type": "text", "text": "我先看看文件。"}), + json!({"type": "tool_use", "id": "toolu_1", "name": "Read", "input": {"file_path": "src/main.rs"}}), + ]), + ); + d.turn( + "/v1/messages", + &body(vec![u1.clone(), a1("sig-1"), with_cache(r1.clone())]), + &anthropic_stream(&[ + json!({"type": "thinking", "thinking": "要改成读环境变量", "signature": "sig-2"}), + json!({"type": "text", "text": "密钥写死了，我来改。"}), + json!({"type": "tool_use", "id": "toolu_2", "name": "Edit", + "input": {"file_path": "src/main.rs", "old": "let key", "new": "let key = env()"}}), + ]), + ); + // 第三轮：第一轮推理的签名换了一个，第二轮的推理被客户端去掉了 —— 都不算改过历史 + d.turn( + "/v1/messages", + &body(vec![u1, a1("sig-1-again"), r1, a2(false), with_cache(r2)]), + &says("好的。"), + ); + + let t = d.transcript(); + assert_eq!(t.session, "s"); + assert_eq!( + t.system.as_deref(), + Some("x-anthropic-billing-header: cc_version=2.0.0\n\nYou are Claude Code.") + ); + assert_eq!(t.turns.len(), 3); + assert_eq!( + t.turns.iter().map(|x| x.id.as_str()).collect::>(), + ["1", "2", "3"] + ); + for x in &t.turns { + assert!(!x.restart, "{x:?}"); + assert_eq!(x.system_changed, None); + assert!(x.gaps.is_empty(), "{x:?}"); + } + + let first = &t.turns[0]; + assert_eq!( + first.input, + [msg( + R::User, + vec![ + text("上下文"), + text("修一下 main.rs 里的 bug") + ] + )] + ); + assert_eq!( + first.output, + [ + thinking("先读文件"), + text("我先看看文件。"), + call("toolu_1", "Read", r#"{"file_path":"src/main.rs"}"#) + ] + ); + + // 第二轮新的只有工具结果：上一轮的回答已经在上一轮里了。结果里的密钥打了码 + let second = &t.turns[1]; + assert_eq!(second.input.len(), 1, "{:?}", second.input); + assert_eq!(second.input[0].role, R::Tool); + let P::ToolResult { + call_id, + text: said, + is_error, + } = &second.input[0].parts[0] + else { + panic!("{:?}", second.input) + }; + assert_eq!((call_id.as_str(), *is_error), ("toolu_1", false)); + assert!(said.contains("sk-an…CRET"), "{said}"); + assert!(!said.contains("SECRETSECRET"), "{said}"); + assert_eq!(second.output[1], text("密钥写死了，我来改。")); + assert!(matches!(&second.output[2], P::ToolCall { name, .. } if name == "Edit")); + + // 第三轮：工具结果、结果里的截图（只有类型和大小）、一句话。有文字就不只是工具结果 + assert_eq!( + t.turns[2].input, + [msg( + R::User, + vec![ + result("toolu_2", "已修改"), + P::Image { + media_type: Some("image/png".into()), + bytes: Some(30) + }, + text("看起来不错，再跑一下测试") + ] + )] + ); + assert_eq!(t.turns[2].output, [text("好的。")]); + + // 图片的数据、推理的签名都不出门 + let wire = serde_json::to_string(&t).unwrap(); + assert!(!wire.contains(&png), "{wire}"); + assert!(!wire.contains("sig-1"), "{wire}"); + } + + /// Chat：`tool_calls` 和 `tool` 消息，对话中途的系统消息照原位置留着，回答有流式的也有整包的 + #[test] + fn an_openai_chat_session_with_tool_calls() { + let mut d = Disk::new(); + let sys = json!({"role": "system", "content": "你是助手"}); + let ask = user("列一下文件"); + let a1 = json!({"role": "assistant", "content": "我来看看", "reasoning_content": "想一想", + "tool_calls": [{"id": "call_1", "type": "function", + "function": {"name": "ls", "arguments": "{\"dir\":\".\"}"}}]}); + let r1 = json!({"role": "tool", "tool_call_id": "call_1", "content": "README.md"}); + let a2 = assistant("有一个 README.md。"); + let body = |messages: Vec| json!({"model": "deepseek-chat", "stream": true, "messages": messages}); + let chunk = |delta: Value, finish: Value| { + ( + "", + json!({"id": "c", "choices": [{"index": 0, "delta": delta, "finish_reason": finish}]}), + ) + }; + let mut stream = sse(&[ + chunk( + json!({"role": "assistant", "reasoning_content": "想一想"}), + Value::Null, + ), + chunk(json!({"content": "我来"}), Value::Null), + chunk(json!({"content": "看看"}), Value::Null), + chunk( + json!({"tool_calls": [{"index": 0, "id": "call_1", "type": "function", + "function": {"name": "ls", "arguments": ""}}]}), + Value::Null, + ), + chunk( + json!({"tool_calls": [{"index": 0, "function": {"arguments": "{\"dir\":"}}]}), + Value::Null, + ), + chunk( + json!({"tool_calls": [{"index": 0, "function": {"arguments": "\".\"}"}}]}), + Value::Null, + ), + chunk(json!({}), json!("tool_calls")), + ]); + stream.extend_from_slice(b"data: [DONE]\n\n"); + let whole = |t: &str| { + json!({"choices": [{"message": {"role": "assistant", "content": t}, "finish_reason": "stop"}]}) + .to_string() + .into_bytes() + }; + d.turn( + "/v1/chat/completions", + &body(vec![sys.clone(), ask.clone()]), + &stream, + ); + d.turn( + "/v1/chat/completions", + &body(vec![sys.clone(), ask.clone(), a1.clone(), r1.clone()]), + &whole("有一个 README.md。"), + ); + d.turn( + "/v1/chat/completions", + &body(vec![ + sys, + ask, + a1, + r1, + a2, + json!({"role": "system", "content": "之后用英文回答"}), + user("谢谢"), + ]), + &whole("You're welcome."), + ); + + let t = d.transcript(); + assert_eq!(t.system.as_deref(), Some("你是助手")); + assert_eq!(t.turns[0].input, [msg(R::User, vec![text("列一下文件")])]); + assert_eq!( + t.turns[0].output, + [ + thinking("想一想"), + text("我来看看"), + call("call_1", "ls", r#"{"dir":"."}"#) + ] + ); + assert_eq!( + t.turns[1].input, + [msg(R::Tool, vec![result("call_1", "README.md")])] + ); + assert_eq!(t.turns[1].output, [text("有一个 README.md。")]); + assert_eq!( + t.turns[2].input, + [ + msg(R::System, vec![text("之后用英文回答")]), + msg(R::User, vec![text("谢谢")]) + ] + ); + assert_eq!(t.turns[2].output, [text("You're welcome.")]); + assert!(t.turns.iter().all(|x| !x.restart && x.gaps.is_empty())); + } + + /// Responses（Codex）：instructions 和开头的 developer 消息是系统提示；推理项、函数调用、 + /// 结果是一个个输入项，一次回答的几项并成一条；托管工具的调用是 `other` + #[test] + fn a_codex_session_on_the_responses_format() { + let mut d = Disk::new(); + let dev = json!({"type": "message", "role": "developer", + "content": [{"type": "input_text", "text": "sandbox: workspace-write"}]}); + let env = json!({"type": "message", "role": "user", + "content": [{"type": "input_text", "text": "cwd"}]}); + let ask = json!({"type": "message", "role": "user", + "content": [{"type": "input_text", "text": "修好构建"}]}); + let args = r#"{"command":["cargo","build"]}"#; + let rs1 = json!({"type": "reasoning", "id": "rs_1", + "summary": [{"type": "summary_text", "text": "看看报错"}], "encrypted_content": "gAAA1"}); + let fc1 = + json!({"type": "function_call", "call_id": "c1", "name": "shell", "arguments": args}); + let out1 = json!({"type": "function_call_output", "call_id": "c1", + "output": "error[E0308]: mismatched types"}); + let patch = "*** Begin Patch\n*** End Patch"; + let rs2 = + json!({"type": "reasoning", "id": "rs_2", "summary": [], "encrypted_content": "gAAA2"}); + let ws = json!({"type": "web_search_call", "id": "ws_1", "status": "completed", + "action": {"type": "search", "query": "E0308"}}); + let cc = json!({"type": "custom_tool_call", "call_id": "c2", "name": "apply_patch", "input": patch}); + let out2 = json!({"type": "custom_tool_call_output", "call_id": "c2", "output": "Done!"}); + let body = |input: Vec| { + json!({"model": "gpt-5.5", "instructions": "You are Codex.", "input": input, "stream": true, + "tools": [{"type": "function", "name": "shell", "parameters": {"type": "object"}}, + {"type": "custom", "name": "apply_patch"}]}) + }; + let ev = |name: &str, mut data: Value| { + data["type"] = json!(name); + (name.to_string(), data) + }; + let stream = |frames: Vec<(String, Value)>| { + let frames: Vec<(&str, Value)> = frames + .iter() + .map(|(n, v)| (n.as_str(), v.clone())) + .collect(); + sse(&frames) + }; + let first = stream(vec![ + ev( + "response.created", + json!({"response": {"id": "resp_1", "model": "gpt-5.5"}}), + ), + ev( + "response.output_item.added", + json!({"output_index": 0, "item": {"type": "reasoning", "id": "rs_1", "summary": []}}), + ), + ev( + "response.reasoning_summary_text.delta", + json!({"output_index": 0, "summary_index": 0, "delta": "看看报错"}), + ), + ev( + "response.output_item.done", + json!({"output_index": 0, "item": rs1}), + ), + ev( + "response.output_item.added", + json!({"output_index": 1, "item": {"type": "function_call", "call_id": "c1", "name": "shell", "arguments": ""}}), + ), + ev( + "response.function_call_arguments.delta", + json!({"output_index": 1, "delta": "{\"command\":"}), + ), + ev( + "response.function_call_arguments.delta", + json!({"output_index": 1, "delta": "[\"cargo\",\"build\"]}"}), + ), + ev( + "response.output_item.done", + json!({"output_index": 1, "item": fc1}), + ), + ev( + "response.completed", + json!({"response": {"usage": {"input_tokens": 1, "output_tokens": 2}}}), + ), + ]); + // 第二轮：推理只有密文（没有摘要），一次网页搜索，一个自由格式的工具调用 + let second = stream(vec![ + ev( + "response.created", + json!({"response": {"id": "resp_2", "model": "gpt-5.5"}}), + ), + ev( + "response.output_item.added", + json!({"output_index": 0, "item": {"type": "reasoning", "id": "rs_2", "summary": []}}), + ), + ev( + "response.output_item.done", + json!({"output_index": 0, "item": rs2}), + ), + ev( + "response.output_item.added", + json!({"output_index": 1, "item": {"type": "web_search_call", "id": "ws_1", "status": "in_progress"}}), + ), + ev( + "response.output_item.done", + json!({"output_index": 1, "item": ws}), + ), + ev( + "response.output_item.added", + json!({"output_index": 2, "item": {"type": "custom_tool_call", "call_id": "c2", "name": "apply_patch", "input": ""}}), + ), + ev( + "response.custom_tool_call_input.delta", + json!({"output_index": 2, "delta": "*** Begin Patch\n"}), + ), + ev( + "response.custom_tool_call_input.delta", + json!({"output_index": 2, "delta": "*** End Patch"}), + ), + ev( + "response.output_item.done", + json!({"output_index": 2, "item": cc}), + ), + ev( + "response.completed", + json!({"response": {"usage": {"input_tokens": 1, "output_tokens": 2}}}), + ), + ]); + let whole = json!({"output": [{"type": "message", "role": "assistant", + "content": [{"type": "output_text", "text": "构建好了。"}]}]}); + let path = "/backend-api/codex/responses"; + d.turn( + path, + &body(vec![dev.clone(), env.clone(), ask.clone()]), + &first, + ); + d.turn( + path, + &body(vec![ + dev.clone(), + env.clone(), + ask.clone(), + rs1.clone(), + fc1.clone(), + out1.clone(), + ]), + &second, + ); + d.turn( + path, + &body(vec![dev, env, ask, rs1, fc1, out1, rs2, ws, cc, out2]), + whole.to_string().as_bytes(), + ); + + let t = d.transcript(); + assert_eq!( + t.system.as_deref(), + Some("You are Codex.\n\nsandbox: workspace-write") + ); + // 用户消息不并：环境说明和要求是两条 + assert_eq!( + t.turns[0].input, + [ + msg( + R::User, + vec![text("cwd")] + ), + msg(R::User, vec![text("修好构建")]) + ] + ); + assert_eq!( + t.turns[0].output, + [thinking("看看报错"), call("c1", "shell", args)] + ); + assert_eq!( + t.turns[1].input, + [msg( + R::Tool, + vec![result("c1", "error[E0308]: mismatched types")] + )] + ); + // 只有密文的推理是一块空的推理 + assert_eq!( + t.turns[1].output, + [ + thinking(""), + P::Other { + label: "web_search_call".into() + }, + call("c2", "apply_patch", patch) + ] + ); + // 推理、搜索、调用三项是同一次回答，并成一条，作为上一轮的回答去掉 + assert_eq!( + t.turns[2].input, + [msg(R::Tool, vec![result("c2", "Done!")])] + ); + assert_eq!(t.turns[2].output, [text("构建好了。")]); + assert!(t.turns.iter().all(|x| !x.restart && x.gaps.is_empty())); + } + + /// Gemini：`functionCall` / `functionResponse`，没有调用号的按函数名对上；推理是带 + /// `thought` 的文字 + #[test] + fn a_gemini_session_with_function_calls() { + let mut d = Disk::new(); + let path = "/v1beta/models/gemini-2.5-pro:streamGenerateContent"; + let sys = json!({"parts": [{"text": "You are a coding agent"}]}); + let ask = json!({"role": "user", "parts": [{"text": "跑一下测试"}]}); + let said = json!({"role": "model", "parts": [{"text": "好，我来跑"}, + {"functionCall": {"name": "run_shell_command", "args": {"command": "npm test"}}}]}); + let ran = json!({"role": "user", "parts": [{"functionResponse": {"name": "run_shell_command", + "response": {"output": "3 passed"}}}]}); + let chunk = |parts: Value| { + ( + "", + json!({"candidates": [{"content": {"role": "model", "parts": parts}}]}), + ) + }; + d.turn( + path, + &json!({"systemInstruction": sys, "contents": [ask.clone()]}), + &sse(&[ + chunk(json!([{"text": "好，"}])), + chunk(json!([{"text": "我来跑"}, + {"functionCall": {"name": "run_shell_command", "args": {"command": "npm test"}}}])), + ]), + ); + d.turn( + path, + &json!({"system_instruction": sys, "contents": [ask, said, ran]}), + &sse(&[ + chunk(json!([{"text": "测试都过了", "thought": true}])), + chunk(json!([{"text": "全部通过。"}])), + ]), + ); + + let t = d.transcript(); + assert_eq!(t.system.as_deref(), Some("You are a coding agent")); + assert_eq!(t.turns[0].input, [msg(R::User, vec![text("跑一下测试")])]); + // 上游没给调用号：回答里读的时候现编了一个，下一轮里客户端记下的是按函数名认的那个，换成它 + assert_eq!( + t.turns[0].output, + [ + text("好，我来跑"), + call( + "run_shell_command", + "run_shell_command", + r#"{"command":"npm test"}"# + ) + ] + ); + assert_eq!( + t.turns[1].input, + [msg(R::Tool, vec![result("run_shell_command", "3 passed")])] + ); + assert_eq!( + t.turns[1].output, + [thinking("测试都过了"), text("全部通过。")] + ); + assert!(t.turns.iter().all(|x| !x.restart && x.gaps.is_empty())); + } + + // ───────────────────────────────────────────────── 前后两个请求怎么比 + + /// 压缩过、改过历史：对不上了，交出整段历史，标 `restart` + #[test] + fn a_compacted_or_edited_history_restarts_with_the_whole_history() { + let mut d = Disk::new(); + d.turn("/v1/messages", &anthropic(&[user("一")]), &says("回一")); + d.turn( + "/v1/messages", + &anthropic(&[user("一"), assistant("回一"), user("二")]), + &says("回二"), + ); + // 压缩：前文换成了一段摘要 + d.turn( + "/v1/messages", + &anthropic(&[user("前文摘要：说过一和二"), user("三")]), + &says("回三"), + ); + // 接着压缩之后的那段往下说 + d.turn( + "/v1/messages", + &anthropic(&[ + user("前文摘要：说过一和二"), + user("三"), + assistant("回三"), + user("四"), + ]), + &says("回四"), + ); + // 改过历史：早先的一条工具结果被清掉了（Claude Code 的 microcompact 就这么做） + d.turn( + "/v1/messages", + &anthropic(&[ + user("前文摘要：说过一和二"), + user("三（已清空）"), + assistant("回三"), + user("四"), + assistant("回四"), + user("五"), + ]), + &says("回五"), + ); + + let t = d.transcript(); + let restart: Vec = t.turns.iter().map(|x| x.restart).collect(); + assert_eq!(restart, [false, false, true, false, true]); + assert_eq!(t.turns[1].input, [msg(R::User, vec![text("二")])]); + // 连着的两条用户消息不并 + assert_eq!( + t.turns[2].input, + [ + msg(R::User, vec![text("前文摘要：说过一和二")]), + msg(R::User, vec![text("三")]) + ] + ); + assert_eq!(t.turns[3].input, [msg(R::User, vec![text("四")])]); + // 整段历史，连同里面的助手消息 + assert_eq!(t.turns[4].input.len(), 6); + assert_eq!(t.turns[4].input[2], msg(R::Assistant, vec![text("回三")])); + } + + /// 系统提示变了：那一轮说出新的那一份；`system` 是第一轮的 + #[test] + fn a_changed_system_prompt_is_said_on_the_turn_it_changed() { + let mut d = Disk::new(); + let with = |system: &str, messages: &[Value]| json!({"model": "m", "system": [{"type": "text", "text": system}], "messages": messages}); + d.turn( + "/v1/messages", + &with("计划模式", &[user("一")]), + &says("回一"), + ); + d.turn( + "/v1/messages", + &with("执行模式", &[user("一"), assistant("回一"), user("开始改")]), + &says("回二"), + ); + d.turn( + "/v1/messages", + &with( + "执行模式", + &[ + user("一"), + assistant("回一"), + user("开始改"), + assistant("回二"), + user("三"), + ], + ), + &says("回三"), + ); + let t = d.transcript(); + assert_eq!(t.system.as_deref(), Some("计划模式")); + let changed: Vec> = t + .turns + .iter() + .map(|x| x.system_changed.as_deref()) + .collect(); + assert_eq!(changed, [None, Some("执行模式"), None]); + // 系统提示变了不等于历史断了 + assert!(t.turns.iter().all(|x| !x.restart)); + assert_eq!(t.turns[1].input, [msg(R::User, vec![text("开始改")])]); + } + + /// 数 token 的请求也是会话里的一轮，但不是对话里的一句：空的，不打断前后的比对 + #[test] + fn a_count_tokens_call_is_an_empty_turn_that_does_not_break_the_chain() { + let mut d = Disk::new(); + d.turn("/v1/messages", &anthropic(&[user("一")]), &says("回一")); + let history = [user("一"), assistant("回一"), user("二")]; + d.turn( + "/v1/messages/count_tokens", + &json!({"model": "m", "messages": history}), + br#"{"input_tokens": 12}"#, + ); + d.turn("/v1/messages", &anthropic(&history), &says("回二")); + let t = d.transcript(); + assert_eq!(t.turns.len(), 3); + let side = &t.turns[1]; + assert!( + side.input.is_empty() + && side.output.is_empty() + && side.gaps.is_empty() + && !side.restart + ); + assert!(!t.turns[2].restart); + assert_eq!(t.turns[2].input, [msg(R::User, vec![text("二")])]); + assert_eq!(t.turns[2].output, [text("回二")]); + } + + /// 失败了的请求没有回答也不算缺；客户端照原样重发，没有新的东西 + #[test] + fn a_failed_turn_has_no_answer_and_its_retry_adds_nothing() { + let mut d = Disk::new(); + let history = [user("一")]; + d.put( + "/v1/messages", + Some(anthropic(&history).to_string().as_bytes()), + Some(br#"{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}"#), + |r| { + r.status = Some(529); + r.error = Some(crate::db::tests::upstream_failed("overloaded")); + }, + ); + d.turn("/v1/messages", &anthropic(&history), &says("回一")); + let t = d.transcript(); + assert!(t.turns[0].output.is_empty() && t.turns[0].gaps.is_empty()); + assert_eq!(t.turns[0].input, [msg(R::User, vec![text("一")])]); + assert!(!t.turns[1].restart); + assert!(t.turns[1].input.is_empty(), "{:?}", t.turns[1].input); + assert_eq!(t.turns[1].output, [text("回一")]); + } + + // ───────────────────────────────────────────────── 读不出来的地方 + + /// 请求体没存下、只存了开头：那一轮没有 `input`，下一个读得懂的从头交出整段历史 + #[test] + fn missing_and_truncated_requests_are_gaps_and_the_next_turn_restarts() { + let mut d = Disk::new(); + d.turn("/v1/messages", &anthropic(&[user("一")]), &says("回一")); + d.put("/v1/messages", None, Some(says("回二").as_slice()), |_| {}); + let h3 = [ + user("一"), + assistant("回一"), + user("二"), + assistant("回二"), + user("三"), + ]; + d.turn("/v1/messages", &anthropic(&h3), &says("回三")); + // 只存了开头：存下来的那一段解析不了 + let mut h4 = h3.to_vec(); + h4.extend([assistant("回三"), user("四")]); + let full = anthropic(&h4).to_string(); + let id = d.put( + "/v1/messages", + Some(&full.as_bytes()[..full.len() / 2]), + Some(says("回四").as_slice()), + |_| {}, + ); + d.cut(id, Which::Request, full.len()); + let mut h5 = h4.clone(); + h5.extend([assistant("回四"), user("五")]); + d.turn("/v1/messages", &anthropic(&h5), &says("回五")); + + let t = d.transcript(); + let gaps: Vec<&[Gap]> = t.turns.iter().map(|x| x.gaps.as_slice()).collect(); + assert_eq!( + gaps, + [ + &[][..], + &[Gap::RequestMissing], + &[], + &[Gap::RequestTruncated], + &[] + ] + ); + let restart: Vec = t.turns.iter().map(|x| x.restart).collect(); + assert_eq!(restart, [false, false, true, false, true]); + // 读不懂的那一轮回答照样读 + assert!(t.turns[1].input.is_empty()); + assert_eq!(t.turns[1].output, [text("回二")]); + assert_eq!(t.turns[2].input.len(), 5); + assert_eq!(t.turns[3].output, [text("回四")]); + assert_eq!(t.turns[4].input.len(), 9); + } + + /// 回答没存下、只存了开头、读不懂：各是一种缺口。没交全的回答，下一轮里客户端记下的那条 + /// 助手消息留着 —— 它是那一轮说过什么的唯一记录 + #[test] + fn missing_truncated_and_unreadable_answers_are_gaps() { + let mut d = Disk::new(); + let mut h = vec![user("一")]; + // 回答没存下：上游回了 200，收到了字节 + d.put( + "/v1/messages", + Some(anthropic(&h).to_string().as_bytes()), + None, + |_| {}, + ); + h.extend([assistant("回一"), user("二")]); + // 只存了开头：读得出第一块，第二块断在半路 + let long = anthropic_stream(&[ + json!({"type": "text", "text": "前一半"}), + json!({"type": "text", "text": "后一半"}), + ]); + let cut = String::from_utf8(long.clone()).unwrap(); + let keep = cut.find("后").unwrap(); + let id = d.put( + "/v1/messages", + Some(anthropic(&h).to_string().as_bytes()), + Some(&long[..keep]), + |_| {}, + ); + d.cut(id, Which::Response, long.len()); + h.extend([assistant("前一半后一半"), user("三")]); + // 读不懂：上游回了 200，正文却是一页 HTML + d.put( + "/v1/messages", + Some(anthropic(&h).to_string().as_bytes()), + Some(b"502 Bad Gateway"), + |_| {}, + ); + h.extend([assistant("回三"), user("四")]); + // 客户端在第一个字节之前就走了：没有回答，也不算缺 + d.put( + "/v1/messages", + Some(anthropic(&h).to_string().as_bytes()), + None, + |r| { + r.cancelled = true; + r.bytes = Some(0); + }, + ); + + let t = d.transcript(); + let gaps: Vec<&[Gap]> = t.turns.iter().map(|x| x.gaps.as_slice()).collect(); + assert_eq!( + gaps, + [ + &[Gap::ResponseMissing][..], + &[Gap::ResponseTruncated], + &[Gap::ResponseUnreadable], + &[] + ] + ); + assert!(t.turns[0].output.is_empty()); + assert_eq!(t.turns[1].output, [text("前一半")]); + assert!(t.turns[2].output.is_empty() && t.turns[3].output.is_empty()); + // 上一轮的回答没交全：客户端记下的那条助手消息留在新的这几条里 + assert!(t.turns.iter().all(|x| !x.restart)); + assert_eq!( + t.turns[1].input, + [ + msg(R::Assistant, vec![text("回一")]), + msg(R::User, vec![text("二")]) + ] + ); + assert_eq!( + t.turns[2].input, + [ + msg(R::Assistant, vec![text("前一半后一半")]), + msg(R::User, vec![text("三")]) + ] + ); + assert_eq!(t.turns[3].input[0].role, R::Assistant); + } + + // ───────────────────────────────────────────────── 打码 + + /// 密钥在工具结果里、在一行的开头（JSON 里它前面是 `\n` 的转义），交出去之前都打了码； + /// 系统提示和回答里的也一样 + #[test] + fn secrets_come_out_masked_wherever_they_are() { + let mut d = Disk::new(); + let key = "sk-ant-api03-SECRETSECRETSECRET"; + let body = json!({"model": "m", "system": format!("备用的 key：{key}"), "messages": [ + user("读一下 .env"), + json!({"role": "assistant", "content": [{"type": "tool_use", "id": "t1", "name": "Read", + "input": {"file_path": ".env", "note": key}}]}), + json!({"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", + "content": format!("# 配置\n{key}\nGITHUB=ghp_bbbbbbbbbbbbbbbbbbbb")}]}) + ]}); + d.turn("/v1/messages", &body, &says(&format!("找到了 {key}"))); + let t = d.transcript(); + let wire = serde_json::to_string(&t).unwrap(); + assert!(!wire.contains("SECRETSECRET"), "{wire}"); + assert!(!wire.contains("bbbbbbbbbbbbbbbb"), "{wire}"); + assert!(wire.contains("sk-an…CRET"), "{wire}"); + let P::ToolResult { text: said, .. } = &t.turns[0].input[2].parts[0] else { + panic!("{:?}", t.turns[0].input) + }; + assert_eq!(said, "# 配置\nsk-an…CRET\nGITHUB=ghp_b…bbbb"); + } + + // ───────────────────────────────────────────────── 转换过格式的 + + /// 客户端说 Anthropic、上游说 Chat 或 Responses：请求按客户端的格式读，回答按上游的 + #[test] + fn a_translated_turn_reads_the_answer_in_the_upstream_format() { + let mut d = Disk::new(); + let translated = |to: &str| { + Some( + json!({"provider": "relay", "from": "anthropic", "to": to, "dropped": []}) + .to_string(), + ) + }; + let chat = sse(&[ + ( + "", + json!({"choices": [{"index": 0, "delta": {"content": "我来查"}}]}), + ), + ( + "", + json!({"choices": [{"index": 0, "delta": {"tool_calls": [{"index": 0, + "type": "function", "function": {"name": "Grep", "arguments": "{\"pattern\":\"TODO\"}"}}]}}]}), + ), + ( + "", + json!({"choices": [{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]}), + ), + ]); + d.put( + "/v1/messages", + Some(anthropic(&[user("找 TODO")]).to_string().as_bytes()), + Some(&chat), + |r| r.translated = translated("openai-chat"), + ); + // 下一轮里客户端记下的调用号，是网关转换时编的那个 + let said = json!({"role": "assistant", "content": [{"type": "text", "text": "我来查"}, + {"type": "tool_use", "id": "call_gateway", "name": "Grep", "input": {"pattern": "TODO"}}]}); + let found = json!({"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_gateway", + "content": "src/a.rs:1: TODO"}]}); + let responses = sse(&[ + ( + "response.output_item.added", + json!({"type": "response.output_item.added", "output_index": 0, + "item": {"type": "message", "role": "assistant", "content": []}}), + ), + ( + "response.output_text.delta", + json!({"type": "response.output_text.delta", "output_index": 0, "content_index": 0, + "delta": "只有一处。"}), + ), + ( + "response.completed", + json!({"type": "response.completed", "response": {"usage": {"input_tokens": 1, "output_tokens": 1}}}), + ), + ]); + d.put( + "/v1/messages", + Some( + anthropic(&[user("找 TODO"), said, found]) + .to_string() + .as_bytes(), + ), + Some(&responses), + |r| r.translated = translated("openai-responses"), + ); + + let t = d.transcript(); + assert!(t.turns.iter().all(|x| x.gaps.is_empty() && !x.restart)); + assert_eq!( + t.turns[0].output, + [ + text("我来查"), + call("call_gateway", "Grep", r#"{"pattern":"TODO"}"#) + ] + ); + assert_eq!( + t.turns[1].input, + [msg( + R::Tool, + vec![result("call_gateway", "src/a.rs:1: TODO")] + )] + ); + assert_eq!(t.turns[1].output, [text("只有一处。")]); + } + + /// Codex 走 Anthropic 的上游：自由格式的 `apply_patch` 被包成了 `{"input": …}`，拆回原文 + #[test] + fn a_freeform_call_answered_by_another_format_is_unwrapped() { + let mut d = Disk::new(); + let patch = "*** Begin Patch\n*** End Patch"; + d.put( + "/v1/responses", + Some( + json!({"model": "claude-sonnet-4-5", "input": "改一下", + "tools": [{"type": "custom", "name": "apply_patch"}]}) + .to_string() + .as_bytes(), + ), + Some(&anthropic_stream(&[json!({"type": "tool_use", "id": "toolu_1", + "name": "apply_patch", "input": {"input": patch}})])), + |r| { + r.translated = Some( + json!({"provider": "anthropic", "from": "openai-responses", "to": "anthropic", "dropped": []}) + .to_string(), + ) + }, + ); + let t = d.transcript(); + assert_eq!(t.turns[0].input, [msg(R::User, vec![text("改一下")])]); + assert_eq!(t.turns[0].output, [call("toolu_1", "apply_patch", patch)]); + } + + /// Bedrock 的二进制帧在网关进门时转成了 SSE；整包的 Converse 也读得出来 + #[test] + fn a_bedrock_answer_is_read_streamed_or_whole() { + let mut d = Disk::new(); + let bedrock = || { + Some( + json!({"provider": "bedrock", "from": "anthropic", "to": "bedrock", "dropped": []}) + .to_string(), + ) + }; + let stream = sse(&[ + ("messageStart", json!({"role": "assistant"})), + ( + "contentBlockDelta", + json!({"contentBlockIndex": 0, "delta": {"reasoningContent": {"text": "想"}}}), + ), + ( + "contentBlockDelta", + json!({"contentBlockIndex": 0, "delta": {"reasoningContent": {"signature": "sig"}}}), + ), + ("contentBlockStop", json!({"contentBlockIndex": 0})), + ( + "contentBlockDelta", + json!({"contentBlockIndex": 1, "delta": {"text": "基岩"}}), + ), + ("contentBlockStop", json!({"contentBlockIndex": 1})), + ( + "contentBlockStart", + json!({"contentBlockIndex": 2, "start": {"toolUse": {"toolUseId": "t1", "name": "Bash"}}}), + ), + ( + "contentBlockDelta", + json!({"contentBlockIndex": 2, "delta": {"toolUse": {"input": "{\"command\":\"ls\"}"}}}), + ), + ("contentBlockStop", json!({"contentBlockIndex": 2})), + ("messageStop", json!({"stopReason": "tool_use"})), + ]); + d.put( + "/v1/messages", + Some(anthropic(&[user("一")]).to_string().as_bytes()), + Some(&stream), + |r| r.translated = bedrock(), + ); + let whole = json!({"output": {"message": {"role": "assistant", "content": [{"text": "整包"}]}}, + "stopReason": "end_turn"}); + d.put( + "/v1/messages", + Some(anthropic(&[user("另一段")]).to_string().as_bytes()), + Some(whole.to_string().as_bytes()), + |r| r.translated = bedrock(), + ); + let t = d.transcript(); + assert_eq!( + t.turns[0].output, + [ + thinking("想"), + text("基岩"), + call("t1", "Bash", r#"{"command":"ls"}"#) + ] + ); + assert_eq!(t.turns[1].output, [text("整包")]); + assert!(t.turns.iter().all(|x| x.gaps.is_empty())); + } + + /// 不带 `alt=sse` 的 Gemini 流是一个 JSON 数组；截断了的读到最后一个完整的元素 + #[test] + fn a_gemini_answer_written_as_a_json_array_is_read_element_by_element() { + let mut d = Disk::new(); + let path = "/v1beta/models/gemini-2.5-pro:streamGenerateContent"; + let body = json!({"contents": [{"role": "user", "parts": [{"text": "画一只猫"}]}]}); + let array = json!([ + {"candidates": [{"content": {"role": "model", "parts": [{"text": "画好了"}]}}]}, + {"candidates": [{"content": {"role": "model", "parts": [ + {"inlineData": {"mimeType": "image/png", "data": "QUJD".repeat(4)}}, + {"executableCode": {"language": "PYTHON", "code": "print(1)"}}]}}]}, + {"candidates": [{"content": {"role": "model", "parts": [{"text": "。"}]}}]} + ]) + .to_string(); + d.turn(path, &body, array.as_bytes()); + let id = d.put( + path, + Some(body.to_string().as_bytes()), + Some(&array.as_bytes()[..array.len() - 30]), + |_| {}, + ); + d.cut(id, Which::Response, array.len()); + + let t = d.transcript(); + let image = P::Image { + media_type: Some("image/png".into()), + bytes: Some(12), + }; + let code = P::Other { + label: "executableCode".into(), + }; + assert_eq!( + t.turns[0].output, + [text("画好了"), image.clone(), code.clone(), text("。")] + ); + assert!(t.turns[0].gaps.is_empty()); + assert_eq!(t.turns[1].output, [text("画好了"), image, code]); + assert_eq!(t.turns[1].gaps, [Gap::ResponseTruncated]); + } + + /// 和会话详情同样的请求、同样的顺序：本地应答的不在，同一刻的按请求号 + #[test] + fn the_turns_are_the_same_requests_in_the_same_order_as_the_session_detail() { + let mut d = Disk::new(); + d.turn("/v1/messages", &anthropic(&[user("一")]), &says("回一")); + d.put("/v1/messages", None, None, |r| { + r.local = true; + }); + d.put("/v1/messages", None, None, |r| r.at_ms = NOW); + d.put("/v1/messages", None, None, |r| r.at_ms = NOW); + let t = d.transcript(); + let want: Vec = + d.db.turns("s") + .unwrap() + .iter() + .map(|r| r.id.to_string()) + .collect(); + assert_eq!(want, ["3", "4", "1"]); + assert_eq!( + t.turns.iter().map(|x| x.id.clone()).collect::>(), + want + ); + } + + /// 别家没有对应物的块不丢：文件、音频、服务端工具的调用和结果、压缩过的前文，按类型名 + /// 记成 `other`，位置不变。图片只说类型和大小，给的是地址的不知道大小 + #[test] + fn blocks_without_a_counterpart_are_named_in_place() { + let mut d = Disk::new(); + let jpeg = format!("data:image/jpeg;base64,{}", "QUJD".repeat(5)); + d.turn( + "/v1/chat/completions", + &json!({"model": "m", "messages": [{"role": "user", "content": [ + {"type": "text", "text": "看图听音"}, + {"type": "image_url", "image_url": {"url": jpeg}}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + {"type": "input_audio", "input_audio": {"data": "UklGRg==", "format": "wav"}}, + {"type": "file", "file": {"file_data": "data:application/pdf;base64,JVBERi0="}} + ]}]}), + br#"{"choices":[{"message":{"role":"assistant","content":"ok"}}]}"#, + ); + let t = d.transcript(); + assert_eq!( + t.turns[0].input, + [msg( + R::User, + vec![ + text("看图听音"), + P::Image { + media_type: Some("image/jpeg".into()), + bytes: Some(15) + }, + P::Image { + media_type: None, + bytes: None + }, + P::Other { + label: "input_audio".into() + }, + P::Other { + label: "file".into() + } + ] + )] + ); + let wire = serde_json::to_string(&t).unwrap(); + assert!(!wire.contains("QUJD") && !wire.contains("UklGRg") && !wire.contains("JVBERi0")); + + // Anthropic：网页搜索是服务端工具，调用和结果都是助手消息里的块；回答的流里也一样 + let mut d = Disk::new(); + let search = json!({"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", + "input": {"query": "rust 2024"}}); + let found = json!({"type": "web_search_tool_result", "tool_use_id": "srvtoolu_1", + "content": [{"type": "web_search_result", "url": "https://x", "title": "x", "encrypted_content": "e"}]}); + d.turn( + "/v1/messages", + &anthropic(&[json!({"role": "user", "content": [ + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "合同全文"}}, + {"type": "text", "text": "查一下"} + ]})]), + &anthropic_stream(&[ + search.clone(), + found.clone(), + json!({"type": "text", "text": "查到了"}), + ]), + ); + let t = d.transcript(); + assert_eq!( + t.turns[0].input, + [msg( + R::User, + vec![ + P::Other { + label: "document".into() + }, + text("查一下") + ] + )] + ); + assert_eq!( + t.turns[0].output, + [ + P::Other { + label: "server_tool_use".into() + }, + P::Other { + label: "web_search_tool_result".into() + }, + text("查到了") + ] + ); + + // Responses：压缩过的前文像一条系统消息；引用服务端内容的项是 `other` + let mut d = Disk::new(); + d.turn( + "/v1/responses", + &json!({"model": "m", "input": [ + {"type": "compaction", "encrypted_content": "gAAAA"}, + {"type": "item_reference", "id": "msg_1"}, + {"role": "user", "content": "接着来"} + ]}), + br#"{"output":[]}"#, + ); + let t = d.transcript(); + assert_eq!( + t.turns[0].input, + [ + msg( + R::System, + vec![P::Other { + label: "compaction".into() + }] + ), + msg( + R::User, + vec![P::Other { + label: "item_reference".into() + }] + ), + msg(R::User, vec![text("接着来")]) + ] + ); + assert!(t.turns[0].output.is_empty() && t.turns[0].gaps.is_empty()); + } + + // ───────────────────────────────────────────────── 量一量 + + /// 读一次几百轮的会话要多久：照 Claude Code 的样子造一段 300 轮的对话，每一轮的请求带着 + /// 整段历史（平均 1 MB 上下，工具结果每条 2.5 KB 上下），流式的回答里有推理、文字和工具调用。 + /// + /// 默认不跑：要写三百多 MB 的正文。发布构建下跑： + /// `cargo test --release -p tw-store transcript_cost -- --ignored --nocapture` + #[test] + #[ignore] + fn transcript_cost() { + let mut d = Disk::new(); + let system = "You are Claude Code. ".repeat(1000); + let tools: Vec = (0..40) + .map(|i| { + json!({"name": format!("Tool{i}"), "description": "x".repeat(600), + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}}}) + }) + .collect(); + let file = "fn main() { println!(\"hello\"); }\n".repeat(75); + let pair = |k: usize| { + [ + json!({"role": "assistant", "content": [ + {"type": "thinking", "thinking": format!("第 {k} 步先读一下"), "signature": "s".repeat(300)}, + {"type": "text", "text": format!("第 {k} 步我来读 src/{k}.rs")}, + {"type": "tool_use", "id": format!("toolu_{k}"), "name": "Read", + "input": {"file_path": format!("src/{k}.rs")}}]}), + json!({"role": "user", "content": [{"type": "tool_result", "tool_use_id": format!("toolu_{k}"), + "content": format!("{file}// 第 {k} 个文件")}]}), + ] + }; + // 开头已经聊了两百轮 + let mut messages = vec![json!({"role": "user", "content": "把这个仓库的 bug 都修了"})]; + for k in 0..200 { + messages.extend(pair(k)); + } + let mut bytes = 0usize; + for turn in 0..300 { + let mut sent = messages.clone(); + if let Some(last) = sent.last_mut() { + *last = with_cache(last.clone()); + } + let body = json!({"model": "claude-sonnet-4-5", "system": system, "tools": tools, + "max_tokens": 32000, "stream": true, "messages": sent}) + .to_string(); + let k = 200 + turn; + let answer = anthropic_stream(&[ + json!({"type": "thinking", "thinking": format!("第 {k} 步先读一下"), "signature": "s".repeat(300)}), + json!({"type": "text", "text": format!("第 {k} 步我来读 src/{k}.rs")}), + json!({"type": "tool_use", "id": format!("toolu_{k}"), "name": "Read", + "input": {"file_path": format!("src/{k}.rs")}}), + ]); + bytes += body.len() + answer.len(); + d.put("/v1/messages", Some(body.as_bytes()), Some(&answer), |_| {}); + messages.extend(pair(k)); + } + let rows = d.db.session_requests("s").unwrap(); + eprintln!( + "{} requests, {} MB of bodies ({} KB per request on average)", + rows.len(), + bytes / 1_000_000, + bytes / rows.len() / 1000 + ); + for _ in 0..3 { + let t0 = std::time::Instant::now(); + let t = build("s", &rows, &d.blobs); + let took = t0.elapsed(); + let wire = serde_json::to_string(&t).unwrap(); + eprintln!( + "built in {took:?} ({:?} per request), {} KB of JSON", + took / rows.len() as u32, + wire.len() / 1000 + ); + assert!( + t.turns + .iter() + .skip(1) + .all(|x| !x.restart && x.input.len() == 1) + ); + assert!( + t.turns + .iter() + .all(|x| x.gaps.is_empty() && x.output.len() == 3) + ); + } + } +} diff --git a/crates/tw-store/src/transcript/read.rs b/crates/tw-store/src/transcript/read.rs new file mode 100644 index 00000000..5f88faf4 --- /dev/null +++ b/crates/tw-store/src/transcript/read.rs @@ -0,0 +1,786 @@ +//! 四种客户端格式里的消息怎么读成对话里的一块一块。 +//! +//! **不经过 tw-dialect 的中间表示。**中间表示只装四种格式之间对应得上的东西：对话中途的 +//! 系统消息并进了系统提示，只装着工具结果的消息和用户说的话是同一个角色，服务端工具的 +//! 调用和结果、文件、音频解码时直接丢掉 —— 转换格式时那样是对的，可对话记录要的恰恰是 +//! 这些：谁在哪一处说了什么，哪里有一块转不过去的东西。所以这里照客户端自己的格式读 +//! JSON，只读消息这一层；取哪个字段和 tw-dialect 的解码器保持一致。 +//! +//! 整包的回答和请求里的消息是同一个形状（Anthropic 的 `content`、Chat 的 `message`、 +//! Responses 的 `output`、Gemini 的 `content`），也按这里读。 + +use std::borrow::Cow; + +use serde_json::Value; +use tw_dialect::ir::Dialect; + +/// 一条消息是谁说的。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum Role { + User, + Assistant, + /// 只装着工具结果 + Tool, + /// 对话中途的 system、developer 消息 + System, +} + +/// 读出来的一块，借着解析好的 JSON。 +/// +/// **已经是规整过的样子**：缓存断点和推理签名不在里面，键的顺序也无所谓了 —— 前后两个 +/// 请求比对时，这些本来就不该算（Claude Code 每一轮都把缓存断点挪到最后一条消息上）。 +#[derive(Debug, Clone)] +pub(super) enum Piece<'a> { + Text(Cow<'a, str>), + /// 推理。只有签名、被打码的，是空串 + Thinking(Cow<'a, str>), + ToolCall { + id: Cow<'a, str>, + name: Cow<'a, str>, + input: Input<'a>, + }, + ToolResult { + call_id: Cow<'a, str>, + text: Cow<'a, str>, + is_error: bool, + }, + /// `data` 只进指纹，**从不交出去** + Image { + media_type: Option<&'a str>, + bytes: Option, + data: Option<&'a str>, + }, + /// 别的块，类型名原样 + Other(&'a str), +} + +/// 工具调用的参数。 +#[derive(Debug, Clone)] +pub(super) enum Input<'a> { + /// 函数工具的参数，模型写出来的 JSON 文本。空的当 `{}` + Args(&'a str), + /// 自由格式工具的原文（Codex 的 `apply_patch`），一个字不改 + Raw(&'a str), + /// 已经是 JSON 的参数（Anthropic、Gemini） + Json(&'a Value), +} + +impl Input<'_> { + /// 交出去的样子：JSON 文本，自由格式的是原文 + pub(super) fn text(&self) -> String { + match self { + Input::Args(s) if s.trim().is_empty() => "{}".into(), + Input::Args(s) | Input::Raw(s) => s.to_string(), + Input::Json(v) => v.to_string(), + } + } +} + +/// 一条消息（Responses 里是一个输入项）读出来的样子。 +#[derive(Debug, Clone)] +pub(super) struct Item<'a> { + pub(super) role: Role, + pub(super) pieces: Vec>, +} + +/// 一个请求体里对话记录要的东西。 +#[derive(Debug, Default)] +pub(super) struct Body<'a> { + /// 系统提示，按出现的先后 + pub(super) system: Vec>, + pub(super) items: Vec>, + /// 定义成自由格式的工具。别的格式的上游会把它们的原文包进 `{"input": …}` + pub(super) freeform: Vec, +} + +/// 按客户端的格式读一个请求体。不是对象、没有对话的读出来是空的，不算读不懂。 +pub(super) fn body(dialect: Dialect, v: &Value) -> Body<'_> { + match dialect { + Dialect::Anthropic => anthropic(v), + Dialect::Chat => chat(v), + Dialect::Responses => responses(v), + Dialect::Gemini => gemini(v), + // 客户端不说 Converse（见 `tw_api::Dialect::Bedrock`） + Dialect::Bedrock => Body::default(), + } +} + +// ───────────────────────────────────────────────────────── 小工具 + +fn str_of<'a>(v: &'a Value, key: &str) -> Option<&'a str> { + v.get(key).and_then(Value::as_str) +} + +fn arr_of<'a>(v: &'a Value, key: &str) -> &'a [Value] { + v.get(key) + .and_then(Value::as_array) + .map(Vec::as_slice) + .unwrap_or(&[]) +} + +/// 一个 `string | [{type: "text", text}]` 形状的内容里的字，几段之间换行。 +fn text_of(v: &Value) -> Cow<'_, str> { + match v { + Value::String(s) => Cow::Borrowed(s), + Value::Array(parts) => { + let texts: Vec<&str> = parts.iter().filter_map(|p| str_of(p, "text")).collect(); + match texts.as_slice() { + [] => Cow::Borrowed(""), + [one] => Cow::Borrowed(one), + many => Cow::Owned(many.join("\n")), + } + } + _ => Cow::Borrowed(""), + } +} + +/// 块的类型名。没写的叫 `unknown` +fn kind_of(v: &Value) -> &str { + str_of(v, "type").unwrap_or("unknown") +} + +/// base64 解出来有多少字节 +fn base64_len(data: &str) -> u64 { + let d = data.trim_end(); + let pad = d.bytes().rev().take_while(|b| *b == b'=').count() as u64; + (d.len() as u64 * 3 / 4).saturating_sub(pad) +} + +/// 一张以 URI 给的图片：`data:` URI 说得出类型和大小，网址说不出 +fn image_uri(uri: &str) -> Piece<'_> { + if let Some(rest) = uri.strip_prefix("data:") + && let Some((head, data)) = rest.split_once(',') + { + let (mime, base64) = match head.strip_suffix(";base64") { + Some(m) => (m, true), + None => (head.split(';').next().unwrap_or_default(), false), + }; + return Piece::Image { + media_type: (!mime.is_empty()).then_some(mime), + bytes: base64.then(|| base64_len(data)), + data: Some(data), + }; + } + Piece::Image { + media_type: None, + bytes: None, + data: Some(uri), + } +} + +/// 文字，空的不算一块 +fn text_piece<'a>(t: Cow<'a, str>, out: &mut Vec>) { + if !t.is_empty() { + out.push(Piece::Text(t)); + } +} + +// ───────────────────────────────────────────────────────── Anthropic Messages + +fn anthropic(v: &Value) -> Body<'_> { + let mut system = Vec::new(); + match v.get("system") { + Some(Value::String(s)) => system.push(Cow::Borrowed(s.as_str())), + Some(Value::Array(blocks)) => system.extend( + blocks + .iter() + .filter_map(|b| str_of(b, "text")) + .map(Cow::Borrowed), + ), + _ => {} + } + Body { + system, + items: arr_of(v, "messages") + .iter() + .map(anthropic_message) + .collect(), + freeform: Vec::new(), + } +} + +fn anthropic_message(m: &Value) -> Item<'_> { + let mut pieces = Vec::new(); + let mut only_results = false; + match m.get("content") { + Some(Value::String(s)) => text_piece(Cow::Borrowed(s), &mut pieces), + Some(Value::Array(blocks)) => { + only_results = !blocks.is_empty() + && blocks + .iter() + .all(|b| str_of(b, "type") == Some("tool_result")); + anthropic_blocks(blocks, &mut pieces); + } + _ => {} + } + let role = match str_of(m, "role") { + Some("assistant") => Role::Assistant, + // DeepSeek Harness 在消息里放的 system 角色 + Some("system") => Role::System, + _ if only_results => Role::Tool, + _ => Role::User, + }; + Item { role, pieces } +} + +/// 一串内容块：消息的 `content`，也是整包回答的 `content`。 +pub(super) fn anthropic_blocks<'a>(blocks: &'a [Value], out: &mut Vec>) { + for b in blocks { + match kind_of(b) { + "text" => text_piece(Cow::Borrowed(str_of(b, "text").unwrap_or_default()), out), + "thinking" => out.push(Piece::Thinking(Cow::Borrowed( + str_of(b, "thinking").unwrap_or_default(), + ))), + "redacted_thinking" => out.push(Piece::Thinking(Cow::Borrowed(""))), + "tool_use" => out.push(Piece::ToolCall { + id: Cow::Borrowed(str_of(b, "id").unwrap_or_default()), + name: Cow::Borrowed(str_of(b, "name").unwrap_or_default()), + input: b.get("input").map_or(Input::Args(""), Input::Json), + }), + "tool_result" => anthropic_result(b, out), + "image" => out.push(anthropic_image(b)), + // document、search_result，服务端工具的 server_tool_use、web_search_tool_result…… + other => out.push(Piece::Other(other)), + } + } +} + +fn anthropic_image(b: &Value) -> Piece<'_> { + let src = b.get("source").unwrap_or(&Value::Null); + match str_of(src, "type") { + Some("base64") => { + let data = str_of(src, "data").unwrap_or_default(); + Piece::Image { + media_type: str_of(src, "media_type"), + bytes: Some(base64_len(data)), + data: Some(data), + } + } + Some("url") => Piece::Image { + media_type: None, + bytes: None, + data: str_of(src, "url"), + }, + // Files API 里的文件：只有一个文件号 + _ => Piece::Image { + media_type: None, + bytes: None, + data: str_of(src, "file_id"), + }, + } +} + +/// 工具结果：里面的字连起来成一块，图片和别的块跟在它后面。 +fn anthropic_result<'a>(b: &'a Value, out: &mut Vec>) { + let mut extra = Vec::new(); + let text = match b.get("content") { + Some(Value::String(s)) => Cow::Borrowed(s.as_str()), + Some(Value::Array(items)) => { + let mut texts = Vec::new(); + for i in items { + match kind_of(i) { + "text" => texts.push(str_of(i, "text").unwrap_or_default()), + "image" => extra.push(anthropic_image(i)), + other => extra.push(Piece::Other(other)), + } + } + match texts.as_slice() { + [one] => Cow::Borrowed(*one), + many => Cow::Owned(many.join("\n")), + } + } + _ => Cow::Borrowed(""), + }; + out.push(Piece::ToolResult { + call_id: Cow::Borrowed(str_of(b, "tool_use_id").unwrap_or_default()), + text, + is_error: b.get("is_error").and_then(Value::as_bool).unwrap_or(false), + }); + out.extend(extra); +} + +// ───────────────────────────────────────────────────────── OpenAI Chat Completions + +fn chat(v: &Value) -> Body<'_> { + let mut system = Vec::new(); + let mut items = Vec::new(); + // 开头连着的 system、developer 消息就是系统提示；对话中途的照原位置留着 + let mut leading = true; + for m in arr_of(v, "messages") { + let role = str_of(m, "role").unwrap_or_default(); + if leading && matches!(role, "system" | "developer") { + system.push(text_of(m.get("content").unwrap_or(&Value::Null))); + continue; + } + leading = false; + items.push(chat_message(m)); + } + Body { + system, + items, + freeform: arr_of(v, "tools") + .iter() + .filter(|t| str_of(t, "type") == Some("custom")) + .filter_map(|t| t.get("custom").and_then(|c| str_of(c, "name"))) + .map(str::to_string) + .collect(), + } +} + +fn chat_message(m: &Value) -> Item<'_> { + let content = m.get("content").unwrap_or(&Value::Null); + match str_of(m, "role").unwrap_or_default() { + "system" | "developer" => { + let mut pieces = Vec::new(); + text_piece(text_of(content), &mut pieces); + Item { + role: Role::System, + pieces, + } + } + "assistant" => Item { + role: Role::Assistant, + pieces: chat_assistant(m), + }, + "tool" => Item { + role: Role::Tool, + pieces: vec![Piece::ToolResult { + call_id: Cow::Borrowed(str_of(m, "tool_call_id").unwrap_or_default()), + text: text_of(content), + is_error: false, + }], + }, + // 旧的函数调用：结果按函数名认 + "function" => Item { + role: Role::Tool, + pieces: vec![Piece::ToolResult { + call_id: Cow::Borrowed(str_of(m, "name").unwrap_or_default()), + text: text_of(content), + is_error: false, + }], + }, + _ => { + let mut pieces = Vec::new(); + match content { + Value::String(s) => text_piece(Cow::Borrowed(s), &mut pieces), + Value::Array(parts) => { + for p in parts { + match kind_of(p) { + "text" => text_piece( + Cow::Borrowed(str_of(p, "text").unwrap_or_default()), + &mut pieces, + ), + "image_url" => { + let url = p + .get("image_url") + .and_then(|i| i.as_str().or_else(|| str_of(i, "url"))); + pieces.push(url.map_or( + Piece::Image { + media_type: None, + bytes: None, + data: None, + }, + image_uri, + )); + } + // input_audio、file…… + other => pieces.push(Piece::Other(other)), + } + } + } + _ => {} + } + Item { + role: Role::User, + pieces, + } + } + } +} + +/// 助手消息：请求里的历史，也是整包回答的 `choices[0].message`。 +pub(super) fn chat_assistant(m: &Value) -> Vec> { + let mut pieces = Vec::new(); + // DeepSeek、Kimi、GLM 这些的推理字段（OpenRouter 叫 `reasoning`） + if let Some(t) = str_of(m, "reasoning_content") + .or_else(|| str_of(m, "reasoning")) + .filter(|t| !t.is_empty()) + { + pieces.push(Piece::Thinking(Cow::Borrowed(t))); + } + match m.get("content") { + Some(Value::String(s)) => text_piece(Cow::Borrowed(s), &mut pieces), + Some(Value::Array(parts)) => { + for p in parts { + match str_of(p, "text").or_else(|| str_of(p, "refusal")) { + Some(t) => text_piece(Cow::Borrowed(t), &mut pieces), + None => pieces.push(Piece::Other(kind_of(p))), + } + } + } + _ => {} + } + if let Some(r) = str_of(m, "refusal") { + text_piece(Cow::Borrowed(r), &mut pieces); + } + for c in arr_of(m, "tool_calls") { + let id = Cow::Borrowed(str_of(c, "id").unwrap_or_default()); + match str_of(c, "type") { + Some("custom") => { + let x = c.get("custom").unwrap_or(&Value::Null); + pieces.push(Piece::ToolCall { + id, + name: Cow::Borrowed(str_of(x, "name").unwrap_or_default()), + input: Input::Raw(str_of(x, "input").unwrap_or_default()), + }); + } + _ => { + let f = c.get("function").unwrap_or(&Value::Null); + pieces.push(Piece::ToolCall { + id, + name: Cow::Borrowed(str_of(f, "name").unwrap_or_default()), + input: Input::Args(str_of(f, "arguments").unwrap_or_default()), + }); + } + } + } + // 旧的函数调用：没有调用号 + if let Some(f) = m.get("function_call") { + pieces.push(Piece::ToolCall { + id: Cow::Borrowed(""), + name: Cow::Borrowed(str_of(f, "name").unwrap_or_default()), + input: Input::Args(str_of(f, "arguments").unwrap_or_default()), + }); + } + if m.get("audio").is_some_and(|a| !a.is_null()) { + pieces.push(Piece::Other("audio")); + } + pieces +} + +// ───────────────────────────────────────────────────────── OpenAI Responses + +fn responses(v: &Value) -> Body<'_> { + let mut system = Vec::new(); + match v.get("instructions") { + Some(Value::String(s)) => system.push(Cow::Borrowed(s.as_str())), + Some(a @ Value::Array(_)) => system.push(text_of(a)), + _ => {} + } + let mut items = Vec::new(); + match v.get("input") { + Some(Value::String(s)) => { + let mut pieces = Vec::new(); + text_piece(Cow::Borrowed(s), &mut pieces); + items.push(Item { + role: Role::User, + pieces, + }); + } + Some(Value::Array(input)) => { + // 开头连着的 system、developer 消息算系统提示，和 Chat 一样 + let mut leading = true; + for it in input { + if leading + && str_of(it, "type").unwrap_or("message") == "message" + && matches!(str_of(it, "role"), Some("system" | "developer")) + { + system.push(text_of(it.get("content").unwrap_or(&Value::Null))); + continue; + } + leading = false; + items.push(responses_item(it)); + } + } + _ => {} + } + let mut freeform = Vec::new(); + for t in arr_of(v, "tools") { + match str_of(t, "type") { + Some("custom") => freeform.extend(str_of(t, "name").map(str::to_string)), + Some("namespace") => { + let ns = str_of(t, "name"); + freeform.extend( + arr_of(t, "tools") + .iter() + .filter(|x| str_of(x, "type") == Some("custom")) + .filter_map(|x| str_of(x, "name")) + .map(|n| flat_name(ns, n).into_owned()), + ); + } + _ => {} + } + } + Body { + system, + items, + freeform, + } +} + +/// namespace 里的工具展开成 `namespace__名字`，和转换给别家时的名字一样 +fn flat_name<'a>(namespace: Option<&str>, name: &'a str) -> Cow<'a, str> { + match namespace { + Some(ns) if !ns.is_empty() => Cow::Owned(format!("{ns}__{name}")), + _ => Cow::Borrowed(name), + } +} + +/// 一个输入项：请求里的历史，也是整包回答的 `output` 里的一项。 +pub(super) fn responses_item(it: &Value) -> Item<'_> { + let kind = str_of(it, "type").unwrap_or("message"); + let one = |role, piece| Item { + role, + pieces: vec![piece], + }; + match kind { + "message" => { + let role = match str_of(it, "role") { + Some("assistant") => Role::Assistant, + Some("system" | "developer") => Role::System, + _ => Role::User, + }; + let mut pieces = Vec::new(); + match it.get("content") { + Some(Value::String(s)) => text_piece(Cow::Borrowed(s), &mut pieces), + Some(Value::Array(parts)) => responses_parts(parts, &mut pieces), + _ => {} + } + Item { role, pieces } + } + "reasoning" => { + let texts = |key: &str| -> Vec<&str> { + arr_of(it, key) + .iter() + .filter_map(|x| str_of(x, "text")) + .collect() + }; + // 推理原文优先，没有就是摘要；几段之间空一行（和 tw-dialect 一样） + let mut t = texts("content"); + if t.is_empty() { + t = texts("summary"); + } + let text = match t.as_slice() { + [] => Cow::Borrowed(""), + [one] => Cow::Borrowed(*one), + many => Cow::Owned(many.join("\n\n")), + }; + one(Role::Assistant, Piece::Thinking(text)) + } + "function_call" | "custom_tool_call" => { + let name = flat_name( + str_of(it, "namespace"), + str_of(it, "name").unwrap_or_default(), + ); + let input = if kind == "function_call" { + Input::Args(str_of(it, "arguments").unwrap_or_default()) + } else { + Input::Raw(str_of(it, "input").unwrap_or_default()) + }; + one( + Role::Assistant, + Piece::ToolCall { + id: Cow::Borrowed(str_of(it, "call_id").unwrap_or_default()), + name, + input, + }, + ) + } + "function_call_output" | "custom_tool_call_output" | "local_shell_call_output" => { + let mut extra = Vec::new(); + let text = match it.get("output") { + Some(Value::String(s)) => Cow::Borrowed(s.as_str()), + Some(Value::Array(parts)) => { + let mut texts = Vec::new(); + for p in parts { + match kind_of(p) { + "input_text" | "output_text" => { + texts.push(str_of(p, "text").unwrap_or_default()) + } + "input_image" => extra.push(responses_image(p)), + other => extra.push(Piece::Other(other)), + } + } + match texts.as_slice() { + [one] => Cow::Borrowed(*one), + many => Cow::Owned(many.join("\n")), + } + } + _ => Cow::Borrowed(""), + }; + let mut pieces = vec![Piece::ToolResult { + call_id: Cow::Borrowed(str_of(it, "call_id").unwrap_or_default()), + text, + is_error: false, + }]; + pieces.extend(extra); + Item { + role: Role::Tool, + pieces, + } + } + // 压缩过的前文：只有 OpenAI 读得懂的一段密文，在对话里的位置像一条系统消息 + "compaction" => one(Role::System, Piece::Other(kind)), + // 别的工具结果（computer_call_output……） + k if k.ends_with("_output") => one(Role::Tool, Piece::Other(k)), + // 托管工具的调用（web_search_call、image_generation_call、mcp_call、local_shell_call……） + k if k.ends_with("_call") || matches!(k, "mcp_list_tools" | "mcp_approval_request") => { + one(Role::Assistant, Piece::Other(k)) + } + // item_reference、mcp_approval_response…… + other => one(Role::User, Piece::Other(other)), + } +} + +fn responses_parts<'a>(parts: &'a [Value], out: &mut Vec>) { + for p in parts { + match kind_of(p) { + "input_text" | "output_text" => { + text_piece(Cow::Borrowed(str_of(p, "text").unwrap_or_default()), out) + } + "refusal" => text_piece(Cow::Borrowed(str_of(p, "refusal").unwrap_or_default()), out), + "input_image" => out.push(responses_image(p)), + // input_file、input_audio…… + other => out.push(Piece::Other(other)), + } + } +} + +fn responses_image(p: &Value) -> Piece<'_> { + match str_of(p, "image_url") { + Some(u) => image_uri(u), + None => Piece::Image { + media_type: None, + bytes: None, + data: str_of(p, "file_id"), + }, + } +} + +// ───────────────────────────────────────────────────────── Gemini + +/// 按驼峰名取字段，取不到再试下划线写法：Gemini 的 REST 接口两种都收 +fn field<'a>(v: &'a Value, camel: &str) -> Option<&'a Value> { + v.get(camel).or_else(|| { + let mut snake = String::with_capacity(camel.len() + 4); + for c in camel.chars() { + if c.is_ascii_uppercase() { + snake.push('_'); + snake.push(c.to_ascii_lowercase()); + } else { + snake.push(c); + } + } + v.get(snake) + }) +} + +fn fstr<'a>(v: &'a Value, camel: &str) -> Option<&'a str> { + field(v, camel).and_then(Value::as_str) +} + +fn gemini(v: &Value) -> Body<'_> { + let mut system = Vec::new(); + match field(v, "systemInstruction") { + Some(Value::String(s)) => system.push(Cow::Borrowed(s.as_str())), + Some(sys) => system.push(text_of(sys.get("parts").unwrap_or(&Value::Null))), + None => {} + } + Body { + system, + items: arr_of(v, "contents").iter().map(gemini_content).collect(), + freeform: Vec::new(), + } +} + +fn gemini_content(c: &Value) -> Item<'_> { + let parts = arr_of(c, "parts"); + let mut pieces = Vec::new(); + gemini_parts(parts, &mut pieces); + let role = match str_of(c, "role") { + Some("model") => Role::Assistant, + Some("function") => Role::Tool, + Some("system") => Role::System, + _ if !parts.is_empty() && parts.iter().all(|p| field(p, "functionResponse").is_some()) => { + Role::Tool + } + _ => Role::User, + }; + Item { role, pieces } +} + +/// 一串 part：请求里的一条消息，也是回答里候选的 `content.parts`。 +/// +/// 没有调用号的函数调用和结果按函数名对上（Gemini 自己不要求调用号）。 +pub(super) fn gemini_parts<'a>(parts: &'a [Value], out: &mut Vec>) { + for p in parts { + if let Some(t) = fstr(p, "text") { + if p.get("thought").and_then(Value::as_bool) == Some(true) { + out.push(Piece::Thinking(Cow::Borrowed(t))); + } else { + text_piece(Cow::Borrowed(t), out); + } + } else if let Some(blob) = field(p, "inlineData") { + let mime = fstr(blob, "mimeType").unwrap_or_default(); + let data = fstr(blob, "data").unwrap_or_default(); + out.push(if mime.starts_with("image/") { + Piece::Image { + media_type: Some(mime), + bytes: Some(base64_len(data)), + data: Some(data), + } + } else { + Piece::Other("inlineData") + }); + } else if let Some(file) = field(p, "fileData") { + let mime = fstr(file, "mimeType").unwrap_or_default(); + out.push(if mime.starts_with("image/") { + Piece::Image { + media_type: Some(mime), + bytes: None, + data: fstr(file, "fileUri"), + } + } else { + Piece::Other("fileData") + }); + } else if let Some(call) = field(p, "functionCall") { + let name = fstr(call, "name").unwrap_or_default(); + out.push(Piece::ToolCall { + id: Cow::Borrowed(fstr(call, "id").unwrap_or(name)), + name: Cow::Borrowed(name), + input: call.get("args").map_or(Input::Args(""), Input::Json), + }); + } else if let Some(resp) = field(p, "functionResponse") { + let name = fstr(resp, "name").unwrap_or_default(); + let body = resp.get("response").unwrap_or(&Value::Null); + out.push(Piece::ToolResult { + call_id: Cow::Borrowed(fstr(resp, "id").unwrap_or(name)), + text: gemini_response_text(body), + is_error: body.get("error").is_some() && body.get("output").is_none(), + }); + } else if field(p, "executableCode").is_some() { + out.push(Piece::Other("executableCode")); + } else if field(p, "codeExecutionResult").is_some() { + out.push(Piece::Other("codeExecutionResult")); + } + // 只带着推理签名的空块、视频片段的说明这些：没有可看的 + } +} + +/// 函数结果写成字：只有一项 `output`（或 `result`、`content`、`error`）且是字符串的取那个 +/// 字符串，别的写成 JSON。和 tw-dialect 转给别家时的写法一样 +fn gemini_response_text(v: &Value) -> Cow<'_, str> { + if let Some(o) = v.as_object() + && o.len() == 1 + && let Some(s) = ["output", "result", "content", "error"] + .iter() + .find_map(|k| o.get(*k).and_then(Value::as_str)) + { + return Cow::Borrowed(s); + } + match v { + Value::String(s) => Cow::Borrowed(s), + Value::Null => Cow::Borrowed(""), + other => Cow::Owned(other.to_string()), + } +} diff --git a/docs/config.md b/docs/config.md index 9157fe3d..70c62abe 100644 --- a/docs/config.md +++ b/docs/config.md @@ -832,6 +832,13 @@ Two limits, because the two kinds of data differ in size by three orders of magnitude: request bodies are tens of kilobytes each, a request's record a few hundred bytes. The byte limit covers bursts. +Each request and response body is kept up to 4 MiB; of a longer one, the +beginning is kept. Bodies are written with credentials and personal numbers +already taken out. A request stored under `enforce` carries the placeholders +the upstream received; anything else the redaction rules +([`security.redact`](#cfg-security-redact)) recognize is masked, in every +mode, `off` included. + @@ -839,7 +846,7 @@ a few hundred bytes. The byte limit covers bursts. |---|---|---|---| | `body_days` | integer | `7` | Days to keep request and response bodies. | | `row_days` | integer | `90` | Days to keep the record of each request (time, model, usage, cost). | -| `body_max_bytes` | integer | `2147483648` | Upper bound on the bytes bodies may take; beyond it the oldest days go first. The default is 2 GiB. | +| `body_max_bytes` | integer | `5368709120` | Upper bound on the bytes bodies may take; beyond it the oldest days go first. The default is 5 GiB. | ### `failover` diff --git a/docs/config.zh-CN.md b/docs/config.zh-CN.md index 2df50c32..d0e1217e 100644 --- a/docs/config.zh-CN.md +++ b/docs/config.zh-CN.md @@ -686,6 +686,8 @@ security: 设两个期限，是因为两类数据的体积相差三个数量级：一条请求的正文有几十 KB，一条请求记录只有几百字节。字节上限用于应对用量突增。 +每份请求和响应正文最多保存 4 MiB，更长的只保存开头。正文写入磁盘前已去掉凭据和个人号码：`enforce` 下保存的请求带着发给上游的占位符，脱敏规则（[`security.redact`](#cfg-security-redact)）认出的其他内容一律打码保存，`off` 时也一样。 + @@ -693,7 +695,7 @@ security: |---|---|---|---| | `body_days` | 整数 | `7` | 请求和响应正文保留的天数。 | | `row_days` | 整数 | `90` | 每条请求记录（时间、模型、用量、费用）保留的天数。 | -| `body_max_bytes` | 整数 | `2147483648` | 正文最多占用的字节数，超出时从最早的日期开始删除。默认 2 GiB。 | +| `body_max_bytes` | 整数 | `5368709120` | 正文最多占用的字节数，超出时从最早的日期开始删除。默认 5 GiB。 | ### `failover` diff --git a/scripts/smoke.sh b/scripts/smoke.sh index 66b48729..8e866779 100755 --- a/scripts/smoke.sh +++ b/scripts/smoke.sh @@ -497,9 +497,10 @@ else bad "历史里一条记录都没有" fi -# 在整份记录里找，也按正文找。**摘录不带密钥出门**：数据面那一节发过「我的 key 是 -# sk-ant-api03-SMOKEKEY…」，按正文找得到那两条，摘录是打过码的；只在密钥里出现的 -# 词找不到。正文是另一条通道异步落盘的，没落下来之前找不到，所以多等几轮 +# 在整份记录里找，也按正文找。**摘录不带密钥出门**：数据面那一节在拦截档下发过「我的 +# key 是 sk-ant-api03-SMOKEKEY…」，按正文找得到那两条，摘录里是发给上游的占位符（正文 +# 落盘之前就换过了）；只在密钥里出现的词找不到。正文是另一条通道异步落盘的，没落下来 +# 之前找不到，所以多等几轮 GOT="" for _ in $(seq 1 20); do C=$(post /history/search '{"q":"我的 key","content":true}') @@ -509,13 +510,21 @@ p = json.loads(raw) hits = p.get("hits") or [] good = (p.get("stopped") in ("end", "full") and len(hits) >= 2 and all(h["side"] == "request" for h in hits) - and all("sk-an" in h["after"] for h in hits) and "SMOKEKEYAAAA" not in raw) + and all("<>" in h["after"] for h in hits) and "SMOKEKEYAAAA" not in raw) print("ok" if good else raw[:600])' < "$TMP/out" 2>&1) [ "$C" = "200" ] && [ "$GOT" = "ok" ] && break sleep 0.25 done -[ "$C" = "200" ] && [ "$GOT" = "ok" ] && ok "POST /history/search 按正文找到了那两条，摘录里的密钥打了码" \ +[ "$C" = "200" ] && [ "$GOT" = "ok" ] && ok "POST /history/search 按正文找到了那两条，摘录里是占位符" \ || bad "按正文找没找对（$C）" "$GOT" +# **落盘的不是原文**：拦截档下存的是发给上游的那一份，盘上没有一个文件带着那把密钥。 +# 读出来时再打码只是第二道。上面那一轮已经等到两条的正文都落了盘 +BLOBS="$THINKWATCH_HOME/blobs" +if [ -d "$BLOBS" ] && ! grep -rqa "SMOKEKEY" "$BLOBS" && grep -rqa "<>" "$BLOBS"; then + ok "落盘的正文里没有那把密钥，存的是发给上游的占位符" +else + bad "落盘的正文不对" "带着密钥的：$(grep -rla "SMOKEKEY" "$BLOBS" 2>/dev/null | head -3)" +fi C=$(post /history/search '{"q":"smokekeyaaaa","content":true}') if [ "$C" = "200" ] && ! grep -q "SMOKEKEYAAAA" "$TMP/out" \ && python3 -c 'import sys, json; sys.exit(0 if not json.load(sys.stdin)["hits"] else 1)' < "$TMP/out"; then @@ -524,6 +533,37 @@ else bad "按密钥里的字找到了东西（$C）" "$(head -c 400 "$TMP/out")" fi +# 会话读成一段对话。**每一次会话都读得出来**，和会话详情一轮对一轮，密钥打了码。数据面 +# 那一节发过的请求都在里面：整包的、流式的、被切断的、中途走掉的、转去 Bedrock 的 —— +# Bedrock 的二进制帧在网关进门时转成 SSE 存下来，这里要从存下的那一份里读回那句话 +SIDS=$(ctl /sessions | python3 -c 'import json,sys;print(" ".join(s["id"] for s in json.load(sys.stdin)))') +TURNS=0; WRONG=""; ALL="" +for SID in $SIDS; do + C=$(get "/sessions/$SID/transcript") + cp "$TMP/out" "$TMP/transcript.json" + ctl "/sessions/$SID" > "$TMP/detail.json" + GOT=$(python3 -c 'import json, sys +raw = open(sys.argv[1]).read() +t, d = json.loads(raw), json.load(open(sys.argv[2])) +same = [x["id"] for x in t["turns"]] == [str(x["id"]) for x in d["turns"]] +print(len(t["turns"]) if same and "SMOKEKEYAAAA" not in raw else raw[:400])' \ + "$TMP/transcript.json" "$TMP/detail.json" 2>&1) + if [ "$C" = "200" ] && [[ "$GOT" =~ ^[0-9]+$ ]]; then + TURNS=$((TURNS + GOT)); ALL="$ALL$(cat "$TMP/transcript.json")" + else + WRONG="$WRONG $SID（$C）：$GOT" + fi +done +if [ -z "$SIDS" ]; then + bad "一次会话都没有，读不了对话记录" +elif [ -n "$WRONG" ]; then + bad "GET /sessions/{id}/transcript 不对" "$WRONG" +elif ! grep -q '"text":"bedrock-stream"' <<<"$ALL"; then + bad "对话记录里没读出 Bedrock 流式回答的那句话" "$(head -c 600 <<<"$ALL")" +else + ok "GET /sessions/{id}/transcript：$(wc -w <<<"$SIDS" | tr -d ' ') 次会话、$TURNS 轮都读得出来，和会话详情一轮对一轮，密钥打了码" +fi + # ---------------------------------------------------------------- 诊断包 step "诊断包不带密钥出门" get /diagnostics >/dev/null