From dc0f88e930c820824b4dc7609283d05c6c1294b9 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Fri, 2 Oct 2026 16:43:29 +0800 Subject: [PATCH 1/2] feat(control): read a session as a conversation transcript (#251) GET /sessions/{id}/transcript turns a session's stored request and response bodies into the conversation the desktop app shows: per turn, the messages new in that request, the answer decoded from the stored response, and what could not be shown (gaps). Control-plane protocol 32. Clients resend the whole history every turn, so a request's new messages are what is left after the previous readable request's messages, compared by fingerprint with cache_control, signatures, key order and reasoning ignored (some clients drop earlier reasoning from the history). The first leftover assistant message is the previous turn's answer and is dropped, unless that answer could not be read in full. No prefix match, or an unreadable request before it, restarts with the whole history. Requests are read from the client's own JSON instead of tw-dialect's IR: the IR drops exactly what a transcript needs (where a system message sits, tool-only messages, server-tool blocks, files). Answers reuse tw-dialect's per-format stream parsers, and the blocks those skip are recognized on the same frame. Every string is masked with mask_body on the way out; images carry only their type and size. 300 turns at about 1.2 MB a request build in about 0.7 s (release). Co-authored-by: Claude Opus 5.5 --- Cargo.lock | 1 + crates/tw-api/src/ep.rs | 2 + crates/tw-api/src/lib.rs | 123 +- crates/tw-api/src/ts.rs | 52 + crates/tw-control/src/lib.rs | 40 +- crates/tw-control/tests/transcript.rs | 222 +++ crates/tw-store/Cargo.toml | 2 + crates/tw-store/src/db.rs | 10 + crates/tw-store/src/lib.rs | 1 + crates/tw-store/src/transcript/answer.rs | 432 +++++ crates/tw-store/src/transcript/mod.rs | 1932 ++++++++++++++++++++++ crates/tw-store/src/transcript/read.rs | 786 +++++++++ scripts/smoke.sh | 31 + 13 files changed, 3627 insertions(+), 7 deletions(-) create mode 100644 crates/tw-control/tests/transcript.rs create mode 100644 crates/tw-store/src/transcript/answer.rs create mode 100644 crates/tw-store/src/transcript/mod.rs create mode 100644 crates/tw-store/src/transcript/read.rs diff --git a/Cargo.lock b/Cargo.lock index 426e1d7..fa4f78e 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/crates/tw-api/src/ep.rs b/crates/tw-api/src/ep.rs index 9802c14..cf15a80 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 7fb6c5d..178cc46 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -671,7 +671,11 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// 31 起出站脱敏有**身份证号和银行卡号**两条出厂就开的规则：[`Matcher`] 多了 /// `cn-resident-id` 和 `bank-card`（[`CardNetwork`]、[`CardPrefix`]），[`SecretKind`] 多了 /// `personal`。照 30 写的界面说不出这两条规则按什么认。 -pub const CONTROL_API_VERSION: u32 = 31; +/// +/// **32 起会话能读成一段对话**：新端点 `GET /sessions/{id}/transcript`（[`Transcript`]） +/// 从存下来的正文里读出每一轮新说的话、回答、推理、工具调用和结果，读不到的地方逐轮说出来 +/// （[`TranscriptGap`]）。照 31 写的界面只有每一轮的用量和金额。 +pub const CONTROL_API_VERSION: u32 = 32; #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] @@ -4026,6 +4030,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 0fd9ad1..f2f2814 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-control/src/lib.rs b/crates/tw-control/src/lib.rs index 7220845..e99265d 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) @@ -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/tests/transcript.rs b/crates/tw-control/tests/transcript.rs new file mode 100644 index 0000000..8a11dde --- /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-store/Cargo.toml b/crates/tw-store/Cargo.toml index 4e7784b..ec40258 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/db.rs b/crates/tw-store/src/db.rs index 5c20a10..9f247eb 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 59c911b..dd1dcff 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/transcript/answer.rs b/crates/tw-store/src/transcript/answer.rs new file mode 100644 index 0000000..81c9d98 --- /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 0000000..d5a8e67 --- /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 0000000..5f88faf --- /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/scripts/smoke.sh b/scripts/smoke.sh index 66b4872..36e5999 100755 --- a/scripts/smoke.sh +++ b/scripts/smoke.sh @@ -524,6 +524,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 From 33b6ae91c5ac5591f3d1c8c67ee89c18e7c40e7e Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Fri, 2 Oct 2026 17:49:34 +0800 Subject: [PATCH 2/2] Store bodies with secrets replaced; keep 4 MiB of each answer; 5 GiB of bodies by default (#252) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Store bodies with secrets replaced; keep 4 MiB of each answer; 5 GiB of bodies by default Request bodies were stored as the client sent them and masked only when read back, so every credential, resident ID and card number a client sent sat on disk as it was. Bodies are now redacted before they are written, off the forwarding path (BodyRecord::for_disk, run on a blocking thread where twcore hands bodies to the store): - Every value the redaction rules recognize is taken out, whatever the mode, `off` included. Under `enforce` a request is stored in the client's format with the placeholders the upstream received; everything else is masked the way the security log shows it (internal addresses included, which the log prints as they are). Then the whole text goes through tw_secret::mask_body, the same masking the read side applies, for what the rules miss. Read-side masking stays. - Placeholders were numbered per hop in the order of that hop's body, so a converted or re-serialized hop could number the same secret differently from the stored request. guard::look now numbers the client body once and every hop continues that ledger: one value, one placeholder across hops, the stored request and the stored answer. - Ledger::avoiding / Scheme::find_in (tw-guard, additive): numbering skips placeholders already written in the text, and guard::hits drops rule hits that sit on one of ours (the connection-string rule took `app:<>@` for a password). Replay and fixture export use it, so a value found again in a stored request never takes a number the stored text already uses and an echo is not restored to the wrong value. - mask_body is idempotent (a 5-character head before `…` and a 4-character tail after it are left alone) and walks bytes instead of collecting a Vec (16 MB for a 4 MB body). Sizes: - The response tap keeps tw_api::BODY_MAX (4 MiB, the store's MAX_ONE is now that constant) plus a 64 KiB margin so a secret across the cut is recognized whole before the store cuts at 4 MiB. The tap grows with the answer and never past its cap. Requests over the window are copied rather than sliced, so a 256 MB request is not kept alive in the queue. - What waits to be written is capped in bytes (QUEUED_MAX, 32 MiB) instead of by count alone: 64 bodies of 4 MiB would be 256 MiB. Each record holds its share until the store accepts it, and the hand-off to the store has a single slot. - Blobs::put_with_len records the original length when it cuts a body itself; a request over 4 MiB used to be stored cut with nothing saying so, and replay sent the half JSON. - retention.body_max_bytes defaults to 5 GiB. tw-store also ran its own hourly gc with hardcoded 7 days, 90 days and 2 GiB next to the one in twcore that reads `retention`, so any setting above the defaults was cut back every hour; that loop and its constants are gone. Consumers: GET /request/{id} shows what was stored (masking it again changes nothing); history search no longer finds text that only occurs inside a secret; replay sends the stored request, with placeholders or masked values instead of secrets; session fingerprints and DeepSeek Harness session_log_bytes still come from the in-memory original. The control-plane protocol and the request store schema are unchanged; old bodies age out within body_days. Co-Authored-By: Claude Opus 5.5 * fix(gateway): keep no answer when there is nowhere to store it With the tap now holding up to 4 MiB per answer, an Ending without a body sink (recording did not start) still buffered every answer only to drop it at the end. It now counts the bytes and keeps nothing. Co-Authored-By: Claude Opus 5.5 * docs(store): an answer to search is now up to 4 MB, not 256 KB Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Claude Opus 5.5 --- bin/twcore/src/main.rs | 38 +- crates/tw-api/src/lib.rs | 15 +- crates/tw-config/src/lib.rs | 2 +- crates/tw-config/src/retention.rs | 4 +- crates/tw-config/tests/manual/schema.rs | 6 +- crates/tw-control/src/lib.rs | 2 +- crates/tw-control/src/replay.rs | 23 +- crates/tw-control/tests/replay.rs | 62 ++- crates/tw-control/tests/stored_bodies.rs | 387 ++++++++++++++ crates/tw-gateway/src/bodies.rs | 519 +++++++++++++++++-- crates/tw-gateway/src/ending.rs | 53 +- crates/tw-gateway/src/fixture.rs | 9 +- crates/tw-gateway/src/guard.rs | 198 ++++++- crates/tw-gateway/src/server/pipeline.rs | 78 ++- crates/tw-gateway/src/server/pipeline/hop.rs | 15 +- crates/tw-gateway/src/state.rs | 10 +- crates/tw-gateway/tests/claude_desktop.rs | 2 +- crates/tw-guard/src/redact/replace.rs | 105 ++++ crates/tw-secret/src/mask.rs | 90 +++- crates/tw-store/src/blobs.rs | 49 +- crates/tw-store/src/search/text.rs | 4 +- crates/tw-store/src/task.rs | 42 +- docs/config.md | 9 +- docs/config.zh-CN.md | 4 +- scripts/smoke.sh | 19 +- 25 files changed, 1549 insertions(+), 196 deletions(-) create mode 100644 crates/tw-control/tests/stored_bodies.rs diff --git a/bin/twcore/src/main.rs b/bin/twcore/src/main.rs index 815ee3f..ebdc988 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/lib.rs b/crates/tw-api/src/lib.rs index 178cc46..494a8ed 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -1706,7 +1706,7 @@ pub struct RetentionView { /// 正文总共最多占多少字节 pub body_max_bytes: u64, /// 正文现在实际占了多少。**不是配置，是现状** —— 没有它， - /// 「2 GB 上限」是个用户无从判断松紧的数字 + /// 「5 GB 上限」是个用户无从判断松紧的数字 pub body_bytes_now: u64, } @@ -3676,14 +3676,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, } diff --git a/crates/tw-config/src/lib.rs b/crates/tw-config/src/lib.rs index 70bf42c..135e111 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 0e95d99..5afb769 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 5cc1142..a05baa6 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 e99265d..93ac21a 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -390,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, }, diff --git a/crates/tw-control/src/replay.rs b/crates/tw-control/src/replay.rs index 13f4415..ace9d55 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 0befd91..647e4fa 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 0000000..5932e3e --- /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-gateway/src/bodies.rs b/crates/tw-gateway/src/bodies.rs index d1f62bf..8261096 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 1464616..efd4fa8 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 6aef83f..c64f586 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 e824170..2b91263 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 67b2837..2d3e7d8 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 2af55cb..7d8456b 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 82278ea..f43527e 100644 --- a/crates/tw-gateway/src/state.rs +++ b/crates/tw-gateway/src/state.rs @@ -145,7 +145,7 @@ pub struct AppState { /// **有界通道，满了就丢。**直接调用意味着文件 I/O 跑在转发那条路上 /// —— 一次慢磁盘写就变成一次慢请求，而观测永远不该有这个权力。 /// `None` 表示观测层没起来，那时什么都不做。 - body_sink: Arc>>, + body_sink: Arc>>, /// 每个上游最近一次报的订阅额度。 /// /// **在内存里，不落库。**它是「现在还剩多少」，不是历史 —— 存一份 @@ -308,9 +308,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); } @@ -364,7 +364,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 94c3db2..5050e64 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 8e962b2..49c245c 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 606b559..6694f1a 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/src/blobs.rs b/crates/tw-store/src/blobs.rs index fe9f2bd..95f472f 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/search/text.rs b/crates/tw-store/src/search/text.rs index 208af6d..7757914 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 e4e33ad..d890bd5 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/docs/config.md b/docs/config.md index 9157fe3..70c62ab 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 2df50c3..d0e1217 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 36e5999..8e86677 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