diff --git a/Cargo.lock b/Cargo.lock index 4827b683..ad3c92ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3252,6 +3252,7 @@ dependencies = [ name = "tw-dialect" version = "0.64.1" dependencies = [ + "memchr", "serde", "serde_json", ] @@ -3287,6 +3288,7 @@ dependencies = [ "http-body-util", "hyper", "hyper-util", + "memchr", "rand 0.10.2", "reqwest", "rustls", @@ -3324,6 +3326,7 @@ version = "0.64.1" dependencies = [ "base64", "bytes", + "memchr", "regex", "serde", "serde_json", @@ -3424,6 +3427,7 @@ dependencies = [ name = "tw-watch" version = "0.64.1" dependencies = [ + "libc", "notify", "tempfile", "thiserror", diff --git a/Cargo.toml b/Cargo.toml index 8391e8b6..b1ad8111 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -103,6 +103,8 @@ flate2 = "1" # regex-automata / regex-syntax 已经在依赖树里（tracing 的 EnvFilter）， # 所以加它的增量很小 regex = "1" +# 在整个请求体里找一小段字节（SIMD）。serde_json、regex 本来就依赖它，不多编任何东西 +memchr = "2" # 保格式地改 TOML。接管 Codex 要写它的 config.toml，而重新序列化 # 会把用户的注释和排版一起抹掉 toml_edit = "0.25" diff --git a/bin/twcore/src/main.rs b/bin/twcore/src/main.rs index 9cf231dc..0db6e860 100644 --- a/bin/twcore/src/main.rs +++ b/bin/twcore/src/main.rs @@ -840,6 +840,9 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> 每一轮都重新读配置，所以改完期限不用重启 —— 下一轮就按新的算。 扫的是目录和一条 DELETE，空跑一次的代价可以忽略。 + + **在阻塞线程上跑，删正文时不拿记录库的锁**（见 `tw_store::task::gc`）：一次删几个 + GB 可以要好几秒，以前那期间这把锁一直攥着，记录和查询都在等。 */ if let Some(rec) = store.clone() { let gw = state.clone(); @@ -847,16 +850,18 @@ fn cmd_serve(path: &Path, port: Option, safe: bool, parent: Option) -> loop { let r = gw.config().retention.clone(); let freed = { - let g = rec.lock().await; - g.gc( - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map(|d| d.as_millis() as i64) - .unwrap_or(0), - r.body_days, - r.row_days, - r.body_max_bytes, - ) + let rec = rec.clone(); + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0); + let (body_days, row_days, max) = + (r.body_days, r.row_days, r.body_max_bytes); + tokio::task::spawn_blocking(move || { + tw_store::task::gc(&rec, now, body_days, row_days, max) + }) + .await + .unwrap_or(0) }; if freed > 0 { tracing::info!( @@ -1013,7 +1018,14 @@ fn build_store( // 每记下一行请求，交给网关的密钥用量上限结算（见 `tw_control::key_limits`） settled: tw_store::SettleHook, ) -> Option>> { - let events = bus.subscribe(); + // **存储层自己的那条通道，不订阅界面的广播**（见 `EventBus::record_feed`）：广播上 + // 掉队丢的是界面上的几个点，记录丢的是那几个请求 + let Some(events) = bus.record_feed() else { + tracing::warn!( + "request recording is already attached to the event bus; not starting a second one" + ); + return None; + }; let (db, blobs) = match tw_store::open(dir) { Ok(opened) => opened, Err(e) => { diff --git a/crates/tw-api/src/ep.rs b/crates/tw-api/src/ep.rs index c4d6c3d6..aaf98bfa 100644 --- a/crates/tw-api/src/ep.rs +++ b/crates/tw-api/src/ep.rs @@ -64,8 +64,9 @@ 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; + /// 一次会话读成一段对话：每一轮新说的话、回答、工具调用和结果（已脱敏）。可以只要 + /// 从某一轮起的那些 + SessionTranscript: GET "/sessions/{id}/transcript" [id], api::TranscriptQuery => 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 122ae583..db1a1903 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -798,7 +798,14 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// **40 起模型规格多了会不会推理、收不收图**：[`ModelRow`] 多了 `reasoning`、`image_input` /// 和它们的来源，[`ModelSpecSave`] 多了 `reasoning`、`image_input`（四项都空才是删掉）。 /// 照 39 写的界面保存规格时会把这两项丢掉。 -pub const CONTROL_API_VERSION: u32 = 40; +/// +/// **41 起会话的对话视图可以只取新的轮次，试算给出每家分到的份额**： +/// `GET /sessions/{id}/transcript` 的请求从 `()` 变成 [`TranscriptQuery`]（`from_turn`， +/// 不带就是整份），[`Transcript`] 多了 `total_turns` 和 `settled_turns`（前面这么多轮 +/// 不会再变，下次从这里取）；[`DryRunCandidate`] 多了 `share`（负载均衡组里这一家按 +/// 当前权重和系数分到的份额，别的组和这一轮不参与的是 None）。照 40 写的界面取对话视图 +/// 时传 `null` 会被拒。 +pub const CONTROL_API_VERSION: u32 = 41; #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] @@ -4697,6 +4704,9 @@ pub struct SessionDetail { /// 太大的只留开头，没存下来的也有。读不到的地方，那一轮的 `gaps` 说出来。 /// /// **已脱敏**，和请求详情里的正文同一套打码。图片只说类型和大小，从不带数据。 +/// +/// **可以只要后面的几轮**（[`TranscriptQuery::from_turn`]）：会话开着时每来一轮重读一次， +/// 前面不会再变的那些（`settled_turns`）不必再传一遍。 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] pub struct Transcript { @@ -4705,10 +4715,31 @@ pub struct Transcript { /// Gemini 的 `systemInstruction`，Chat 和 Responses 还有开头连着的 system、developer /// 消息，几段之间空一行。没有是 null pub system: Option, - /// 和 [`SessionDetail::turns`] 同样的请求，同样的顺序 + /// 这次会话一共几轮，就是 [`SessionDetail::turns`] 有几条。给了 + /// [`TranscriptQuery::from_turn`] 时，`turns` 是其中从那一轮起的那些 + pub total_turns: u32, + /// 开头这么多轮不会再变了：它们的正文都落了盘（或者不会再来），后面的请求不会再改写 + /// 它们，这次会话里还在跑的请求落库时也不会排到它们前面。**下一次从这里要起** + /// （`from_turn = settled_turns`），留着手上的前这么多轮，换掉后面的。 + /// + /// 后面那几轮还会变：最后一轮的回答可能还没落盘、工具调用的号会被下一轮换成客户端 + /// 记下的那个、一个开始得早结束得晚的请求会插到它们中间 + pub settled_turns: u32, + /// 和 [`SessionDetail::turns`] 同样的请求，同样的顺序。给了 + /// [`TranscriptQuery::from_turn`] 的，只有从那一轮起的那些 pub turns: Vec, } +/// `GET /sessions/{id}/transcript` 的查询串。 +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct TranscriptQuery { + /// 只要从这一轮起的那些（从 0 数，[`Transcript::turns`] 里的第几条）。不给是整段； + /// 比总轮数还大的，`turns` 是空的 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub from_turn: Option, +} + /// 对话里的一轮，就是会话里的一个请求。 /// /// 客户端每一轮都把整段历史发上来：请求 i 的消息 = 请求 i-1 的消息 + 上一轮的回答 + 新的 @@ -5066,6 +5097,12 @@ pub struct DryRunCandidate { /// 分得少，没有样本的那一项算 1。只在 `balance_by` 不是 `weights` 的 `load-balance` 里有 #[serde(default, skip_serializing_if = "Option::is_none")] pub balance_factor: Option, + /// 照此刻的数字轮下去，它分到的新对话占这个组的多少，0 到 1：有效权重（权重 × 系数） + /// 除以这一轮参加的各家之和，和数据面挑排头是同一份数字。只在 `load-balance` 里有， + /// 同一个组里有份额的各家加起来是 1。熔断着、冷却着、并发数满着的这一轮不参加，没有 + /// 份额 —— 全都这样时都参加 + #[serde(default, skip_serializing_if = "Option::is_none")] + pub share: Option, } /// 一个要转换格式的候选上游。 diff --git a/crates/tw-api/src/ts.rs b/crates/tw-api/src/ts.rs index 6ba442ac..4843bf7f 100644 --- a/crates/tw-api/src/ts.rs +++ b/crates/tw-api/src/ts.rs @@ -301,7 +301,12 @@ mod tests { #[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: { req: TranscriptQuery; res: Transcript };")); + assert!( + decl_of(&ts, "TranscriptQuery").contains("from_turn?: number"), + "{}", + decl_of(&ts, "TranscriptQuery") + ); assert!(ts.contains( " SessionTranscript: { method: \"GET\", path: \"/sessions/{id}/transcript\", params: [\"id\"], format: \"json\" }," )); @@ -309,6 +314,8 @@ mod tests { for field in [ "session: string", "system: string | null", + "total_turns: number", + "settled_turns: number", "turns: Array", ] { assert!(transcript.contains(field), "{field}: {transcript}"); diff --git a/crates/tw-control/src/dryrun.rs b/crates/tw-control/src/dryrun.rs index 98f2563d..6ea33aca 100644 --- a/crates/tw-control/src/dryrun.rs +++ b/crates/tw-control/src/dryrun.rs @@ -51,17 +51,18 @@ fn order_like_the_data_plane( } /// 一家候选的排序依据：典型的快慢（从发出去到回答的第一段内容）、最近的成功率、按它们算出 -/// 的系数。**只给顺序真用到的那几样**（[`runtime_facts`] 只取用得上的）—— `url-test` 看快慢； -/// `load-balance` 按 `balance_by` 看快慢、成败，系数乘在权重上。用不着的、没有样本的 -/// 是 None。 +/// 的系数，和 `load-balance` 照这些数字分给它的份额。**只给顺序真用到的那几样** +/// （[`runtime_facts`] 只取用得上的）—— `url-test` 看快慢；`load-balance` 按 `balance_by` +/// 看快慢、成败，系数乘在权重上。用不着的、没有样本的是 None。 struct Basis { ttfb_ms: Option, success_rate: Option, balance_factor: Option, + share: Option, } -/// 每一家候选的排序依据，按 `d.candidates` 的名字查。系数和数据面一样按这份候选、 -/// 这份数字算（`tw_engine::balance_factors`） +/// 每一家候选的排序依据，按 `d.candidates` 的名字查。系数、份额和数据面一样按这份候选、 +/// 这份数字算（`tw_engine::balance_factors`、`tw_engine::weighted::shares`） fn bases( d: &tw_engine::Decision, runtime: Option<&(&tw_engine::Group, tw_engine::Facts)>, @@ -69,11 +70,18 @@ fn bases( let Some((g, f)) = runtime else { return Default::default(); }; - let factors = if g.kind == tw_engine::GroupType::LoadBalance && !g.balance_by.is_weights() { + let balanced = g.kind == tw_engine::GroupType::LoadBalance; + let factors = if balanced && !g.balance_by.is_weights() { tw_engine::balance_factors(g.balance_by, &d.candidates, f) } else { Vec::new() }; + // 份额只按权重分时也有：权重本身就是比例 + let shares = if balanced { + tw_engine::weighted::shares(g, &d.candidates, f) + } else { + Vec::new() + }; d.candidates .iter() .enumerate() @@ -82,6 +90,7 @@ fn bases( ttfb_ms: f.ttfb_ms.get(name).copied(), success_rate: f.success.get(name).copied(), balance_factor: factors.get(i).copied(), + share: shares.get(i).copied().flatten(), }; (name.clone(), basis) }) @@ -325,7 +334,7 @@ pub async fn dry_run( // `load-balance` 的候选带上各自的权重：排头的为什么是它，一半在这个数里 let balanced = group.filter(|g| g.kind == tw_engine::GroupType::LoadBalance); // 每一家收到的模型名，和为什么不是请求里写的那个；顺序看运行时数字的，再加上 - // 每一家的那几个数字（权重、快慢、成功率、系数）—— 「为什么轮到它」 + // 每一家的那几个数字（权重、快慢、成功率、系数、份额）—— 「为什么轮到它」 // 要能从这里看出来 let mut basis_of = bases(&d, runtime.as_ref()); out.candidate_models = out @@ -344,6 +353,7 @@ pub async fn dry_run( ttfb_ms: basis.as_ref().and_then(|b| b.ttfb_ms), success_rate: basis.as_ref().and_then(|b| b.success_rate), balance_factor: basis.as_ref().and_then(|b| b.balance_factor), + share: basis.as_ref().and_then(|b| b.share), } }) .collect(); diff --git a/crates/tw-control/src/key_limits.rs b/crates/tw-control/src/key_limits.rs index 891d447e..08df8fce 100644 --- a/crates/tw-control/src/key_limits.rs +++ b/crates/tw-control/src/key_limits.rs @@ -57,10 +57,16 @@ pub fn rebuild(gw: &tw_gateway::AppState, db: &tw_store::Db) { pub fn follow(gw: &tw_gateway::AppState, store: Arc>) { // 弱引用：网关的账拿着这个办法，办法再拿着账就是一个圈，谁都放不掉 let limits = Arc::downgrade(&gw.key_limits); + // **接上时的运行时，记在这里**：结算在存储层自己的记录线程上，那条线程不在运行时里，在 + // 那儿问「当前的运行时」问不到，新的一期就永远不会从记录里加回来 + let here = tokio::runtime::Handle::try_current().ok(); gw.key_limits .reread_with(Arc::new(move |m: tw_gateway::key_limits::Moved| { let (limits, store) = (limits.clone(), store.clone()); - let Ok(rt) = tokio::runtime::Handle::try_current() else { + let Some(rt) = here + .clone() + .or_else(|| tokio::runtime::Handle::try_current().ok()) + else { return; }; rt.spawn(async move { diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index 9ee2f0fe..59af3e47 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -305,10 +305,15 @@ async fn overview(State(s): State) -> Json { let cfg = s.config(); let cfg = &*cfg; let engine = cfg.engine(); - // 正文现在占了多少。**扫的是目录，所以放在这儿算一次**，不塞进 - // 下面那个视图表达式里 + // 正文现在占了多少。**扫的是目录**：每份正文问一次大小，几万份要几十毫秒，所以在阻塞 + // 线程上扫，也不拿着记录库的锁扫 —— 正文目录只是一个路径 let body_bytes_now = match &s.store { - Some(st) => st.lock().await.blobs().total_bytes(), + Some(st) => { + let root = st.lock().await.blobs().root().to_path_buf(); + tokio::task::spawn_blocking(move || tw_store::Blobs::new(root).total_bytes()) + .await + .unwrap_or(0) + } None => 0, }; Json(tw_api::Overview { @@ -724,8 +729,8 @@ async fn route_stats( ) -> Result, Fail> { let (from, to) = range(q.from_ms, q.to_ms); let store = need_store(&s)?; - let g = store.lock().await; - Ok(Json(g.db().route_stats(from, to).map_err(records)?)) + let stats = on_store(store, move |g| g.db().route_stats(from, to)).await?; + Ok(Json(stats.map_err(records)?)) } /// 一段时间的汇总。不给参数就是「今天」。 @@ -1226,11 +1231,17 @@ async fn storage(State(s): State) -> Json { forwarding_affected: false, }); }; - let g = store.lock().await; + let root = store.lock().await.blobs().root().to_path_buf(); + let rows = on_store(store, |g| g.db().count().unwrap_or(0)) + .await + .unwrap_or(0); + let blob_bytes = tokio::task::spawn_blocking(move || tw_store::Blobs::new(root).total_bytes()) + .await + .unwrap_or(0); Json(tw_api::StorageStatus { recording: true, - rows: g.db().count().unwrap_or(0), - blob_bytes: g.blobs().total_bytes(), + rows, + blob_bytes, // **永远是 false。**观测挂了，代理照跑。哪天有人想改成 // true，先回去读那一节。 forwarding_affected: false, @@ -1252,6 +1263,21 @@ pub(crate) fn need_store( }) } +/// 在阻塞线程上拿着记录库做 `f`。 +/// +/// **查库、扫正文目录都不占异步线程**：一条聚合几毫秒到几十毫秒，同一个线程上的别的请求 +/// （控制面的、转发的）都得等它。锁也在那个线程上等：记录器正拿着它写一批行时，等的是那个 +/// 阻塞线程，不是异步线程。 +pub(crate) async fn on_store( + store: &Arc>, + f: impl FnOnce(&tw_store::Recorder) -> T + Send + 'static, +) -> Result { + let store = store.clone(); + tokio::task::spawn_blocking(move || f(&store.blocking_lock())) + .await + .map_err(internal) +} + fn history_row( r: tw_store::RequestRow, security: Vec, @@ -1652,15 +1678,13 @@ async fn sessions( let Some(store) = &s.store else { return Json(Vec::new()); }; - let g = store.lock().await; - Json( - g.db() - .sessions(within(&q), list_limit(&q)) - .unwrap_or_default() - .iter() - .map(session_view) - .collect(), - ) + let (within, limit) = (within(&q), list_limit(&q)); + let rows = on_store(store, move |g| g.db().sessions(within, limit)) + .await + .ok() + .and_then(Result::ok) + .unwrap_or_default(); + Json(rows.iter().map(session_view).collect()) } async fn session_detail( @@ -1673,23 +1697,20 @@ async fn session_detail( msg!("control.store_off" => "Request recording is not running."), )); }; - let g = store.lock().await; - let session = g - .db() - .sessions(None, 500) - .unwrap_or_default() - .iter() - .find(|x| x.id == id) - .map(session_view) - .ok_or_else(|| no_such_session(&id))?; - let turns = g - .db() - .turns(&id) - .unwrap_or_default() - .iter() - .map(turn_view) - .collect(); - Ok(Json(tw_api::SessionDetail { session, turns })) + // **按会话号直接取**：以前取最近五百次会话再从里面找，更早的会话点开是 404 + let (session, turns) = { + let id = id.clone(); + on_store(store, move |g| { + Ok::<_, tw_store::DbError>((g.db().session(&id)?, g.db().turns(&id)?)) + }) + .await? + .map_err(records)? + }; + let session = session.ok_or_else(|| no_such_session(&id))?; + Ok(Json(tw_api::SessionDetail { + session: session_view(&session), + turns: turns.iter().map(turn_view).collect(), + })) } fn no_such_session(id: &str) -> Fail { @@ -1699,28 +1720,53 @@ fn no_such_session(id: &str) -> Fail { ) } -/// 一次会话读成一段对话（见 `tw_store::transcript`）。 +/// 一次会话读成一段对话（见 `tw_store::transcript`）。给了 `from_turn` 的只交出从那一轮起的 +/// 那些（见 [`tw_api::Transcript::settled_turns`]）。 /// /// **放到阻塞线程上跑**，库只在取行的时候锁一下：几百轮的会话要读几百份正文、解析几百 MB -/// 的 JSON，和按正文找是同一个道理（见 [`history_search`]）—— 记录和正文落盘走的是同一把锁。 +/// 的 JSON，和按正文找是同一个道理（见 [`history_search`]）—— 记录落库走的是同一把锁。读过 +/// 的记着（`Recorder::transcripts`），下一次只读新来的那几轮。 async fn session_transcript( State(s): State, axum::extract::Path(id): axum::extract::Path, + axum::extract::Query(q): axum::extract::Query, ) -> 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)?; + let store = need_store(&s)?.clone(); + // 这次会话里还在跑的请求：落库时排在它开始的那一刻，不一定排在最后 + let running = s + .gateway + .bus + .live() + .running + .into_iter() + .filter(|r| r.session.as_deref() == Some(id.as_str())) + .map(|r| (r.at_ms as i64, r.id as i64)) + .min(); + let from_turn = q.from_turn.unwrap_or(0) as usize; + let transcript = tokio::task::spawn_blocking(move || { + let (rows, blobs, transcripts) = { + let g = store.blocking_lock(); + let rows = g.db().session_requests(&id).map_err(records)?; + // 正文目录只是一个路径，读它不需要锁 + let blobs = tw_store::Blobs::new(g.blobs().root().to_path_buf()); + (rows, blobs, g.transcripts()) + }; + if rows.is_empty() { + return Err(no_such_session(&id)); + } + Ok(transcripts.read( + &blobs, + &tw_store::transcript::Ask { + session: &id, + rows: &rows, + from_turn, + running, + now_ms: now_ms(), + }, + )) + }) + .await + .map_err(internal)??; Ok(Json(transcript)) } diff --git a/crates/tw-control/tests/dryrun.rs b/crates/tw-control/tests/dryrun.rs index 7cdebd2d..cefc9aac 100644 --- a/crates/tw-control/tests/dryrun.rs +++ b/crates/tw-control/tests/dryrun.rs @@ -804,6 +804,88 @@ async fn only_the_numbers_the_order_uses_are_shown() { ); } +/// 每一家的份额：`load-balance` 组里各家照此刻的数字分到的新对话的比例，加起来是 1。 +/// 界面不再自己拿权重和系数去乘 —— 那是数据面的算法，各算各的迟早对不上（停着的、满着的 +/// 哪些不参加，系数取整之后是多少） +#[tokio::test] +async fn each_balanced_candidate_carries_the_share_the_data_plane_gives_it() { + let shares = |r: &tw_api::DryRunResult| { + r.candidate_models + .iter() + .map(|c| (c.provider.clone(), c.share)) + .collect::>() + }; + let weighted = CFG.replace( + " providers: [官方, 中转]\n", + " providers: [{ name: 官方, weight: 3 }, 中转]\n", + ); + + // 只按权重分：份额就是权重的比例 + let (_d, app, gw) = app_and_gateway(&weighted); + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + let s = shares(&r); + assert_eq!((s["官方"], s["中转"]), (Some(0.75), Some(0.25)), "{r:?}"); + + // 熔断着的这一轮不参加：没有份额，剩下的那一家拿全部 + while gw.health.is_available("中转") { + gw.health.record_failure("中转"); + } + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + assert_eq!(r.circuit_open, ["中转"]); + let s = shares(&r); + assert_eq!((s["官方"], s["中转"]), (Some(1.0), None), "{r:?}"); + // 全都熔断着：都参加，照权重分（数据面照样一家家试） + while gw.health.is_available("官方") { + gw.health.record_failure("官方"); + } + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + let s = shares(&r); + assert_eq!((s["官方"], s["中转"]), (Some(0.75), Some(0.25)), "{r:?}"); + + // 并发数满着的同样不参加 + let busy = weighted.replace( + " base_url: https://relay.example.com\n", + " base_url: https://relay.example.com\n max_concurrent: 1\n", + ); + let (_d, app, gw) = app_and_gateway(&busy); + let held = gw.slots.try_take("中转").expect("中转有空位"); + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + let s = shares(&r); + assert_eq!((s["官方"], s["中转"]), (Some(1.0), None), "{r:?}"); + drop(held); + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + assert_eq!(shares(&r)["中转"], Some(0.25), "空出来之后回到轮里"); + + // 按快慢、成败分：有效权重是权重 × 系数放大一千倍取整，份额按它算 + let (_d, app, gw) = app_and_gateway(&weighted.replace( + " type: load-balance\n", + " type: load-balance\n balance_by: latency\n", + )); + for _ in 0..3 { + gw.latency.record("官方", 400); + gw.latency.record("中转", 100); + } + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + // 中位数 250：官方 3 × (250/400)² → 1172，中转 1 × (250/100)² → 6250 + let s = shares(&r); + let (a, b) = (s["官方"].unwrap(), s["中转"].unwrap()); + assert!((a - 1172.0 / 7422.0).abs() < 1e-12, "{r:?}"); + assert!((b - 6250.0 / 7422.0).abs() < 1e-12, "{r:?}"); + assert!((a + b - 1.0).abs() < 1e-12); + + // 不经过 `load-balance` 的没有份额：指定的上游、按顺序的组 + let r = run(&app, r#"{"model":"claude-sonnet-4-5","cache":true}"#).await; + assert_eq!(shares(&r)["官方"], None); + for kind in ["fallback", "url-test", "cheapest"] { + let (_d, app) = app_with(&CFG.replace("type: load-balance", &format!("type: {kind}"))); + let r = run(&app, r#"{"model":"claude-sonnet-4-5"}"#).await; + assert!( + r.candidate_models.iter().all(|c| c.share.is_none()), + "{kind}: {r:?}" + ); + } +} + /// 按快慢、成败分时，试算说的排头就是数据面下一个新对话真的去的那一家：同一份延迟表、 /// 成功率、轮询状态，同一个有效权重。起真网关、真上游，每发一个新对话之前先试算一次， /// 一次都不能对不上 —— 真请求会添新的首字节样本和成败，系数一直在变，对得上才说明两边 diff --git a/crates/tw-control/tests/error_answers.rs b/crates/tw-control/tests/error_answers.rs index 35697a55..73248ccb 100644 --- a/crates/tw-control/tests/error_answers.rs +++ b/crates/tw-control/tests/error_answers.rs @@ -131,7 +131,7 @@ async fn world() -> World { tw_store::Blobs::new(blobs.clone()), tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), ), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let state = ControlState { diff --git a/crates/tw-control/tests/key_limits.rs b/crates/tw-control/tests/key_limits.rs index 4d2032b4..f518e93c 100644 --- a/crates/tw-control/tests/key_limits.rs +++ b/crates/tw-control/tests/key_limits.rs @@ -587,3 +587,49 @@ async fn a_time_zone_change_adds_the_new_day_up_again_from_the_records() { } assert_eq!(used(&b), 4); } + +/// 同上，只是**发现这一期变了的是一次结算**，而结算在存储层自己的记录线程上 —— 那条线程 +/// 不在异步运行时里。以前读回新一期的任务要在「当前的运行时」上起，在那儿起不来，这一天 +/// 就从 0 起了 +#[tokio::test] +async fn a_period_change_found_while_settling_on_the_recorder_thread_is_read_back() { + let clock = Arc::new(Moving { + east: tw_gateway::key_limits::TestClock::new(ms(NOON), 8 * 3600), + west: tw_gateway::key_limits::TestClock::new(ms(NOON), -5 * 3600), + moved: Default::default(), + }); + let b = bed_with(None, clock.clone()); + tw_control::key_limits::follow(&b.gw, b.rec.clone()); + for e in [ + request(1, "plain", ms("2026-10-04T23:00:00+08:00"), 10), + request(2, "plain", ms("2026-10-05T09:00:00+08:00"), 10), + ] + .iter() + .flatten() + { + b.rec.lock().await.on_event(e); + } + let limits: Vec = + serde_yaml_ng::from_str("[{per: day, requests: 100}]").unwrap(); + let used = |b: &Bed| b.gw.key_limits.view("plain", &limits)[0].used; + assert_eq!(used(&b), 1, "东八区的今天"); + clock.moved.store(true, std::sync::atomic::Ordering::SeqCst); + // 换了时区之后第一件事是结算一个请求，在一条不属于运行时的线程上 + let rec = b.rec.clone(); + std::thread::spawn(move || { + for e in request(3, "plain", ms(NOON), 10) { + rec.blocking_lock().on_event(&e); + } + }) + .join() + .unwrap(); + let mut seen = 0; + for _ in 0..100 { + seen = used(&b); + if seen == 3 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + assert_eq!(seen, 3, "西五区的今天：三个请求都在里面"); +} diff --git a/crates/tw-control/tests/route_stats.rs b/crates/tw-control/tests/route_stats.rs index 1f84e97f..e373bb0f 100644 --- a/crates/tw-control/tests/route_stats.rs +++ b/crates/tw-control/tests/route_stats.rs @@ -87,7 +87,7 @@ async fn every_route_and_rule_counts_what_went_through_it_including_denials() { std::fs::write(&p, &text).unwrap(); let cfg: tw_config::Config = serde_yaml_ng::from_str(&text).unwrap(); let gw = tw_gateway::AppState::new(cfg).unwrap(); - // 和 twcore 一样：存储层订阅总线，事件落库 + // 和 twcore 一样：存储层从总线给它的那条通道收事件，落库 let (_bodies, rx) = tokio::sync::mpsc::channel(8); let store = tw_store::task::spawn( tw_store::Recorder::new( @@ -95,7 +95,7 @@ async fn every_route_and_rule_counts_what_went_through_it_including_denials() { tw_store::Blobs::new(d.path().join("blobs")), tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), ), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let state = ControlState { @@ -187,7 +187,7 @@ async fn a_store_with_no_history_yet_says_nothing_is_covered() { tw_store::Blobs::new(d.path().join("blobs")), tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), ), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let state = ControlState { diff --git a/crates/tw-control/tests/sessions.rs b/crates/tw-control/tests/sessions.rs index 5c47be1c..084dc6c9 100644 --- a/crates/tw-control/tests/sessions.rs +++ b/crates/tw-control/tests/sessions.rs @@ -194,3 +194,34 @@ async fn the_endpoints_take_a_time_window_off_the_query_string() { let none = get(&app, "/sessions?from_ms=9000&to_ms=9999").await; assert!(none.as_array().unwrap().is_empty(), "{none}"); } + +/// **会话详情按会话号取，多老都取得到。**以前它取最近五百次会话再从里面找：更早的会话 +/// 在列表里点得到（翻到那么远的话），点开却是 404 +#[tokio::test] +async fn a_session_behind_the_newest_five_hundred_still_opens() { + let mut rows = vec![turn(1, "per-token"), turn(2, "per-token")]; + for i in 0..600 { + let mut r = turn(100 + i, "per-token"); + r.session = Some(format!("later-{i}")); + rows.push(r); + } + let (_d, app) = app(&rows); + + let detail = get(&app, "/sessions/s1").await; + assert_eq!(detail["session"]["id"], "s1"); + assert_eq!(detail["session"]["turns"], 2, "{detail}"); + assert_eq!(detail["turns"].as_array().unwrap().len(), 2); + + // 没有的照样是 404 + let r = app + .clone() + .oneshot( + Request::builder() + .uri("/sessions/nope") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(r.status(), StatusCode::NOT_FOUND); +} diff --git a/crates/tw-control/tests/stored_bodies.rs b/crates/tw-control/tests/stored_bodies.rs index c1861802..d561ffa4 100644 --- a/crates/tw-control/tests/stored_bodies.rs +++ b/crates/tw-control/tests/stored_bodies.rs @@ -125,7 +125,7 @@ async fn world(mode: &str) -> World { tw_store::Blobs::new(blobs.clone()), tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), ), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let state = ControlState { @@ -212,8 +212,15 @@ fn files(root: &std::path::Path) -> Vec<(String, String)> { .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)); + let name = f.file_name().to_string_lossy().to_string(); + // 正文先写进 `.tmp` 再改名：还在写的那份不算，列目录之后才改名走的也跳过 + if name.ends_with(".tmp") { + continue; + } + let Ok(bytes) = std::fs::read(f.path()) else { + continue; + }; + out.push((name, String::from_utf8_lossy(&bytes).to_string())); } } out diff --git a/crates/tw-control/tests/transcript.rs b/crates/tw-control/tests/transcript.rs index 8a11dde4..baeaec18 100644 --- a/crates/tw-control/tests/transcript.rs +++ b/crates/tw-control/tests/transcript.rs @@ -148,6 +148,9 @@ async fn the_transcript_is_the_shape_the_desktop_app_reads() { json!({ "session": "s1", "system": "sys", + "total_turns": 2, + // 第二轮没有回答（不会再被改写），两轮都早就结束了 + "settled_turns": 2, "turns": [ { "id": "1", @@ -220,3 +223,43 @@ async fn a_session_whose_bodies_are_gone_still_lists_its_turns() { assert_eq!(x["input"], json!([])); } } + +/// `from_turn`：只要从那一轮起的那些，形状不变，`total_turns` 说一共几轮。比总轮数还大的 +/// 是空的。不给和给 0 一样，是整段 +#[tokio::test] +async fn from_turn_gives_the_turns_from_there_on() { + let first = + json!({"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "一"}]}); + let answer = json!({"type": "message", "content": [{"type": "text", "text": "回一"}]}); + let second = json!({"model": "claude-sonnet-4-5", "messages": [ + {"role": "user", "content": "一"}, + {"role": "assistant", "content": "回一"}, + {"role": "user", "content": "二"} + ]}); + let (_d, app) = app( + &[turn(1), turn(2)], + &[(1, Some(first), Some(answer)), (2, Some(second), None)], + ); + + let (st, all) = get(&app, "/sessions/s1/transcript").await; + assert_eq!(st, StatusCode::OK, "{all}"); + let (_, zero) = get(&app, "/sessions/s1/transcript?from_turn=0").await; + assert_eq!(zero, all); + + let (st, tail) = get(&app, "/sessions/s1/transcript?from_turn=1").await; + assert_eq!(st, StatusCode::OK, "{tail}"); + assert_eq!(tail["total_turns"], 2); + assert_eq!(tail["settled_turns"], all["settled_turns"]); + assert_eq!(tail["system"], all["system"]); + assert_eq!(tail["turns"], json!([all["turns"][1]])); + assert_eq!(tail["turns"][0]["id"], "2"); + + let (st, none) = get(&app, "/sessions/s1/transcript?from_turn=9").await; + assert_eq!(st, StatusCode::OK); + assert_eq!(none["total_turns"], 2); + assert_eq!(none["turns"], json!([])); + + // 不是数的 `from_turn` 是请求写错了 + let (st, body) = get(&app, "/sessions/s1/transcript?from_turn=x").await; + assert_eq!(st, StatusCode::BAD_REQUEST, "{body}"); +} diff --git a/crates/tw-control/tests/ws_turns.rs b/crates/tw-control/tests/ws_turns.rs index 7553b31a..7f63bfe1 100644 --- a/crates/tw-control/tests/ws_turns.rs +++ b/crates/tw-control/tests/ws_turns.rs @@ -79,7 +79,7 @@ providers: gw.pricing.clone(), ) .settling_to(tw_control::key_limits::settle_hook(&gw)), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let addr = tw_gateway::serve(gw.clone(), ([127, 0, 0, 1], 0).into()) @@ -243,7 +243,7 @@ providers: gw.pricing.clone(), ) .settling_to(tw_control::key_limits::settle_hook(&gw)), - gw.bus.subscribe(), + gw.bus.record_feed().unwrap(), rx, ); let addr = tw_gateway::serve(gw.clone(), ([127, 0, 0, 1], 0).into()) diff --git a/crates/tw-dialect/Cargo.toml b/crates/tw-dialect/Cargo.toml index 7a3e72a2..06a6fd99 100644 --- a/crates/tw-dialect/Cargo.toml +++ b/crates/tw-dialect/Cargo.toml @@ -13,3 +13,4 @@ description = "Request, response and stream conversion between Anthropic Message [dependencies] serde = { workspace = true } serde_json = { workspace = true } +memchr = { workspace = true } diff --git a/crates/tw-dialect/src/convert.rs b/crates/tw-dialect/src/convert.rs index d0775178..793b36e1 100644 --- a/crates/tw-dialect/src/convert.rs +++ b/crates/tw-dialect/src/convert.rs @@ -852,9 +852,8 @@ impl Collector { /// `tw1.` 前缀的签名不是这个上游签发的：Anthropic 会以签名无效拒绝整个请求， /// 对话从此卡死。所以直通前把它们去掉 —— 那段推理本来就不是这个上游产生的。 pub fn strip_carried(client: Dialect, body: &[u8]) -> Option> { - if !body.windows(CARRIED.len()).any(|w| w == CARRIED.as_bytes()) { - return None; - } + // 每一跳同格式直通都要看一眼整个请求体：按字节找（SIMD），绝大多数请求在这里就回去了 + memchr::memmem::find(body, CARRIED.as_bytes())?; let mut v: Value = serde_json::from_slice(body).ok()?; let carried = |x: Option<&Value>| { x.and_then(Value::as_str) diff --git a/crates/tw-dialect/src/usage.rs b/crates/tw-dialect/src/usage.rs index 248fd7a5..dab624a0 100644 --- a/crates/tw-dialect/src/usage.rs +++ b/crates/tw-dialect/src/usage.rs @@ -492,10 +492,7 @@ fn find_at(hay: &[u8], needle: &[u8], from: usize) -> Option { if from >= hay.len() || needle.len() > hay.len() - from { return None; } - hay[from..] - .windows(needle.len()) - .position(|w| w == needle) - .map(|i| i + from) + memchr::memmem::find(&hay[from..], needle).map(|i| i + from) } /// 从 `open`（一个 `{`）开始配对，返回对应 `}` 的下标。 diff --git a/crates/tw-engine/src/weighted.rs b/crates/tw-engine/src/weighted.rs index 951c5922..f50c6fb0 100644 --- a/crates/tw-engine/src/weighted.rs +++ b/crates/tw-engine/src/weighted.rs @@ -238,6 +238,28 @@ pub fn advance(g: &Group, members: &[String], f: &Facts, leader: &str) -> HashMa out } +/// 长期看每一家排头的份额，0 到 1，和 `members` 一一对应：**这一轮参加的**（[`round`]， +/// 停着的、满着的不算，全不能参加时都算）按各自的有效权重（[`effective`]）分，加起来是 1； +/// 不参加的是 `None`。 +/// +/// 平滑加权轮询每排有效权重之和那么多次，各家正好各排头自己的有效权重那么多次，所以份额 +/// 就是有效权重 ÷ 这一轮的总和。和 [`lead`] 读的是同一份 [`round`]：试算说的份额，就是 +/// 数据面照此刻的数字轮下去各家分到的新对话的比例（进行中的对话留在原来那一家，那一份 +/// 由之后的新对话补回来，见 [`advance`]）。 +pub fn shares(g: &Group, members: &[String], f: &Facts) -> Vec> { + let round = round(g, members, f); + let total: i64 = round.iter().map(|(_, w)| w).sum(); + members + .iter() + .map(|m| { + round + .iter() + .find(|(r, _)| *r == m.as_str()) + .map(|(_, w)| *w as f64 / total as f64) + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; @@ -563,6 +585,46 @@ mod tests { assert_eq!((got["慢"], got["快"]), (1172, 1563)); } + /// 份额就是一整圈里各家排头的比例：轮一整圈数出来的次数 ÷ 圈长，和 [`shares`] 给的一样。 + /// 停着的、满着的不参加，份额是空的，剩下的加起来还是 1；全都不能参加时都算 + #[test] + fn the_share_is_how_often_each_member_leads_over_a_full_round() { + let g = balanced(&[("慢", 3), ("快", 1), ("停", 2)], BalanceBy::Latency); + let members = names(&g); + let mut f = Facts { + ttfb_ms: ttfb(&[("慢", 400), ("快", 100)]), + paused: ["停".to_string()].into_iter().collect(), + ..Default::default() + }; + let got = shares(&g, &members, &f); + assert_eq!(got[2], None, "停着的不参加"); + // 中位数 250：慢 3 × 0.390625 → 1172，快 1 × 6.25 → 6250 + let n = 1172 + 6250; + let led = count(&g, &mut f, n); + for (i, m) in members.iter().take(2).enumerate() { + let want = led[m] as f64 / n as f64; + assert!((got[i].unwrap() - want).abs() < 1e-12, "{m}: {got:?}"); + } + assert!((got[0].unwrap() + got[1].unwrap() - 1.0).abs() < 1e-12); + + // 满着的同样不参加 + let f = Facts { + busy: ["快".to_string()].into_iter().collect(), + ..Default::default() + }; + assert_eq!(shares(&g, &members, &f), [Some(0.6), None, Some(0.4)]); + // 全都停着、满着：都参加，按权重分 + let f = Facts { + paused: ["慢".to_string(), "停".to_string()].into_iter().collect(), + busy: ["快".to_string()].into_iter().collect(), + ..Default::default() + }; + assert_eq!( + shares(&g, &members, &f), + [Some(0.5), Some(1.0 / 6.0), Some(1.0 / 3.0)] + ); + } + /// 各家系数一样时，次序和只按权重分一模一样：7:3 照样穿插着来，不是先七后三 #[test] fn equal_factors_keep_the_exact_interleaving() { diff --git a/crates/tw-gateway/Cargo.toml b/crates/tw-gateway/Cargo.toml index b2859387..3ddbe4c1 100644 --- a/crates/tw-gateway/Cargo.toml +++ b/crates/tw-gateway/Cargo.toml @@ -39,6 +39,7 @@ tokio-rustls = { workspace = true } base64 = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +memchr = { workspace = true } blake3 = { workspace = true } sha2 = { workspace = true } rand = { workspace = true } diff --git a/crates/tw-gateway/src/bodies.rs b/crates/tw-gateway/src/bodies.rs index 1c83f322..df5c11aa 100644 --- a/crates/tw-gateway/src/bodies.rs +++ b/crates/tw-gateway/src/bodies.rs @@ -25,7 +25,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use bytes::Bytes; use tw_guard::redact::replace::{Ledger, Scheme}; -use tw_guard::redact::rules::{Kind, Rule, RuleSet}; +use tw_guard::redact::rules::{Hit, Kind, Rule, RuleSet}; /// 打码要多看的那一截。 /// @@ -101,12 +101,25 @@ impl Redaction { /// /// **读的时候还会再打一遍**，打第二遍不再改动什么（`mask_body` 认得自己打过的码）。 pub fn apply(&self, text: &str) -> String { + self.apply_found(text, None) + } + + /// [`apply`](Self::apply)，`found` 是按这套规则在 `text` 上已经找过的命中（[`crate::guard::hits`] + /// 的结果）：不再找一遍。没找过的是 `None` + fn apply_found(&self, text: &str, found: Option<&[Hit]>) -> String { // 我们自己的占位符不算（见 `crate::guard::hits`）：回答里回显的、重放过来的，原样存 - let hits = crate::guard::hits(text, &self.rules); + let fresh; + let hits = match found { + Some(hits) => hits, + None => { + fresh = crate::guard::hits(text, &self.rules); + &fresh[..] + } + }; 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 { + for h in hits { out.push_str(&text[at..h.bytes.start]); let value = &text[h.bytes.clone()]; match sent.get(value) { @@ -148,6 +161,8 @@ pub struct BodyRecord { pub original_len: usize, /// 落盘之前怎么换、怎么打码 pub redaction: Redaction, + /// 请求路上按 `redaction` 的规则在这份正文上已经找过的命中（见 [`Self::found`]） + found: Option>, /// 占着的那份额度（见 [`QUEUED_MAX`]） held: Option, } @@ -168,10 +183,21 @@ impl BodyRecord { body, original_len, redaction, + found: None, held: None, } } + /// 请求路上出站脱敏已经按 `redaction` 的规则在 `body` 上找过一遍（见 + /// [`crate::guard::look_hits`]）：落盘前打码用它找到的，不再把整份正文扫一遍。 + /// + /// 只在交来的是整份正文、它又是 UTF-8 时用得上 —— 那时落盘前看的文字就是请求路上看的 + /// 那一份；截过的、不是 UTF-8 的照旧再找（[`Self::for_disk`]）。 + pub fn found(mut self, hits: Option>) -> Self { + self.found = hits; + self + } + /// 落盘的那一份：换过、打过码（[`Redaction::apply`]），和要记下的原本长度。 /// /// **在阻塞线程上调**：一份 4 MB 的正文要扫好几遍。 @@ -180,7 +206,13 @@ impl BodyRecord { /// 那样，而二进制的正文里没有能看的东西。 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)); + let text = String::from_utf8_lossy(window); + // 请求路上找过的那一遍看的就是这些字：整份都在窗口里，原文就是 UTF-8 + let entire = window.len() == self.body.len() && self.original_len <= self.body.len(); + let found = self + .found + .filter(|_| entire && matches!(text, std::borrow::Cow::Borrowed(_))); + let body = self.redaction.apply_found(&text, found.as_deref()); // 截过的（交来的只是开头）报原本的长度。没截过的就是换过、打过码的这一份的长度： // 比存储层的上限还长的，由存储层截、由它记下（`tw_store::Blobs::put_with_len`） let whole = self.original_len.max(self.body.len()); @@ -266,6 +298,8 @@ 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]); + // 找过的是整份；落盘前看的只是开头这一截，要重新找（见 `BodyRecord::found`） + rec.found = None; } let n = rec.body.len(); // 先占额度，占不下就丢 @@ -461,6 +495,44 @@ mod tests { serde_json::from_str::(&stored).expect("存下来的还是 JSON"); } + /// 请求路上找过的命中拿来打码，落盘的和自己再找一遍的一字不差；交来的只是开头一截时 + /// 不用它（它说的是整份），照旧自己找 + #[test] + fn the_hits_found_on_the_way_in_are_used_only_for_the_whole_body() { + let body = format!( + r#"{{"messages":[{{"role":"user","content":"key {KEY} 库 postgres://app:hunter2hunter2@db/x"}}]}}"# + ); + for mode in [ + tw_config::SecurityMode::Observe, + tw_config::SecurityMode::Enforce, + ] { + let rules = RuleSet::defaults(); + let seen = crate::guard::look_hits(mode, &rules, body.as_bytes()); + assert_eq!(seen.hits.as_ref().map(Vec::len), Some(2)); + let r = Redaction { + rules: Arc::new(rules), + ledger: seen.ledger, + }; + let again = written(record(BodyKind::Request, &body, r.clone())); + let reused = + written(record(BodyKind::Request, &body, r).found(seen.hits.map(Arc::from))); + assert_eq!(reused, again); + assert!(!reused.contains(KEY), "{reused}"); + } + // 说「什么都没找到」的命中：整份交来时照它（证明真的用上了），截过的不照它 + let lie = || Some(Arc::from(Vec::new())); + let whole = written(record(BodyKind::Request, &body, Redaction::default()).found(lie())); + assert!(whole.contains("sk-an…AAAA"), "{whole}"); + assert!( + whole.contains("hunter2hunter2"), + "规则那一道真的没再找：{whole}" + ); + let mut cut = record(BodyKind::Request, &body, Redaction::default()).found(lie()); + cut.original_len = body.len() + 1; + let cut = written(cut); + assert!(!cut.contains("hunter2hunter2"), "{cut}"); + } + #[test] fn what_the_rules_miss_is_masked_by_its_shape() { // 网关自己的钥匙（`tw-`）没有一条脱敏规则认：读的时候那一道兜住它，写的时候也一样 diff --git a/crates/tw-gateway/src/client_api.rs b/crates/tw-gateway/src/client_api.rs index 3c1b5229..8b72581a 100644 --- a/crates/tw-gateway/src/client_api.rs +++ b/crates/tw-gateway/src/client_api.rs @@ -224,6 +224,9 @@ pub struct Reading { pub facts: tw_engine::RequestFacts, /// DeepSeek Harness 发的请求（见 [`tw_dialect::harness`]）。要看请求头，由管线填 pub harness: Option, + /// 请求体解得开（是 JSON）。之后每一跳按字节改它的地方（见 [`crate::egress`]）凭这个 + /// 知道不用再解一遍 + pub json: bool, } pub fn read(path: &str, query: Option<&str>, body: Option<&serde_json::Value>) -> Reading { @@ -259,6 +262,7 @@ pub fn read(path: &str, query: Option<&str>, body: Option<&serde_json::Value>) - decoded, facts, harness: None, + json: body.is_some(), } } diff --git a/crates/tw-gateway/src/clientprobe.rs b/crates/tw-gateway/src/clientprobe.rs index ea981f20..db8ba3a0 100644 --- a/crates/tw-gateway/src/clientprobe.rs +++ b/crates/tw-gateway/src/clientprobe.rs @@ -54,15 +54,16 @@ const TOPIC_MARK: &str = "Analyze if this message indicates a new conversation t const SUGGEST_MARK: &str = "[SUGGESTION MODE:"; const WARMUP_BODY: &str = "Warmup"; -/// 认一认这是不是客户端自己的辅助请求。 +/// 认一认这是不是客户端自己的辅助请求。`parsed` 是 `body` 解析出来的 JSON（解不开是 +/// `None`）：管线为路由本来就要解析一遍，这里不再解一遍。 /// -/// **先做整包字节扫描快速排除**，不命中就直接返回，连 JSON 都不解析。 +/// **先做整包字节扫描快速排除**，不命中就直接返回，连解析出来的 JSON 都不看。 /// 绝大多数请求会在这一步出去 —— 而这段代码跑在每一个请求的关键路径上。 /// /// `is_claude_code` 由调用方判断。`max_tokens: 1` 那条**必须**同时要求 /// 它，否则会误伤别人真实的 `max_tokens: 1` 请求。 -pub fn classify(body: &Bytes, is_claude_code: bool) -> Option { - let has = |m: &str| memfind(body, m.as_bytes()); +pub fn classify(body: &[u8], parsed: Option<&Value>, is_claude_code: bool) -> Option { + let has = |m: &str| memchr::memmem::find(body, m.as_bytes()).is_some(); // 快速排除。四个 B 类标记各有一句独特的原文；A 类里 warmup 也有， // 而 health check 只能靠 `max_tokens` 这个字段名先粗筛。 @@ -75,8 +76,7 @@ pub fn classify(body: &Bytes, is_claude_code: bool) -> Option { return None; } - let v: Value = serde_json::from_slice(body).ok()?; - let obj = v.as_object()?; + let obj = parsed?.as_object()?; // B 类的三条：标记出现在 system 或者第一个 user 消息里就算。 // **先判 B 类** —— 一个带标记的请求即使 max_tokens 是 1，它也是 @@ -107,14 +107,6 @@ pub fn classify(body: &Bytes, is_claude_code: bool) -> Option { None } -/// 朴素子串查找。请求体是几十 KB 的量级，不值得为它引一个 SIMD 库。 -fn memfind(hay: &[u8], needle: &[u8]) -> bool { - if needle.is_empty() || hay.len() < needle.len() { - return false; - } - hay.windows(needle.len()).any(|w| w == needle) -} - /// 把 system 和所有消息里的文本摊平了看。 fn all_text(obj: &serde_json::Map) -> String { let mut out = String::new(); @@ -339,6 +331,12 @@ mod tests { Bytes::from(s.to_string()) } + /// 和管线一样：先解析一遍，再认 + fn classify(body: &Bytes, is_claude_code: bool) -> Option { + let parsed = serde_json::from_slice::(body).ok(); + super::classify(body, parsed.as_ref(), is_claude_code) + } + // ── 该认出来的 ────────────────────────────────────────────────── #[test] fn a_health_check_is_recognised() { @@ -449,10 +447,14 @@ mod tests { #[test] fn an_ordinary_request_does_not_even_get_parsed() { - // 快速排除跑在每个请求的关键路径上。这条测的是它确实先排除了。 + // 快速排除跑在每个请求的关键路径上。这条测的是它确实先排除了：原文里一个标记都 + // 没有，解析出来的 JSON 就不看（这里故意交一份「是预热」的进去） let body = b(r#"{"model":"claude-sonnet-4-5","max_tokens":8000, "messages":[{"role":"user","content":"帮我看看这段代码"}]}"#); assert_eq!(classify(&body, true), None); + let warmup: Value = + serde_json::json!({"messages": [{"role": "user", "content": "Warmup"}]}); + assert_eq!(super::classify(&body, Some(&warmup), true), None); } #[test] diff --git a/crates/tw-gateway/src/egress.rs b/crates/tw-gateway/src/egress.rs index be407ae9..c294ebc0 100644 --- a/crates/tw-gateway/src/egress.rs +++ b/crates/tw-gateway/src/egress.rs @@ -150,19 +150,41 @@ pub fn gateway_headers(hop: &Hop, client: &HeaderMap) -> Vec<(String, String)> { /// （[`crate::chatgpt::IDENTITY_METADATA`]）；会话、线程这些照发 /// /// 返回去掉之后的请求体；没有可去的是 `None`，请求体一个字节都不动。这不是为了兼容去掉 -/// 请求的内容，所以不算丢弃的字段 +/// 请求的内容，所以不算丢弃的字段。 +/// +/// **只剪掉那几个成员，别的字节一个不动**（见 [`cut_members`]）：解析再写回去的话，没开 +/// `preserve_order` 的 serde_json 会把每个对象的键按字母重排 —— 工具定义、工具参数、整个 +/// 请求都变了样（上游的提示缓存按字节认），而 Claude Code 的请求每一个都带着 `user_id`。 +/// 剪不了的怪样子（同一个键写了两遍）才照旧解析、去掉、写回。 +/// +/// `body` 要是解得开的 JSON：调用方解过它（见管线的 `read`），这里不再解一遍 pub fn strip_body_identity(dialect: Dialect, body: &Bytes) -> Option { - // 大多数请求没有这些字段：先便宜地看一眼，免得每个请求都解析、重写一遍请求体 - let text = std::str::from_utf8(body).ok()?; + // 大多数请求没有这些字段：先便宜地看一眼，免得每个请求都把整个请求体走一遍 + let (parent, keys, drop_empty): (&str, &[&str], bool) = match dialect { + Dialect::Anthropic => ("metadata", &["user_id"], true), + Dialect::Responses => ("client_metadata", &crate::chatgpt::IDENTITY_METADATA, false), + _ => return None, + }; let marks: &[&str] = match dialect { Dialect::Anthropic => &["\"user_id\""], - Dialect::Responses => &crate::chatgpt::IDENTITY_METADATA, - _ => return None, + _ => keys, }; - if !marks.iter().any(|m| text.contains(m)) { + if !marks + .iter() + .any(|m| memchr::memmem::find(body, m.as_bytes()).is_some()) + { return None; } - let mut v: Value = serde_json::from_str(text).ok()?; + match cut_members(body, parent, keys, drop_empty) { + Cut::Done(out) => Some(Bytes::from(out)), + Cut::Nothing => None, + Cut::Unusual => reserialize(dialect, body), + } +} + +/// 解析、去掉、写回：[`strip_body_identity`] 剪不了的时候走这条（键会按字母重排）。 +fn reserialize(dialect: Dialect, body: &Bytes) -> Option { + let mut v: Value = serde_json::from_slice(body).ok()?; let obj = v.as_object_mut()?; let removed = match dialect { Dialect::Anthropic => { @@ -186,6 +208,74 @@ pub fn strip_body_identity(dialect: Dialect, body: &Bytes) -> Option { serde_json::to_vec(&v).ok().map(Bytes::from) } +/// [`cut_members`] 的结果。 +#[derive(Debug, PartialEq)] +enum Cut { + /// 剪过的请求体 + Done(Vec), + /// 没有要剪的 + Nothing, + /// 同一个键写了两遍，或者不像 JSON 对象：剪的结果说不准和解析出来的一样，交给 + /// [`reserialize`] + Unusual, +} + +/// 顶层对象里 `parent` 那个对象，剪掉它的 `keys` 成员（连同分隔它的逗号）；剪空了而 +/// `drop_empty` 的话，`parent` 整个剪掉。别的字节原样留着。 +/// +/// 只走顶层和 `parent` 这两层：别的值整个跳过，字符串里按字节找下一个引号或反斜杠。键里 +/// 写了转义的（`"user\u005fid"`）照解析出来的样子比，和 serde_json 认的一样。 +fn cut_members(body: &[u8], parent: &str, keys: &[&str], drop_empty: bool) -> Cut { + let Some(top) = crate::splice::object(body, crate::splice::ws(body, 0)) else { + return Cut::Unusual; + }; + let mut parents = top.iter().filter(|m| m.is(body, parent)); + let (Some(p), None) = (parents.next(), parents.next()) else { + // 一个都没有就没得剪；有两个的话 serde_json 留后一个，交给它 + return if top.iter().any(|m| m.is(body, parent)) { + Cut::Unusual + } else { + Cut::Nothing + }; + }; + if body[p.value.start] != b'{' { + return Cut::Nothing; + } + let Some(inner) = crate::splice::object(body, p.value.start) else { + return Cut::Unusual; + }; + let gone: Vec = (0..inner.len()) + .filter(|&i| keys.iter().any(|k| inner[i].is(body, k))) + .collect(); + if gone.is_empty() { + return Cut::Nothing; + } + // 同一个键写了两遍 + if keys + .iter() + .any(|k| gone.iter().filter(|&&i| inner[i].is(body, k)).count() > 1) + { + return Cut::Unusual; + } + let cuts = if drop_empty && gone.len() == inner.len() { + let at = top + .iter() + .position(|m| m.key.start == p.key.start) + .unwrap_or_default(); + crate::splice::cuts(&top, &[at]) + } else { + crate::splice::cuts(&inner, &gone) + }; + let mut out = Vec::with_capacity(body.len()); + let mut at = 0; + for c in cuts { + out.extend_from_slice(&body[at..c.start]); + at = c.end; + } + out.extend_from_slice(&body[at..]); + Cut::Done(out) +} + #[cfg(test)] mod tests { use super::*; @@ -437,4 +527,220 @@ mod tests { ); assert_eq!(strip_body_identity(Dialect::Chat, &body), None); } + + /// 去掉 `user_id` 之后，请求体别的字节一个不动：键的先后（工具定义、工具参数、整个请求）、 + /// 空白、转义都照客户端写的。以前解析再写回，每个对象的键都按字母重排了 + #[test] + fn the_forwarded_body_is_the_client_body_less_the_identity() { + let tools = r#""tools":[{"name":"Read","description":"读文件 \"quoted\" \\ path","input_schema":{"type":"object","properties":{"file_path":{"type":"string"},"offset":{"type":"number"},"limit":{"type":"number"}},"required":["file_path"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}}]"#; + let messages = r#""messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_01","content":"{\"user_id\":\"not this one\"} }]"}]},{"role":"assistant","content":[{"type":"tool_use","id":"toolu_02","name":"Read","input":{"zeta":1,"alpha":[2,{"b":3,"a":4}]}}]}]"#; + let identity = r#""metadata":{"user_id":"user_abc_account_123_session_456"}"#; + for (before, after) in [ + // Claude Code 的顺序：metadata 在工具定义后面 + ( + format!( + r#"{{"model":"claude-sonnet-5",{messages},"system":"s",{tools},{identity},"max_tokens":32000,"stream":true}}"# + ), + format!( + r#"{{"model":"claude-sonnet-5",{messages},"system":"s",{tools},"max_tokens":32000,"stream":true}}"# + ), + ), + // 排在最前、最后 + ( + format!(r#"{{{identity},{messages}}}"#), + format!(r#"{{{messages}}}"#), + ), + ( + format!(r#"{{{messages},{identity}}}"#), + format!(r#"{{{messages}}}"#), + ), + // 带空白、换行的 + ( + format!("{{\n {messages},\n {identity} ,\n \"stream\": true\n}}"), + format!("{{\n {messages},\n \"stream\": true\n}}"), + ), + // metadata 里还有别的：只剪 user_id + ( + r#"{"metadata":{"user_id":"u","note":"x"},"messages":[]}"#.to_string(), + r#"{"metadata":{"note":"x"},"messages":[]}"#.to_string(), + ), + ( + r#"{"metadata": {"note": "x", "user_id": "u"}, "messages": []}"#.to_string(), + r#"{"metadata": {"note": "x"}, "messages": []}"#.to_string(), + ), + ( + r#"{"metadata":{"a":1,"user_id":{"nested":["}"]},"b":2}}"#.to_string(), + r#"{"metadata":{"a":1,"b":2}}"#.to_string(), + ), + // 键写成转义的：serde_json 认它是 user_id，这里也认（原文里别处得写着 `"user_id"`， + // 先看的那一眼才放它过去，和以前一样）。反斜杠由 `char::from(92)` 拼：测试里直接 + // 写出来的转义，经过某些编辑工具会变成真字符 + ( + format!( + r#"{{"metadata":{{"user{b}u005fid":"u"}},"user_id":1}}"#, + b = char::from(92) + ), + r#"{"user_id":1}"#.to_string(), + ), + ] { + let out = strip_body_identity(Dialect::Anthropic, &Bytes::from(before.clone())) + .unwrap_or_else(|| panic!("nothing stripped from {before}")); + assert_eq!(std::str::from_utf8(&out).unwrap(), after); + // 和解析、去掉、写回的意思一样 + assert_eq!( + serde_json::from_slice::(&out).unwrap(), + serde_json::from_slice::( + &reserialize(Dialect::Anthropic, &Bytes::from(before)).unwrap() + ) + .unwrap() + ); + } + + let codex = r#"{"model":"gpt-5.5","input":[],"client_metadata":{"x-codex-turn-metadata":"{}","session_id":"conv-1","x-codex-installation-id":"inst-1"},"stream":true}"#; + let out = strip_body_identity(Dialect::Responses, &Bytes::from(codex)).unwrap(); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + r#"{"model":"gpt-5.5","input":[],"client_metadata":{"session_id":"conv-1"},"stream":true}"# + ); + // 剪空了也留着这个对象，和以前一样 + let codex = r#"{"client_metadata":{"x-codex-installation-id":"inst-1","x-codex-turn-metadata":"{}"}}"#; + let out = strip_body_identity(Dialect::Responses, &Bytes::from(codex)).unwrap(); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + r#"{"client_metadata":{}}"# + ); + } + + /// 同一个键写了两遍：剪的结果说不准和 serde_json 认的一样（它留后一个），照旧解析、 + /// 去掉、写回 + #[test] + fn a_key_written_twice_falls_back_to_reserializing() { + for body in [ + r#"{"metadata":{"user_id":"a"},"x":1,"metadata":{"user_id":"b","n":1}}"#, + r#"{"metadata":{"user_id":"a","user_id":"b"},"x":1}"#, + ] { + let body = Bytes::from(body); + assert_eq!( + strip_body_identity(Dialect::Anthropic, &body), + reserialize(Dialect::Anthropic, &body) + ); + } + assert_eq!( + cut_members( + br#"{"metadata":1,"metadata":{"user_id":"b"}}"#, + "metadata", + &["user_id"], + true + ), + Cut::Unusual + ); + } + + /// 随机拼出来的请求：剪出来的和解析、去掉、写回的意思一样，剪没剪也一样；剪过的那一份 + /// 只少了几段，别的字节按原样、原来的先后都在 + #[test] + fn cutting_means_what_reserializing_meant() { + struct Rng(u64); + impl Rng { + fn below(&mut self, n: usize) -> usize { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 % n as u64) as usize + } + fn pick<'a>(&mut self, xs: &[&'a str]) -> &'a str { + xs[self.below(xs.len())] + } + } + let mut r = Rng(0x51_7cc1_b727_220a); + const KEYS: &[&str] = &[ + "\"user_id\"", + "\"user\\u005fid\"", + "\"note\"", + "\"a\"", + "\"x-codex-installation-id\"", + "\"x-codex-turn-metadata\"", + "\"session_id\"", + ]; + const VALUES: &[&str] = &[ + "1", + "-2.5e3", + "true", + "null", + "\"s\"", + "\"q\\\"}]{[,\\\\\"", + "\"\\\"user_id\\\"\"", + "[]", + "{}", + "[1,{\"user_id\":\"deep\"},\"]\"]", + "{\"metadata\":{\"user_id\":\"deep\"}}", + ]; + const WS: &[&str] = &["", " ", "\n ", "\t"]; + let inner = |r: &mut Rng| { + let members: Vec = (0..r.below(4)) + .map(|_| { + let (k, a, b, v) = (r.pick(KEYS), r.pick(WS), r.pick(WS), r.pick(VALUES)); + format!("{k}{a}:{b}{v}") + }) + .collect(); + format!("{{{}}}", members.join(",")) + }; + for _ in 0..5_000 { + let mut top = Vec::new(); + for _ in 0..r.below(5) { + let (k, v) = match r.below(4) { + 0 => ("\"metadata\"", inner(&mut r)), + 1 => ("\"client_metadata\"", inner(&mut r)), + 2 => ("\"messages\"", r.pick(VALUES).to_string()), + _ => (r.pick(KEYS), r.pick(VALUES).to_string()), + }; + let (a, b, c, d) = (r.pick(WS), r.pick(WS), r.pick(WS), r.pick(WS)); + top.push(format!("{a}{k}{b}:{c}{v}{d}")); + } + let (a, b) = (r.pick(WS), r.pick(WS)); + let body = Bytes::from(format!("{a}{{{}}}{b}", top.join(","))); + if serde_json::from_slice::(&body).is_err() { + continue; + } + for d in [Dialect::Anthropic, Dialect::Responses] { + let cut = strip_body_identity(d, &body); + let old = reserialize(d, &body); + // 旧的写法先看一眼原文里有没有那几个字：没有就不解析 + let marked = match d { + Dialect::Anthropic => memchr::memmem::find(&body, b"\"user_id\"").is_some(), + _ => crate::chatgpt::IDENTITY_METADATA + .iter() + .any(|m| memchr::memmem::find(&body, m.as_bytes()).is_some()), + }; + let old = old.filter(|_| marked); + let value = |b: &Option| { + b.as_ref() + .map(|b| serde_json::from_slice::(b).unwrap()) + }; + assert_eq!( + value(&cut), + value(&old), + "{d:?} {}", + String::from_utf8_lossy(&body) + ); + // 剪过的那一份是原文的一个子序列：别的字节原样、按原来的先后 + if let Some(cut) = &cut + && cut_members(&body, "metadata", &["user_id"], true) != Cut::Unusual + && cut_members( + &body, + "client_metadata", + &crate::chatgpt::IDENTITY_METADATA, + false, + ) != Cut::Unusual + { + let mut rest = body.iter(); + assert!( + cut.iter().all(|c| rest.any(|b| b == c)), + "{}", + String::from_utf8_lossy(&body) + ); + } + } + } + } } diff --git a/crates/tw-gateway/src/forward.rs b/crates/tw-gateway/src/forward.rs index 1af97948..99cfe6e1 100644 --- a/crates/tw-gateway/src/forward.rs +++ b/crates/tw-gateway/src/forward.rs @@ -75,29 +75,58 @@ pub fn gemini_path_with_model(path: &str, model: &str) -> String { /// 字段名按客户端的格式写（[`tw_dialect::params`]）：同一个「最大输出」在四种格式里叫 /// 四个名字。认不出格式时按 Anthropic 写。Gemini 的模型在路径里，见 /// [`gemini_path_with_model`]。 +/// +/// **只动改了的那几个字段，别的字节一个不动**（见 [`crate::splice`]）：解析再写回的话， +/// 每个对象的键都按字母重排了，工具定义、工具参数一起变了样。改出来和原来一样的（规则 +/// 写的值客户端本来就是这么写的）原样发。改到的对象里同一个键写了两遍的，才照旧写回。 pub fn apply_set( body: &Bytes, set: &tw_engine::SetAction, client: Option, ) -> Bytes { - use tw_dialect::params; if set.is_empty() { return body.clone(); } - let Ok(mut v) = serde_json::from_slice::(body) else { + let Ok(before) = serde_json::from_slice::(body) else { // 解不开就别动。我们的解析器不认识的东西，上游可能完全认识。 tracing::warn!("the request body is not JSON; skipping the parameter rewrites"); return body.clone(); }; - if !v.is_object() { + if !before.is_object() { return body.clone(); } + let mut v = before.clone(); + set_params(&mut v, set, client); + match crate::splice::rewrite(body, &before, &v) { + crate::splice::Rewrite::Changed(b) => return Bytes::from(b), + crate::splice::Rewrite::Same => return body.clone(), + crate::splice::Rewrite::Unusual => {} + } + match serde_json::to_vec(&v) { + Ok(b) => Bytes::from(b), + // 序列化不该失败，但真失败了宁可发原文也不要发半个 body + Err(e) => { + tracing::error!( + "the rewritten request body could not be serialized; sending the original instead: {e}" + ); + body.clone() + } + } +} + +/// 把 `set` 点名的参数写进解析出来的请求体 `v`（见 [`apply_set`]）。 +fn set_params( + v: &mut serde_json::Value, + set: &tw_engine::SetAction, + client: Option, +) { + use tw_dialect::params; let client = client.unwrap_or(tw_dialect::ir::Dialect::Anthropic); if let Some(m) = &set.model { - params::set_model(client, &mut v, m); + params::set_model(client, v, m); } if let Some(t) = set.max_tokens { - params::set_max_output_tokens(client, &mut v, t); + params::set_max_output_tokens(client, v, t); } match set.thinking { // 开启思考需要一个 budget，而我们没有一个合理的值可以编。 @@ -105,19 +134,9 @@ pub fn apply_set( Some(true) => tracing::warn!( "set.thinking: true is not supported yet (it needs budget_tokens); ignoring it" ), - Some(false) => params::disable_thinking(client, &mut v), + Some(false) => params::disable_thinking(client, v), None => {} } - match serde_json::to_vec(&v) { - Ok(b) => Bytes::from(b), - // 序列化不该失败，但真失败了宁可发原文也不要发半个 body - Err(e) => { - tracing::error!( - "the rewritten request body could not be serialized; sending the original instead: {e}" - ); - body.clone() - } - } } pub fn map_reqwest_error(e: reqwest::Error) -> GatewayError { @@ -272,4 +291,201 @@ mod tests { "/v1beta/models/small:streamGenerateContent" ); } + + fn set( + model: Option<&str>, + max_tokens: Option, + thinking: Option, + ) -> tw_engine::SetAction { + tw_engine::SetAction { + model: model.map(str::to_string), + max_tokens, + thinking, + ..Default::default() + } + } + + /// 以前的写法：解析、改、整份写回（键按字母重排） + fn whole(body: &[u8], s: &tw_engine::SetAction, d: Option) -> Bytes { + let mut v: serde_json::Value = serde_json::from_slice(body).unwrap(); + set_params(&mut v, s, d); + Bytes::from(serde_json::to_vec(&v).unwrap()) + } + + /// 规则改了模型、输出上限、关了思考：上游收到的是客户端的那些字节，只有这几个字段变了 + /// —— 工具定义里的键、工具参数里的键、整个请求的先后都照客户端写的 + #[test] + fn a_rewritten_body_is_the_client_body_but_for_the_fields_set() { + use tw_dialect::ir::Dialect; + let tools = r#""tools":[{"name":"Read","description":"读 \"文件\"","input_schema":{"type":"object","properties":{"file_path":{"type":"string"},"offset":{"type":"number"},"limit":{"type":"number"}},"required":["file_path"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"}}]"#; + let messages = r#""messages":[{"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"Read","input":{"zeta":1,"alpha":{"b":2,"a":3}}}]}]"#; + let body = format!( + r#"{{"model":"claude-opus-5",{messages},"system":"s",{tools},"metadata":{{"user_id":"u"}},"max_tokens":32000,"thinking":{{"type":"enabled","budget_tokens":31999}},"stream":true}}"# + ); + let out = apply_set( + &Bytes::from(body.clone()), + &set(Some("claude-haiku-5"), Some(4096), Some(false)), + Some(Dialect::Anthropic), + ); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + format!( + r#"{{"model":"claude-haiku-5",{messages},"system":"s",{tools},"metadata":{{"user_id":"u"}},"max_tokens":4096,"stream":true}}"# + ) + ); + // 客户端没写的输出上限补在最后 + let out = apply_set( + &Bytes::from(format!(r#"{{"model":"m",{tools}}}"#)), + &set(None, Some(10), None), + Some(Dialect::Anthropic), + ); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + format!(r#"{{"model":"m",{tools},"max_tokens":10}}"#) + ); + // Gemini 的输出上限和思考在 generationConfig 里：只动里面那两个，别的照旧 + let gemini = format!( + r#"{{"contents":[],{tools},"generationConfig": {{"temperature": 0.2, "thinkingConfig": {{"thinkingBudget": 9}}, "maxOutputTokens": 9, "topK": 3}}}}"# + ); + let out = apply_set( + &Bytes::from(gemini), + &set(Some("ignored"), Some(512), Some(false)), + Some(Dialect::Gemini), + ); + assert_eq!( + std::str::from_utf8(&out).unwrap(), + format!( + r#"{{"contents":[],{tools},"generationConfig": {{"temperature": 0.2, "maxOutputTokens": 512, "topK": 3}}}}"# + ) + ); + } + + /// 改出来和原来一样（规则写的值客户端本来就这么写）：原样发，连拷贝都没有 + #[test] + fn a_set_that_changes_nothing_sends_the_same_bytes() { + let b = Bytes::from_static(br#"{ "max_tokens" : 100, "b":{"z":1,"a":2} }"#); + let out = apply_set(&b, &set(None, Some(100), None), None); + assert_eq!(out.as_ptr(), b.as_ptr()); + } + + /// 改到的对象里同一个键写了两遍：剪的结果说不准和 serde_json 认的一样（它留后一个）， + /// 照旧整份写回 + #[test] + fn a_key_written_twice_falls_back_to_writing_the_body_back() { + let b = br#"{"model":"a","x":1,"model":"b"}"#; + let s = set(Some("c"), None, None); + assert_eq!( + apply_set(&Bytes::from_static(b), &s, None), + whole(b, &s, None) + ); + } + + /// 随机拼出来的请求、随机的改写、五种格式：改出来的和整份写回的意思一样；没改到的 + /// 顶层成员原样、按原来的先后都在 + #[test] + fn rewriting_in_place_means_what_writing_back_meant() { + use tw_dialect::ir::Dialect; + struct Rng(u64); + impl Rng { + fn below(&mut self, n: usize) -> usize { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + (self.0 % n as u64) as usize + } + fn pick<'a>(&mut self, xs: &[&'a str]) -> &'a str { + xs[self.below(xs.len())] + } + } + let mut r = Rng(0x2545_f491_4f6c_dd1d); + // 转义的键（`"model"`）由 `char::from(92)` 拼：测试里直接写出来的转义，经过 + // 某些编辑工具会变成真字符 + let escaped_model = format!("\"mod{}u0065l\"", char::from(92)); + let config = [ + r#"{"temperature":0.2,"maxOutputTokens":9,"thinkingConfig":{"thinkingBudget":1}}"#, + r#"{"max_output_tokens":9,"thinking_config":{"a":1},"top_k":2}"#, + r#"{ }"#, + r#"{"thinkingConfig":1,"thinkingConfig":2}"#, + r#""not an object""#, + ]; + let members = [ + "\"model\":\"big\"", + "\"max_tokens\":9", + "\"max_completion_tokens\":7", + "\"max_output_tokens\":5", + "\"thinking\":{\"type\":\"enabled\",\"budget_tokens\":9}", + "\"reasoning_effort\":\"high\"", + "\"reasoning\":{\"effort\":\"low\"}", + "\"inferenceConfig\":{\"temperature\":1,\"maxTokens\":3}", + "\"additionalModelRequestFields\":{\"thinking\":{\"type\":\"enabled\"},\"k\":1}", + "\"tools\":[{\"input_schema\":{\"type\":\"object\",\"properties\":{\"b\":{},\"a\":{}}}}]", + "\"messages\":[{\"role\":\"user\",\"content\":\"max_tokens \\\"model\\\" }{\"}]", + ]; + const WS: &[&str] = &["", " ", "\n "]; + let dialects = [ + Dialect::Anthropic, + Dialect::Chat, + Dialect::Responses, + Dialect::Gemini, + Dialect::Bedrock, + ]; + let (mut changed, mut unusual) = (0, 0); + for _ in 0..5_000 { + let mut top: Vec = Vec::new(); + for _ in 0..r.below(7) { + let m = match r.below(14) { + 0 => format!("\"generationConfig\":{}", r.pick(&config)), + 1 => format!("\"generation_config\":{}", r.pick(&config)), + 2 => format!("{escaped_model}:\"esc\""), + _ => r.pick(&members).to_string(), + }; + let (a, b) = (r.pick(WS), r.pick(WS)); + top.push(format!("{a}{m}{b}")); + } + let body = format!("{{{}}}", top.join(",")); + let s = set( + [None, Some("small"), Some("big")][r.below(3)], + [None, Some(9), Some(512)][r.below(3)], + [None, Some(false), Some(true)][r.below(3)], + ); + if s.is_empty() { + continue; + } + let d = Some(dialects[r.below(dialects.len())]); + let out = apply_set(&Bytes::from(body.clone()), &s, d); + let value = |b: &[u8]| serde_json::from_slice::(b).unwrap(); + assert_eq!( + value(&out), + value(&whole(body.as_bytes(), &s, d)), + "{body} {s:?} {d:?}" + ); + let before = value(body.as_bytes()); + let mut after = before.clone(); + set_params(&mut after, &s, d); + if crate::splice::rewrite(body.as_bytes(), &before, &after) + == crate::splice::Rewrite::Unusual + { + unusual += 1; + continue; + } + changed += 1; + // 没改到的顶层成员：原文那一段照样在，先后不变 + let b = body.as_bytes(); + let mut from = 0; + for m in crate::splice::object(b, 0).unwrap() { + let name: String = serde_json::from_slice(&b[m.key.clone()]).unwrap(); + if before.get(&name) != after.get(&name) { + continue; + } + let piece = &body[m.key.start..m.value.end]; + let at = std::str::from_utf8(&out).unwrap()[from..] + .find(piece) + .unwrap_or_else(|| { + panic!("{piece} moved in {}", String::from_utf8_lossy(&out)) + }); + from += at + piece.len(); + } + } + assert!(changed > 1_000 && unusual > 50, "{changed} {unusual}"); + } } diff --git a/crates/tw-gateway/src/guard.rs b/crates/tw-gateway/src/guard.rs index d4eb6c2c..b0379ae8 100644 --- a/crates/tw-gateway/src/guard.rs +++ b/crates/tw-gateway/src/guard.rs @@ -51,6 +51,17 @@ pub fn look(mode: Mode, rules: &RuleSet, body: &[u8]) -> (Vec, Ledger) flow::look(mode, rules, body) } +/// [`look`]，连同找到的命中（见 [`flow::Look::hits`]）：同一份正文落盘前打码时不再找一遍 +/// （[`crate::bodies::BodyRecord::found`]）。 +pub fn look_hits(mode: Mode, rules: &RuleSet, body: &[u8]) -> flow::Look { + flow::look_hits( + mode, + rules, + body, + Ledger::new(tw_guard::redact::replace::Scheme::SECRET), + ) +} + /// [`look`]，接着 `seed` 的账编号（见 [`flow::look_from`]）。 pub fn look_from(mode: Mode, rules: &RuleSet, body: &[u8], seed: Ledger) -> (Vec, Ledger) { flow::look_from(mode, rules, body, seed) @@ -66,6 +77,30 @@ pub fn replace( flow::replace(mode, rules, body, ledger) } +/// [`replace`]，`found` 是按这套规则在 `body` 上已经找过的命中（[`look_hits`] 在同一份字节上 +/// 找的）：照它换，不再找一遍。没找过的是 `None`，和 [`replace`] 一样。 +pub fn replace_found( + mode: Mode, + rules: &RuleSet, + body: bytes::Bytes, + ledger: &Ledger, + found: Option<&[Hit]>, +) -> (bytes::Bytes, Ledger) { + let Some(hits) = found.filter(|_| mode.acts() && !rules.is_empty()) else { + return replace(mode, rules, body, ledger); + }; + // 没命中就原样返回，连一次拷贝都不做 + if hits.is_empty() { + return (body, ledger.clone()); + } + // 找得到命中的一定是 UTF-8（见 `flow::look_hits`） + let Ok(text) = std::str::from_utf8(&body) else { + return replace(mode, rules, body, ledger); + }; + let r = tw_guard::redact::replace::apply(text, hits, ledger.clone()); + (bytes::Bytes::from(r.text), r.ledger) +} + /// `after` 里 `before` 没有的那些值：插件写进请求里的（见 [`crate::plugin::request`]）。 /// /// 按规则和打过码的样子比：同一个值在两份里打出来的码一样。客户端原话里就有的值，开头 @@ -140,6 +175,16 @@ pub fn screen(s: &Screen, dialect: tw_dialect::ir::Dialect, body: &[u8]) -> Scre tw_guard::content::screen(s.mode, &s.rules, dialect, body) } +/// [`screen`]，请求体已经解析好了：`v` 是它解出来的样子（见 +/// [`tw_guard::content::screen_value`]）。 +pub fn screen_value( + s: &Screen, + dialect: tw_dialect::ir::Dialect, + v: &serde_json::Value, +) -> Screening { + tw_guard::content::screen_value(s.mode, &s.rules, dialect, v) +} + /// 把一次查下来的结论报出去：每条命中的规则一条 [`tw_api::Event::ContentMatched`]，挂在 /// 请求 `id` 上。要拒绝时返回告诉客户端的那句话。 pub fn report( @@ -435,6 +480,37 @@ mod tests { assert_eq!(ledger.len(), 150); } + /// 照开头那一遍找到的命中换，和自己再找一遍换的一字不差、账本一样 + #[test] + fn replacing_with_the_hits_already_found_is_replacing() { + let other = "ghp_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; + let body = bytes::Bytes::from(format!( + r#"{{"system":"<> {KEY}","messages":[{{"content":"{other} {KEY}"}}]}}"# + )); + let rules = RuleSet::defaults(); + let seen = look_hits(Mode::Enforce, &rules, &body); + let hits = seen.hits.unwrap(); + let again = replace(Mode::Enforce, &rules, body.clone(), &seen.ledger); + let known = replace_found( + Mode::Enforce, + &rules, + body.clone(), + &seen.ledger, + Some(&hits), + ); + assert_eq!(known.0, again.0); + assert_eq!(known.1.table(), again.1.table()); + assert!(!String::from_utf8_lossy(&known.0).contains(KEY)); + // 观察档不换，找过也不换 + let (out, _) = replace_found(Mode::Observe, &rules, body.clone(), &fresh(), Some(&hits)); + assert_eq!(out, body); + // 什么都没找到：原样，不拷贝 + let plain = bytes::Bytes::from_static(b"{\"messages\":[]}"); + let (out, l) = replace_found(Mode::Enforce, &rules, plain.clone(), &fresh(), Some(&[])); + assert_eq!(out.as_ptr(), plain.as_ptr()); + assert!(l.is_empty()); + } + /// 查表比和两两比，留下的一模一样：同一条规则下打码一样的算见过，换一条规则就不算 #[test] fn more_found_keeps_exactly_what_comparing_every_pair_kept() { diff --git a/crates/tw-gateway/src/lib.rs b/crates/tw-gateway/src/lib.rs index 3a111a64..7cb53c2f 100644 --- a/crates/tw-gateway/src/lib.rs +++ b/crates/tw-gateway/src/lib.rs @@ -45,6 +45,7 @@ pub mod sent; pub mod server; pub mod session; pub mod slots; +mod splice; pub mod state; pub mod translate; pub mod wire; diff --git a/crates/tw-gateway/src/server/pipeline.rs b/crates/tw-gateway/src/server/pipeline.rs index 07f56865..f2e8addf 100644 --- a/crates/tw-gateway/src/server/pipeline.rs +++ b/crates/tw-gateway/src/server/pipeline.rs @@ -59,6 +59,9 @@ struct Started { ledger: tw_guard::redact::replace::Ledger, /// 出站脱敏在客户端原文里找到的。插件改过的那一跳只再报插件写进来的（见 [`plug`]） found: Vec, + /// 拦截档下，出站脱敏在客户端原文（`req.body`）里找到的命中。一跳发出去的就是原文的 + /// 那些字节时（同格式直通、一个字节都没改），换的时候照它换，不再找一遍（见 [`hop`]） + hits: Option>, /// 这个请求最多等到什么时候：准入时定下（见 [`admission`]），等密钥的分钟、小时上限和 /// 等上游的空位共用这一段（`failover.slot_wait_secs`） wait_until: tokio::time::Instant, @@ -71,7 +74,14 @@ pub(super) async fn pipeline( live: crate::live::Pass, ending: &mut Option, ) -> Result { - let intent = match probe(&state, &rt, &req) { + // 请求体**只解析一次**：本地应答的判定、路由事实、对话的指纹、内容过滤都用这一份（见 + // `read`）。解不开是 None —— 照样往下走，同格式直通照样发 + let (parsed, probed) = heavy(&req.body, || { + let parsed = serde_json::from_slice::(&req.body).ok(); + let probed = probe(&state, &rt, &req, parsed.as_ref()); + (parsed, probed) + }); + let intent = match probed { Probe::Answered(resp) => return Ok(resp), Probe::Intent(intent) => intent, }; @@ -89,7 +99,7 @@ pub(super) async fn pipeline( // 管线第 2 步：读出路由事实，路由。**看的是客户端的原话**：插件的请求钩子排在路由 // 之后（每发往一个上游跑一次，见 `plug`），左右不了请求去哪一家 - let (mut reading, fp) = read(&req, intent); + let (mut reading, fp) = heavy(&req.body, || read(&req, intent, parsed.as_ref())); let conv = conversation(&rt, &req, &reading, fp.as_deref()); let (choice, decision) = match route(&state, &rt, &req, &reading, conv.as_ref())? { Routed::Go(choice, decision) => (choice, decision), @@ -98,8 +108,8 @@ pub(super) async fn pipeline( Routed::Refused(choice, why) => { let to = ("", tw_api::Billing::PerToken); // 一个字节都没发出去，也没什么可报的；存下来的请求照样按这一档换、打码 - let (_, ledger) = look(&rt, &req); - let redaction = redaction(&rt, ledger); + let seen = heavy(&req.body, || look(&rt, &req)); + let redaction = redaction(&rt, seen.ledger); let (id, _) = open( &state, &req, @@ -108,7 +118,7 @@ pub(super) async fn pipeline( to, fp.as_deref(), ending, - redaction, + (redaction, seen.hits.map(Into::into)), ); state.bus.emit(super::routed_nowhere(id, choice)); return Err(why); @@ -137,18 +147,24 @@ pub(super) async fn pipeline( .await?; // 管线第 4 步：内容过滤先下结论，不发事件。删过的话，后面一律用删过的那一份 - let screening = screen(&rt, &mut req, &mut reading); - let started = start( - &state, - &rt, - &req, - &reading, - choice, - wait_until, - &decision, - fp.as_deref(), - ending, - ); + let size = req.body.len(); + let (screening, started) = heavy_for(size, || { + let screening = screen(&rt, &mut req, &mut reading, parsed.as_ref()); + // 解析出来的那一份到这里就用完了：删过字的话它也不再是请求体的样子 + drop(parsed); + let started = start( + &state, + &rt, + &req, + &reading, + choice, + wait_until, + &decision, + fp.as_deref(), + ending, + ); + (screening, started) + }); // 用量上限的预留跟着请求号走，等存储层记下这一行时换成实数 hold.bind(started.id); // 结论挂在请求号上报。**拒绝的也在开始之后**：被拒是一次来源为 `denied` 的失败， @@ -291,10 +307,17 @@ enum Probe { /// 于是断网时健康检查照样失败，白白丢掉这个功能最有价值的场景。 /// 同样的理由让它排在「一个 provider 都没有」那一条之前：那一条也是 /// 一种「没有可用上游」，而本地应答本来就不需要上游。 -fn probe(state: &AppState, rt: &Runtime, req: &Inbound) -> Probe { - let Some(kind) = - crate::clientprobe::classify(&req.body, is_claude_code(&req.client_name, &req.headers)) - else { +fn probe( + state: &AppState, + rt: &Runtime, + req: &Inbound, + parsed: Option<&serde_json::Value>, +) -> Probe { + let Some(kind) = crate::clientprobe::classify( + &req.body, + parsed, + is_claude_code(&req.client_name, &req.headers), + ) else { return Probe::Intent(String::new()); }; use tw_config::ProbeAction::*; @@ -321,27 +344,30 @@ fn probe(state: &AppState, rt: &Runtime, req: &Inbound) -> Probe { /// 管线第 2 步的前半：读出路由事实，和这段对话的指纹。 /// -/// **只解析一次，指纹也只算一次。**路由要它，认对话（见 `crate::affinity`）要指纹， -/// 开始事件归会话也要指纹，而 body 可能有几百 KB —— 解两遍、哈希两遍是白付一份钱。 +/// **只解析一次，指纹也只算一次。**`parsed` 是管线开头解出来的那一份（本地应答的判定、 +/// 内容过滤用的也是它）。路由要它，认对话（见 `crate::affinity`）要指纹，开始事件归会话 +/// 也要指纹，而 body 可能有几 MB —— 解两遍、哈希两遍是白付一份钱。 /// /// 生成回答的请求解码成中间表示，**四种格式的客户端读出同一份路由事实**。 /// body 解不开时用空的性质走兜底规则。**不要因此拒绝请求** —— 我们的解析器 /// 不认识的东西，上游可能完全认识（只有需要转换时才用得上解码结果） -fn read(req: &Inbound, intent: String) -> (crate::client_api::Reading, Option) { - let parsed = serde_json::from_slice::(&req.body).ok(); - let mut reading = - crate::client_api::read(req.uri.path(), req.query.as_deref(), parsed.as_ref()); +fn read( + req: &Inbound, + intent: String, + parsed: Option<&serde_json::Value>, +) -> (crate::client_api::Reading, Option) { + let mut reading = crate::client_api::read(req.uri.path(), req.query.as_deref(), parsed); reading.facts.client = req.client_name.clone(); reading.facts.intent = intent; reading.harness = tw_dialect::harness::detect( req.headers .get(axum::http::header::USER_AGENT) .and_then(|v| v.to_str().ok()), - parsed.as_ref(), + parsed, ); // 认出「这几十个请求是同一次任务」。**认不出来就是 None** —— 硬凑一个会把 // 互不相干的请求并成一个「会话」 - let fp = parsed.as_ref().and_then(crate::session::fingerprint); + let fp = parsed.and_then(crate::session::fingerprint); (reading, fp) } @@ -817,7 +843,9 @@ fn start( // 是同一条记录**，差别只在换没换 —— 真正的替换在每一跳发出去之前做， // 那一跳的请求体可能是转换过格式的。拦截档下账本在这里就编好号：每一跳、 // 存下来的那份请求都按它换，同一个值处处是同一个占位符 - let (found, ledger) = look(rt, req); + let seen = look(rt, req); + let (found, ledger) = (seen.found, seen.ledger); + let hits: Option> = seen.hits.map(Into::into); let (id, at_ms) = open( state, req, @@ -826,7 +854,7 @@ fn start( (first, billing.into()), fp, ending, - redaction(rt, ledger.clone()), + (redaction(rt, ledger.clone()), hits.clone()), ); let redact_mode = rt.config.security.redact.mode; if !found.is_empty() { @@ -846,6 +874,7 @@ fn start( conversation: crate::affinity::identity(&req.headers, fp), ledger, found, + hits: hits.filter(|_| redact_mode.acts()), wait_until, } } @@ -854,14 +883,11 @@ fn start( /// /// 插件的密钥映射按同一份原文、同一个找法编号（见 [`crate::plugin::bridge`]）：插件看到的 /// 占位符和这本账里的是同一个号。插件往某一跳写进新的值，那一跳接着编（见 [`plug`]）。 -fn look( - rt: &Runtime, - req: &Inbound, -) -> ( - Vec, - tw_guard::redact::replace::Ledger, -) { - crate::guard::look(rt.config.security.redact.mode, &rt.redact, &req.body) +/// +/// 找到的命中跟着请求体交去留档：落盘前打码不再把同一份正文找一遍（见 +/// [`crate::bodies::BodyRecord::found`]）。 +fn look(rt: &Runtime, req: &Inbound) -> tw_guard::redact::flow::Look { + crate::guard::look_hits(rt.config.security.redact.mode, &rt.redact, &req.body) } /// 这个请求的正文落盘之前怎么换、怎么打码：此刻生效的规则，和这个请求的账本。 @@ -874,7 +900,8 @@ fn redaction(rt: &Runtime, ledger: tw_guard::redact::replace::Ledger) -> crate:: /// 发 `RequestStarted`、把这个请求欠着的结局放进 `ending`、把请求体交去留档， /// 交回这个请求的号和开始的时刻。`to` 是要发往的那一家和它怎么收钱；一家都不会去的 -/// （被规则拒绝了）是空的名字。`redaction` 是请求体、响应体落盘之前怎么换、打码。 +/// （被规则拒绝了）是空的名字。`redaction` 是请求体、响应体落盘之前怎么换、打码，和出站 +/// 脱敏在请求体上已经找到的命中（见 [`look`]）。 /// /// **会话在这里定**（见 [`crate::session::Sessions`]）：开始事件带着它，落库的 /// 那一行记的也是它。 @@ -887,7 +914,10 @@ fn open( to: (&str, tw_api::Billing), fp: Option<&str>, ending: &mut Option, - redaction: crate::bodies::Redaction, + (redaction, found): ( + crate::bodies::Redaction, + Option>, + ), ) -> (u64, u64) { let facts = &reading.facts; let id = state.bus.next_id(); @@ -948,7 +978,8 @@ fn open( req.body.clone(), req.body.len(), redaction, - ), + ) + .found(found), ); (id, at_ms) } @@ -965,17 +996,21 @@ fn open( /// 模型，查了只会在真正的请求之前把同一处命中多记一遍、还可能把计数请求拒掉（理由见 /// [`crate::client_api::ClientApi::screened`]）。在原文上查，中间表示解不开的请求照样查 /// （同格式直通照样发，上游可能认得它）。 +/// +/// `parsed` 是请求体解析出来的样子（管线开头解的那一份），解不开的是 None：解不开的请求体 +/// 没有调用方的正文可言，什么都不报。 fn screen( rt: &Runtime, req: &mut Inbound, reading: &mut crate::client_api::Reading, + parsed: Option<&serde_json::Value>, ) -> tw_guard::content::Screening { let screened = crate::client_api::ClientApi::screened(req.uri.path()); - let Some(api) = req.api.filter(|_| screened) else { + let (Some(api), Some(parsed)) = (req.api.filter(|_| screened), parsed) else { return Default::default(); }; let dialect = api.dialect(); - let sc = crate::guard::screen(&crate::guard::Screen::of(rt), dialect, &req.body); + let sc = crate::guard::screen_value(&crate::guard::Screen::of(rt), dialect, parsed); if let Some(body) = &sc.body { req.body = body.clone(); if reading.decoded.is_some() { @@ -998,6 +1033,32 @@ fn screen( sc } +/// 请求体到这么大，管线上整份解它、扫它的那几步挪出异步线程（见 [`heavy`]）。 +const HEAVY_BODY: usize = 1024 * 1024; + +/// 在整个请求体上做的 CPU 活：解析、本地应答的判定、内容过滤、出站脱敏、每一跳改写请求体。 +/// +/// **大的请求体挪出异步线程做**（`block_in_place`）：一个 8 MB 的请求要算几十毫秒，就地算的话， +/// 同一个线程上别的连接 —— 正在流的回答 —— 跟着停住这么久。挪出去时这个线程上排着的活交给 +/// 别的线程接着干。小的就地做：挪一次的开销比它本身还大。 +/// +/// 只在多线程的运行时上挪：单线程的运行时（测试用的就是它）里 `block_in_place` 会 panic。 +pub(super) fn heavy(body: &[u8], f: impl FnOnce() -> T) -> T { + heavy_for(body.len(), f) +} + +/// [`heavy`]，按请求体的长度：要在 `f` 里改请求体的调用方先量好长度 +pub(super) fn heavy_for(size: usize, f: impl FnOnce() -> T) -> T { + use tokio::runtime::{Handle, RuntimeFlavor}; + let threads = + Handle::try_current().is_ok_and(|h| h.runtime_flavor() == RuntimeFlavor::MultiThread); + if size >= HEAVY_BODY && threads { + tokio::task::block_in_place(f) + } else { + f() + } +} + /// 这个请求是 Claude Code 发的吗。 /// /// **`max_tokens: 1` 那条判定必须同时要求它**，否则会误伤别人 diff --git a/crates/tw-gateway/src/server/pipeline/admission.rs b/crates/tw-gateway/src/server/pipeline/admission.rs index 41f28cfd..e6cc42a1 100644 --- a/crates/tw-gateway/src/server/pipeline/admission.rs +++ b/crates/tw-gateway/src/server/pipeline/admission.rs @@ -140,7 +140,7 @@ fn refused( tracing::info!(key = %req.client_name, per = r.limit.per.word(), "a usage limit of the key refused the request"); let why = r.error(); // 一个字节都没发出去；存下来的请求照样按这一档换、打码 - let (_, ledger) = look(rt, req); + let seen = super::heavy(&req.body, || look(rt, req)); let (id, _) = open( state, req, @@ -149,7 +149,7 @@ fn refused( ("", tw_api::Billing::PerToken), fp, ending, - redaction(rt, ledger), + (redaction(rt, seen.ledger), seen.hits.map(Into::into)), ); state .bus diff --git a/crates/tw-gateway/src/server/pipeline/hop.rs b/crates/tw-gateway/src/server/pipeline/hop.rs index ece5c1de..2c29542d 100644 --- a/crates/tw-gateway/src/server/pipeline/hop.rs +++ b/crates/tw-gateway/src/server/pipeline/hop.rs @@ -64,6 +64,9 @@ struct Asked<'r> { body: &'r Bytes, path: &'r str, decoded: Option<&'r Result>, + /// `body` 解得开（是 JSON）：客户端的原话照管线开头解的那一遍，插件改过的那一份是 + /// 网关自己写出来的，一定解得开 + json: bool, } impl<'r> Asked<'r> { @@ -77,11 +80,13 @@ impl<'r> Asked<'r> { body: &r.body, path: &r.path, decoded: r.decoded.as_ref(), + json: true, }, None => Asked { body: &req.body, path: req.uri.path(), decoded: reading.decoded.as_ref(), + json: reading.json, }, } } @@ -402,15 +407,17 @@ pub(super) async fn try_upstreams<'a>( let model = Some(plugged.model.clone().unwrap_or(sent)).filter(asked_other); let asked = Asked::of(req, reading, &plugged); - let out = match prepare( - state, - req, - reading, - &asked, - provider, - &effective_set, - Some(id), - ) { + let out = match super::heavy(asked.body, || { + prepare( + state, + req, + reading, + &asked, + provider, + &effective_set, + Some(id), + ) + }) { Ok(out) => out, Err(err) => { chain.push(hop_failed( @@ -424,8 +431,6 @@ pub(super) async fn try_upstreams<'a>( } }; - // 这一家在这段对话里拒过的别家封存的推理：发之前先去掉（见 `crate::seal`） - let unsealed = unseal_upfront(state, req, started, provider, &out); // 出站脱敏的拦截档：换掉**这一跳真正发出去的那一份**（可能转换过 // 格式）。规则是全局的，每一跳换掉的是同一批东西；**接着原文那本账换**， // 同一个值在每一跳、在存下来的那份请求里都是同一个占位符。插件改过的一跳接着 @@ -434,8 +439,19 @@ pub(super) async fn try_upstreams<'a>( .rewritten .as_ref() .map_or(&started.ledger, |r| &r.ledger); - let (body, ledger) = - crate::guard::replace(rt.config.security.redact.mode, &rt.redact, unsealed, seed); + let (body, ledger) = super::heavy(&out.body, || { + // 这一家在这段对话里拒过的别家封存的推理：发之前先去掉（见 `crate::seal`） + let unsealed = unseal_upfront(state, req, started, provider, &out); + // 发出去的就是客户端原文的那些字节（同格式直通、一个字节都没改）：开头那一遍在 + // 它上面找到的就是这一跳要换的，不再找一遍 + let known = started.hits.as_deref().filter(|_| { + plugged.rewritten.is_none() + && unsealed.as_ptr() == req.body.as_ptr() + && unsealed.len() == req.body.len() + }); + let mode = rt.config.security.redact.mode; + crate::guard::replace_found(mode, &rt.redact, unsealed, seed, known) + }); // 用这个 provider 自己的 Client —— 它带着该走的代理。**在取密钥 // 之前拿到**：OAuth 换 token 也要走这条代理。 @@ -1133,6 +1149,7 @@ fn successor<'r>( body: &req.body, path: req.uri.path(), decoded: reading.decoded.as_ref(), + json: reading.json, }; rest.any(|name| { let Some(provider) = rt.config.providers.iter().find(|p| &p.name == name) else { @@ -1232,8 +1249,11 @@ fn prepare( } else if provider.forward_client_identity { out } else { - // 客户端自动填的身份字段不发（见 `egress` 模块）。**不算丢弃的字段** + // 客户端自动填的身份字段不发（见 `egress` 模块）。**不算丢弃的字段**。按字节 + // 剪，要请求体解得开：上面几步要么原样交回、要么照解出来的写回，解不开的照旧 + // 解不开、原样发 client_dialect + .filter(|_| asked.json) .and_then(|d| crate::egress::strip_body_identity(d, &out)) .unwrap_or(out) }; diff --git a/crates/tw-gateway/src/splice.rs b/crates/tw-gateway/src/splice.rs new file mode 100644 index 00000000..268b230e --- /dev/null +++ b/crates/tw-gateway/src/splice.rs @@ -0,0 +1,334 @@ +//! 在 JSON 原文上按字节改请求体：只动改了的那几个成员，别的字节一个不动。 +//! +//! 解析再写回不行：没开 `preserve_order` 的 serde_json 会把每个对象的键按字母重排 —— +//! 工具定义、工具参数、整个请求都变了样，而上游的提示缓存按字节认。所以去掉客户端的身份 +//! 字段（[`crate::egress`]）、按路由规则改参数（[`crate::forward::apply_set`]）都在原文上 +//! 剪、换、补。 +//! +//! **只给解得开的 JSON 用**：调用方解析过它。写坏了的地方这里不报错，只保证不越界。 + +use std::borrow::Cow; +use std::ops::Range; + +use serde_json::{Map, Value}; + +/// 对象里的一个成员：键（连同引号）和值在原文里的区间 +#[derive(Debug, Clone)] +pub(crate) struct Member { + pub(crate) key: Range, + pub(crate) value: Range, +} + +impl Member { + /// 键解出来的样子：写了转义的照 serde_json 解 + fn name<'b>(&self, b: &'b [u8]) -> Option> { + let raw = &b[self.key.start + 1..self.key.end - 1]; + if !raw.contains(&b'\\') { + return std::str::from_utf8(raw).ok().map(Cow::Borrowed); + } + serde_json::from_slice::(&b[self.key.clone()]) + .ok() + .map(Cow::Owned) + } + + /// 键解出来是不是 `name` + pub(crate) fn is(&self, b: &[u8], name: &str) -> bool { + self.name(b).is_some_and(|k| k == name) + } +} + +/// 从 `i` 起跳过空白 +pub(crate) fn ws(b: &[u8], mut i: usize) -> usize { + while i < b.len() && matches!(b[i], b' ' | b'\t' | b'\n' | b'\r') { + i += 1; + } + i +} + +/// `b[i]` 是一个字符串的开头引号：返回结尾引号之后的位置 +fn string_end(b: &[u8], i: usize) -> Option { + if b.get(i) != Some(&b'"') { + return None; + } + let mut j = i + 1; + loop { + j += memchr::memchr2(b'"', b'\\', b.get(j..)?)?; + if b[j] == b'"' { + return Some(j + 1); + } + j += 2; + } +} + +/// `b[i]` 起的一个值：返回它之后的位置 +fn value_end(b: &[u8], i: usize) -> Option { + match *b.get(i)? { + b'"' => string_end(b, i), + b'{' | b'[' => { + let mut depth = 0usize; + let mut j = i; + loop { + match *b.get(j)? { + b'"' => { + j = string_end(b, j)?; + continue; + } + b'{' | b'[' => depth += 1, + b'}' | b']' => { + depth -= 1; + if depth == 0 { + return Some(j + 1); + } + } + _ => {} + } + j += 1; + } + } + // 数、true、false、null：到下一个分隔为止 + _ => { + let end = b[i..] + .iter() + .position(|c| matches!(c, b',' | b'}' | b']' | b' ' | b'\t' | b'\n' | b'\r')) + .map_or(b.len(), |n| i + n); + (end > i).then_some(end) + } + } +} + +/// `b[open]` 是一个对象的 `{`：它的成员，按先后 +pub(crate) fn object(b: &[u8], open: usize) -> Option> { + if b.get(open) != Some(&b'{') { + return None; + } + let mut out = Vec::new(); + let mut i = ws(b, open + 1); + if b.get(i) == Some(&b'}') { + return Some(out); + } + loop { + let key = i..string_end(b, i)?; + let colon = ws(b, key.end); + if b.get(colon) != Some(&b':') { + return None; + } + let v = ws(b, colon + 1); + let value = v..value_end(b, v)?; + let next = ws(b, value.end); + out.push(Member { key, value }); + match *b.get(next)? { + b',' => i = ws(b, next + 1), + b'}' => return Some(out), + _ => return None, + } + } +} + +/// 剪掉 `members` 里第 `gone` 几个（按先后）要剪的区间，按先后、互不重叠。剪完还是 +/// 合法的 JSON：剪掉一个成员连同它后面的逗号；排在最后的，连同它前面的逗号 +pub(crate) fn cuts(members: &[Member], gone: &[usize]) -> Vec> { + let mut out = Vec::new(); + let mut k = 0; + while k < gone.len() { + // 连着要剪的一串 + let first = gone[k]; + let mut last = first; + while k + 1 < gone.len() && gone[k + 1] == last + 1 { + k += 1; + last += 1; + } + k += 1; + out.push(if last + 1 < members.len() { + members[first].key.start..members[last + 1].key.start + } else if first > 0 { + members[first - 1].value.end..members[last].value.end + } else { + members[first].key.start..members[last].value.end + }); + } + out +} + +/// [`rewrite`] 的结果。 +#[derive(Debug, PartialEq)] +pub(crate) enum Rewrite { + /// 改了的那一份 + Changed(Vec), + /// 改完和原来一样：原文一个字节都不用动 + Same, + /// 改到的对象里同一个键写了两遍（serde_json 留后一个，剪的结果说不准和它一样），或者 + /// 原文不像一个 JSON 对象：交给调用方解析、写回 + Unusual, +} + +/// 原文 `body`（解出来是 `before`）改成 `after` 的样子，**只动改了的地方**： +/// +/// - 值变了的成员只换它的值，两边都是对象的往里一层一层地比，只换里面变了的； +/// - 没了的成员连同分隔它的逗号剪掉； +/// - 新加的成员补在那个对象的末尾（按键名的先后）。 +/// +/// 没变的成员、它们的先后、空白和转义都照原文。改出来的那一份解析出来就是 `after`（测试 +/// 拿它和解析、写回的那一份比过）。 +pub(crate) fn rewrite(body: &[u8], before: &Value, after: &Value) -> Rewrite { + let (Some(before), Some(after)) = (before.as_object(), after.as_object()) else { + return Rewrite::Unusual; + }; + let mut edits = Vec::new(); + if diff(body, ws(body, 0), before, after, &mut edits).is_none() { + return Rewrite::Unusual; + } + if edits.is_empty() { + return Rewrite::Same; + } + // 同一处起的：先补（区间是空的），再剪、再换 + edits.sort_by_key(|(r, _): &(Range, Vec)| (r.start, r.end)); + let mut out = Vec::with_capacity(body.len() + 64); + let mut at = 0; + for (r, text) in edits { + debug_assert!(r.start >= at, "edits overlap"); + out.extend_from_slice(&body[at..r.start]); + out.extend_from_slice(&text); + at = r.end; + } + out.extend_from_slice(&body[at..]); + Rewrite::Changed(out) +} + +/// `b[open]` 处的对象从 `before` 改成 `after` 要做的改动，记进 `edits`（原文的区间换成 +/// 什么）。同一个键写了两遍、或者原文对不上的是 `None` +fn diff( + b: &[u8], + open: usize, + before: &Map, + after: &Map, + edits: &mut Vec<(Range, Vec)>, +) -> Option<()> { + let members = object(b, open)?; + let names: Vec> = members.iter().map(|m| m.name(b)).collect::>()?; + let mut sorted: Vec<&str> = names.iter().map(|n| n.as_ref()).collect(); + sorted.sort_unstable(); + if sorted.windows(2).any(|w| w[0] == w[1]) { + return None; + } + let mut gone = Vec::new(); + for (i, (m, name)) in members.iter().zip(&names).enumerate() { + let Some(now) = after.get(name.as_ref()) else { + gone.push(i); + continue; + }; + match (before.get(name.as_ref()), now) { + (Some(was), now) if was == now => {} + (Some(Value::Object(was)), Value::Object(now)) if b[m.value.start] == b'{' => { + diff(b, m.value.start, was, now, edits)?; + } + _ => edits.push((m.value.clone(), serde_json::to_vec(now).ok()?)), + } + } + let mut added = Vec::new(); + for (k, v) in after { + if !names.iter().any(|n| n == k) { + added.push(format!( + "{}:{}", + serde_json::to_string(k).ok()?, + serde_json::to_string(v).ok()? + )); + } + } + let added = added.join(","); + let kept = members.len() - gone.len(); + if !added.is_empty() && kept == 0 && !members.is_empty() { + // 原来的成员都没了：整段换成新加的 + let last = members.len() - 1; + edits.push(( + members[0].key.start..members[last].value.end, + added.into_bytes(), + )); + return Some(()); + } + for c in cuts(&members, &gone) { + edits.push((c, Vec::new())); + } + if !added.is_empty() { + match members + .iter() + .enumerate() + .rev() + .find(|(i, _)| !gone.contains(i)) + { + Some((_, m)) => { + let at = m.value.end; + edits.push((at..at, format!(",{added}").into_bytes())); + } + // 空对象：写在 `{` 后面 + None => edits.push((open + 1..open + 1, added.into_bytes())), + } + } + Some(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rewritten(body: &str, after: impl FnOnce(&mut Value)) -> Rewrite { + let before: Value = serde_json::from_str(body).unwrap(); + let mut now = before.clone(); + after(&mut now); + rewrite(body.as_bytes(), &before, &now) + } + + fn text(r: Rewrite) -> String { + match r { + Rewrite::Changed(b) => String::from_utf8(b).unwrap(), + other => panic!("{other:?}"), + } + } + + #[test] + fn only_what_changed_is_touched() { + let body = r#"{ "zeta": 1, "model": "opus", "tools": [{"name":"t","input_schema":{"type":"object","properties":{"b":{},"a":{}},"required":["b"]}}], "conf": {"y": 1, "x": 2, "drop": 3}, "gone": true }"#; + let out = text(rewritten(body, |v| { + v["model"] = "haiku".into(); + v["conf"]["x"] = 20.into(); + v["conf"].as_object_mut().unwrap().remove("drop"); + v["conf"]["new"] = "n".into(); + v.as_object_mut().unwrap().remove("gone"); + v["added"] = 7.into(); + })); + assert_eq!( + out, + r#"{ "zeta": 1, "model": "haiku", "tools": [{"name":"t","input_schema":{"type":"object","properties":{"b":{},"a":{}},"required":["b"]}}], "conf": {"y": 1, "x": 20,"new":"n"},"added":7 }"# + ); + } + + #[test] + fn nothing_changed_is_the_same_and_an_emptied_object_takes_the_new_members() { + assert_eq!(rewritten(r#"{"a":1}"#, |_| {}), Rewrite::Same); + assert_eq!( + text(rewritten(r#"{"c":{"old":1}}"#, |v| v["c"] = + serde_json::json!({"new": 2}))), + r#"{"c":{"new":2}}"# + ); + assert_eq!( + text(rewritten(r#"{"c":{ }}"#, |v| v["c"]["k"] = 1.into())), + r#"{"c":{"k":1 }}"# + ); + assert_eq!( + text(rewritten(r#"{}"#, |v| v["k"] = 1.into())), + r#"{"k":1}"# + ); + } + + #[test] + fn a_key_written_twice_where_something_changes_is_unusual() { + assert_eq!( + rewritten(r#"{"a":1,"a":2}"#, |v| v["b"] = 1.into()), + Rewrite::Unusual + ); + // 写了两遍的在没改到的那一层里：照原文留着，和客户端发的一样 + assert_eq!( + text(rewritten(r#"{"m":{"a":1,"a":2},"x":1}"#, |v| v["x"] = 2.into())), + r#"{"m":{"a":1,"a":2},"x":2}"# + ); + } +} diff --git a/crates/tw-gateway/tests/headers_and_paths.rs b/crates/tw-gateway/tests/headers_and_paths.rs index 764876bd..843e08a0 100644 --- a/crates/tw-gateway/tests/headers_and_paths.rs +++ b/crates/tw-gateway/tests/headers_and_paths.rs @@ -440,6 +440,34 @@ async fn an_anthropic_upstream_gets_the_request_and_nothing_about_the_client() { assert_eq!(v["messages"][0]["content"], "hi"); } +/// 去掉 `user_id` 之后，上游收到的就是客户端发的那些字节：键的先后（工具定义、工具参数、 +/// 整个请求）、空白、转义一个都没动。以前解析再写回，每个对象的键都按字母重排了 —— 而 +/// Claude Code 的请求每一个都带着 `user_id`。 +/// +/// 体过了 1 MiB、跑在多线程的运行时上：这么大的请求体，管线上那几步是挪出异步线程算的 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn the_upstream_gets_the_client_body_byte_for_byte_less_the_identity() { + let (up, seen) = start_upstream().await; + let gw = start_gateway(anthropic(up)).await; + let file = "fn main() {\n println!(\"读了 \\\"一个\\\" 文件\");\n}\n".repeat(20_000); + let file = serde_json::to_string(&file).unwrap(); + let head = format!( + r#"{{"model":"claude-sonnet-4-5","messages":[{{"role":"user","content":[{{"type":"tool_result","tool_use_id":"toolu_01","content":{file}}}]}},{{"role":"assistant","content":[{{"type":"tool_use","id":"toolu_02","name":"Read","input":{{"zeta":1,"alpha":[2,{{"b":3,"a":4}}]}}}}]}}], "tools":[{{"name":"Read","input_schema":{{"type":"object","properties":{{"file_path":{{"type":"string"}},"limit":{{"type":"number"}}}},"required":["file_path"],"$schema":"http://json-schema.org/draft-07/schema#"}}}}],"# + ); + let identity = r#""metadata":{"user_id":"user_abc_account_123_session_456"},"#; + let tail = r#""max_tokens":16,"stream":false}"#; + let body = format!("{head}{identity}{tail}"); + assert!(body.len() > 1024 * 1024, "{}", body.len()); + post(gw, "/v1/messages", CLAUDE_CODE, &body).await; + + let g = seen.lock().unwrap(); + assert!( + g.body == format!("{head}{tail}").as_bytes(), + "the upstream got a different body: {}", + String::from_utf8_lossy(&g.body[g.body.len().saturating_sub(300)..]) + ); +} + #[tokio::test] async fn an_upstream_that_asks_for_it_gets_the_clients_own_identity() { let (up, seen) = start_upstream().await; diff --git a/crates/tw-guard/Cargo.toml b/crates/tw-guard/Cargo.toml index e37cbffd..9a35191b 100644 --- a/crates/tw-guard/Cargo.toml +++ b/crates/tw-guard/Cargo.toml @@ -16,6 +16,7 @@ base64 = { workspace = true } bytes = { workspace = true } serde_json = { workspace = true } regex = { workspace = true } +memchr = { workspace = true } thiserror = { workspace = true } serde_yaml_ng = { workspace = true } tw-dialect = { workspace = true } diff --git a/crates/tw-guard/src/content.rs b/crates/tw-guard/src/content.rs index 79e4e7b5..12ceba1d 100644 --- a/crates/tw-guard/src/content.rs +++ b/crates/tw-guard/src/content.rs @@ -34,7 +34,7 @@ mod screen; pub(crate) use codepoints::visible as codepoints_visible; pub use codepoints::{CodepointError, Codepoints, MAX_ITEMS as MAX_CODEPOINT_ITEMS}; pub(crate) use screen::evaluate; -pub use screen::{Outcome, ScreenHit, Screening, screen, screen_text}; +pub use screen::{Outcome, ScreenHit, Screening, screen, screen_text, screen_value}; /// 一条规则怎么认。 #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize)] @@ -259,10 +259,54 @@ impl<'a> Lower<'a> { Lower { text, lower: None } } fn get(&mut self) -> &str { - self.lower.get_or_insert_with(|| self.text.to_lowercase()) + self.lower.get_or_insert_with(|| lowercase(self.text)) } } +/// 和 `text.to_lowercase()` 一样的结果，快一些：ASCII 的一段整段转，别的字符逐个转（标准库 +/// 碰到头一个 ASCII 以外的字符之后就逐个字符地转了，夹着中文、`→` 的正文几乎整段都是）。 +/// +/// 按上下文转的只有 `Σ`（在词尾是 `ς`）：正文里有它就整段交给标准库。 +fn lowercase(text: &str) -> String { + if memchr::memmem::find(text.as_bytes(), "Σ".as_bytes()).is_some() { + return text.to_lowercase(); + } + let b = text.as_bytes(); + let mut out = String::with_capacity(text.len()); + let mut i = 0; + while i < b.len() { + let start = i; + i = ascii_end(b, i); + if i > start { + let at = out.len(); + out.push_str(&text[start..i]); + out[at..].make_ascii_lowercase(); + } + // 停在一个多字节字符的头一个字节上：字符边界 + if let Some(c) = text.get(i..).and_then(|rest| rest.chars().next()) { + out.extend(c.to_lowercase()); + i += c.len_utf8(); + } + } + out +} + +/// 从 `i` 起连着的 ASCII 字节到哪儿为止。八个一组地看 +fn ascii_end(b: &[u8], mut i: usize) -> usize { + const HIGH: u64 = 0x8080_8080_8080_8080; + while let Some(word) = b.get(i..i + 8) { + let word = u64::from_ne_bytes(word.try_into().unwrap_or_default()); + if word & HIGH != 0 { + break; + } + i += 8; + } + while i < b.len() && b[i].is_ascii() { + i += 1; + } + i +} + /// 一条规则在一段正文里的全部命中 #[derive(Debug, Default)] struct Found { @@ -289,7 +333,9 @@ fn find_contains(text: &str, lower: &str, needle: &str) -> Vec> { return out; } if lower.len() == text.len() { - for (at, _) in lower.match_indices(needle) { + // 按字节找（SIMD）：每段正文、每条规则都要从头找到尾。和 `match_indices` 一样是 + // 一个接一个、互不重叠的 + for at in memchr::memmem::find_iter(lower.as_bytes(), needle.as_bytes()) { // 区间落回原文要对齐到字符边界 let start = floor(text, at); let end = ceil(text, at + needle.len()); @@ -373,19 +419,46 @@ fn find_points(p: &Codepoints, text: &str, escapes: bool) -> Found { if p.min() > 0x7F && text.is_ascii() && (!escapes || !text.contains("\\u")) { return out; } - for (r, c) in chars(text, escapes) { - if !p.contains(c) { - continue; - } + let mut hit = |r: Range| { out.count += 1; match out.ranges.last_mut() { Some(last) if last.end == r.start => last.end = r.end, _ => out.ranges.push(r), } + }; + // 码位都在 ASCII 以外、又不认转义：ASCII 的字节一个都命中不了，跳过去，只解码别的 + // 字符。夹着中文、`→` 的正文（工具读回来的代码）就不再逐个字符地解一遍 + if p.min() > 0x7F && !escapes { + for (r, c) in non_ascii(text) { + if p.contains(c) { + hit(r); + } + } + } else { + for (r, c) in chars(text, escapes) { + if p.contains(c) { + hit(r); + } + } } out } +/// 一段正文里 ASCII 以外的每一个字符和它的字节区间，按出现的先后。ASCII 的字节八个一组 +/// 地跳过（[`ascii_end`]）。 +fn non_ascii(text: &str) -> impl Iterator, char)> + '_ { + let b = text.as_bytes(); + let mut i = 0; + std::iter::from_fn(move || { + i = ascii_end(b, i); + // 停在一个多字节字符的头一个字节上：字符边界 + let c = text.get(i..)?.chars().next()?; + let at = i; + i += c.len_utf8(); + Some((at..i, c)) + }) +} + /// 一处命中。 #[derive(Debug, Clone, PartialEq, Eq)] pub struct Hit { @@ -978,4 +1051,108 @@ mod tests { assert_eq!(find_points(&p, &text, false).count, 0); assert!(text.is_ascii(), "例子里不该有真的不可见字符"); } + + /// 码位都在 ASCII 以外时只解码 ASCII 以外的字符：找到的和逐个字符看的一样 —— 几段、 + /// 每段从哪到哪、一共几个 + #[test] + fn skipping_ascii_finds_what_looking_at_every_character_found() { + let mut seed = 0x2545_f491_4f6c_dd1d_u64; + let mut next = |n: usize| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + (seed % n as u64) as usize + }; + // 码位用 `char::from_u32` 拼：测试里直接写出来的转义，经过某些编辑工具会变成真字符 + let ch = |n: u32| char::from_u32(n).unwrap().to_string(); + let pieces: Vec = [ + "a", + "Z", + " ", + "0", + "\"", + "{", + "→", + "密钥", + "é", + "😀", + "abcdefghijklmnop", + ] + .iter() + .map(|s| s.to_string()) + .chain( + [ + 0x200B, 0x202E, 0x2066, 0xE0041, 0xE007F, 0xE0000, 0x7F, 0x80, 0xFEFF, + ] + .into_iter() + .map(ch), + ) + .collect(); + let sets = [ + "U+E0000–U+E007F", + "U+202A–U+202E, U+2066–U+2069", + "U+200B, U+FEFF", + "U+0080–U+10FFFF", + ] + .map(|p| Codepoints::parse(p).unwrap()); + for _ in 0..5_000 { + let text: String = (0..next(60)) + .map(|_| pieces[next(pieces.len())].as_str()) + .collect(); + for p in &sets { + assert!(p.min() > 0x7F); + let fast = find_points(p, &text, false); + let mut slow = Found::default(); + for (r, c) in chars(&text, false) { + if p.contains(c) { + slow.count += 1; + match slow.ranges.last_mut() { + Some(last) if last.end == r.start => last.end = r.end, + _ => slow.ranges.push(r), + } + } + } + assert_eq!( + (fast.count, fast.ranges), + (slow.count, slow.ranges), + "{text:?}" + ); + } + } + } + + /// 快一些的小写和标准库的一字不差：`Σ` 在词尾、词中，会变长的（`İ`）、四个字节的、 + /// 组合字符、开尔文符号这类转成 ASCII 的 + #[test] + fn lowercasing_by_runs_gives_what_the_standard_library_gives() { + let mut seed = 0x9e37_79b9_7f4a_7c15_u64; + let mut next = |n: usize| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + (seed % n as u64) as usize + }; + // 码位用 `char::from_u32` 拼：测试里直接写出来的转义，经过某些编辑工具会变成真字符 + let ch = |n: u32| char::from_u32(n).unwrap().to_string(); + let pieces: Vec = ["A", "z", " ", "IGNORE", "Previous", ".", "→", "密钥", "😀"] + .iter() + .map(|s| s.to_string()) + .chain( + [ + 0x03A3, 0x0391, 0x0130, 0x212A, 0x1E9E, 0x00DF, 0x01C5, 0xFB00, 0x0301, 0x2167, + 0x10400, 0x00C9, 0x0049, 0x0131, 0x03C2, + ] + .into_iter() + .map(ch), + ) + .collect(); + for _ in 0..20_000 { + let text: String = (0..next(30)) + .map(|_| pieces[next(pieces.len())].as_str()) + .collect(); + assert_eq!(lowercase(&text), text.to_lowercase(), "{text:?}"); + } + let long = format!("{}{}", "Ignore Previous ".repeat(10), ch(0x212A)); + assert_eq!(lowercase(&long), long.to_lowercase()); + } } diff --git a/crates/tw-guard/src/content/screen.rs b/crates/tw-guard/src/content/screen.rs index 91337f92..bd7bc38f 100644 --- a/crates/tw-guard/src/content/screen.rs +++ b/crates/tw-guard/src/content/screen.rs @@ -88,17 +88,26 @@ pub fn screen(mode: Mode, rules: &Rules, dialect: Dialect, body: &[u8]) -> Scree if !mode.detects() || rules.is_empty() { return Screening::default(); } - let Ok(mut v) = serde_json::from_slice::(body) else { + let Ok(v) = serde_json::from_slice::(body) else { return Screening::default(); }; - let spots = tw_dialect::caller::spots(dialect, &v); + screen_value(mode, rules, dialect, &v) +} + +/// [`screen`]，请求体已经解析好了（调用方为别的事解过一遍）：`v` 是它解出来的样子，不再 +/// 解一遍。删过的请求体照它写出来，和 [`screen`] 写的一样。 +pub fn screen_value(mode: Mode, rules: &Rules, dialect: Dialect, v: &Value) -> Screening { + if !mode.detects() || rules.is_empty() { + return Screening::default(); + } + let spots = tw_dialect::caller::spots(dialect, v); if spots.is_empty() { return Screening::default(); } let e = { let segments: Vec<(&str, bool)> = spots .iter() - .map(|s| (s.get(&v).unwrap_or_default(), s.in_tool_result)) + .map(|s| (s.get(v).unwrap_or_default(), s.in_tool_result)) .collect(); evaluate(rules, &segments, Scope::default()) }; @@ -107,6 +116,8 @@ pub fn screen(mode: Mode, rules: &Rules, dialect: Dialect, body: &[u8]) -> Scree && out.refused.is_none() && let Some(texts) = e.texts { + // 要删的时候才抄一份来改 + let mut v = v.clone(); for (s, t) in spots.iter().zip(texts) { if let Some(x) = s.get_mut(&mut v) { *x = t; diff --git a/crates/tw-guard/src/redact/flow.rs b/crates/tw-guard/src/redact/flow.rs index 42e8a2eb..ef8aeea7 100644 --- a/crates/tw-guard/src/redact/flow.rs +++ b/crates/tw-guard/src/redact/flow.rs @@ -128,25 +128,51 @@ pub fn look(mode: Mode, rules: &RuleSet, body: &[u8]) -> (Vec, Ledger) /// 请求），改过的这一份接着原文那本账编，同一个值还是同一个号，新出现的值接着往后编。 /// 不在替换档时账本是空的，`seed` 用不上。 pub fn look_from(mode: Mode, rules: &RuleSet, body: &[u8], seed: Ledger) -> (Vec, Ledger) { - let empty = || Ledger::new(Scheme::SECRET); + let l = look_hits(mode, rules, body, seed); + (l.found, l.ledger) +} + +/// [`look_from`] 的结果，连同找到的那些命中。 +#[derive(Debug, Clone)] +pub struct Look { + /// 报出去的记录（见 [`look`]） + pub found: Vec, + /// 这个请求的账本。不在替换档时是空的 + pub ledger: Ledger, + /// 这一遍在 `body` 上找到的命中（同 [`hits`]，区间是 `body` 里的）。没找的是 `None`： + /// 关着、一条规则都没开、不是 UTF-8。 + /// + /// 同一份正文之后还要再找一遍的调用方（桌面版落盘前打码）拿它省掉那一遍：规则一样、 + /// 正文一样，找出来的就一样 + pub hits: Option>, +} + +/// [`look_from`]，把找到的命中一起交回去（见 [`Look::hits`]）。 +pub fn look_hits(mode: Mode, rules: &RuleSet, body: &[u8], seed: Ledger) -> Look { + let empty = || Look { + found: Vec::new(), + ledger: Ledger::new(Scheme::SECRET), + hits: None, + }; if !mode.detects() || rules.is_empty() { - return (Vec::new(), empty()); + return empty(); } let Ok(text) = std::str::from_utf8(body) else { - return (Vec::new(), empty()); + return empty(); }; let hits = hits(text, rules); let found = crate::redact::rules::findings(text, &hits); - if !mode.acts() { - return (found, empty()); - } - let seed = seed.avoiding(text); - let ledger = if hits.is_empty() { - seed + let ledger = if !mode.acts() { + Ledger::new(Scheme::SECRET) } else { - crate::redact::replace::apply(text, &hits, seed).ledger + // 只要账本：换过的那一份在每一跳发出去之前才写（见 `replace`），这里不抄一份 + crate::redact::replace::number(text, &hits, seed.avoiding(text)) }; - (found, ledger) + Look { + found, + ledger, + hits: Some(hits), + } } /// 替换档下换掉要发出去的这一份，**接着 `ledger` 的账**（见 [`look`]）。返回换过的体和 @@ -398,6 +424,43 @@ mod tests { assert!(l.is_empty()); } + /// 看一遍只编号、不抄一份：账本和真换一遍得到的一样（让开原文里写着的占位符、同一个值 + /// 一个号），交回的命中就是 [`hits`] 找到的 + #[test] + fn looking_numbers_like_replacing_and_hands_back_its_hits() { + let other = "ghp_BBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBBB"; + let body = format!( + r#"{{"system":"<> {KEY}","messages":[{{"content":"{other} {KEY}"}}]}}"# + ); + let rules = RuleSet::defaults(); + let l = look_hits(Mode::Enforce, &rules, body.as_bytes(), fresh()); + let found = hits(&body, &rules); + assert_eq!(found.len(), 3); + assert_eq!(l.hits.as_deref(), Some(&found[..])); + let replaced = crate::redact::replace::apply(&body, &found, fresh().avoiding(&body)); + assert_eq!(l.ledger.table(), replaced.ledger.table()); + assert_eq!(l.ledger.len(), 2); + assert_eq!((l.found.clone(), l.ledger.table().clone()), { + let (f, ledger) = look(Mode::Enforce, &rules, body.as_bytes()); + (f, ledger.table().clone()) + }); + // 观察档照样找、照样交回命中，只是不编号 + let o = look_hits(Mode::Observe, &rules, body.as_bytes(), fresh()); + assert_eq!((o.found, o.hits), (l.found, l.hits)); + assert!(o.ledger.is_empty()); + // 关着的、不是 UTF-8 的不找 + assert!( + look_hits(Mode::Off, &rules, body.as_bytes(), fresh()) + .hits + .is_none() + ); + assert!( + look_hits(Mode::Enforce, &rules, &[0xff, 0xfe], fresh()) + .hits + .is_none() + ); + } + /// 连接串里的占位符长得像口令，可它不是凭据：不再换一次、不报、原样留着 #[test] fn a_placeholder_where_a_password_would_be_is_not_a_password() { diff --git a/crates/tw-guard/src/redact/replace.rs b/crates/tw-guard/src/redact/replace.rs index af29e3d9..3c7880a2 100644 --- a/crates/tw-guard/src/redact/replace.rs +++ b/crates/tw-guard/src/redact/replace.rs @@ -207,6 +207,20 @@ pub fn apply(text: &str, hits: &[Hit], mut ledger: Ledger) -> Redacted { Redacted { text: out, ledger } } +/// 只编号、不抄：[`apply`] 会给 `hits` 发的那些号，不写出换过的那一份。只要账本时用它 +/// （拦截档下看一遍客户端原文，见 [`crate::redact::flow::look`]）—— 换过的那一份在每一跳 +/// 发出去之前才写。 +pub(crate) fn number(text: &str, hits: &[Hit], mut ledger: Ledger) -> Ledger { + let default = ledger.scheme.label; + for h in hits { + ledger.issue( + &text[h.bytes.clone()], + h.label.as_deref().unwrap_or(default), + ); + } + ledger +} + /// 扫 + 换，一步到位。`text` 是 JSON 请求体（见 [`crate::redact::rules::scan`]）。 pub fn redact(text: &str, rules: &RuleSet, ledger: Ledger) -> Redacted { let hits = crate::redact::rules::scan(text, rules); diff --git a/crates/tw-guard/src/redact/rules.rs b/crates/tw-guard/src/redact/rules.rs index a4ef6ec5..648714ea 100644 --- a/crates/tw-guard/src/redact/rules.rs +++ b/crates/tw-guard/src/redact/rules.rs @@ -523,12 +523,68 @@ impl RuleSet { } } -fn is_tok(c: char) -> bool { - c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' +/// 一个字节是不是 token 的一部分。**token 只由 ASCII 组成**，所以按字节切和按字符切切出来的 +/// 一样：多字节字符的每个字节都不小于 0x80，既落不进 token，也不是反斜杠和引号。 +const fn is_tok(b: u8) -> bool { + b.is_ascii_alphanumeric() || b == b'-' || b == b'_' || b == b'.' } -/// 一段 token 是哪种 API 密钥。**只看开着的那几条。** -fn classify_token(tok: &str, set: &RuleSet) -> Option<&'static str> { +/// [`is_tok`] 查表：每个请求体的每个字节都要问一次 +const TOK: [bool; 256] = { + let mut t = [false; 256]; + let mut b = 0; + while b < 256 { + t[b] = is_tok(b as u8); + b += 1; + } + t +}; + +/// 一把密钥最短有多长：前缀规则里最短的那一种（前缀加上至少要有的尾巴），和 OpenAI 老式 +/// 密钥的下限，取小的。比它短的 token 哪一种都不是 —— 一段正文里绝大多数 token 在这里就 +/// 回去了 +const KEY_MIN: usize = { + let mut min = OPENAI_MIN; + let mut i = 0; + while i < BUILTINS.len() { + if let Matcher::Prefix { prefix, min_tail } = BUILTINS[i].matcher + && prefix.len() + min_tail < min + { + min = prefix.len() + min_tail; + } + i += 1; + } + min +}; + +/// 一把密钥能以哪些字节开头：每条前缀规则的头一个字节，和 OpenAI 老式密钥的 `s` +const KEY_FIRST: [bool; 256] = { + let mut t = [false; 256]; + t[b's' as usize] = true; + let mut i = 0; + while i < BUILTINS.len() { + if let Matcher::Prefix { prefix, .. } = BUILTINS[i].matcher { + t[prefix.as_bytes()[0] as usize] = true; + } + i += 1; + } + t +}; + +/// 一段 token 是哪种 API 密钥。**只看开着的那几条。**`openai`：OpenAI 老式密钥那条开着 +/// （扫之前问一次，不在每个 token 上查一遍表）。 +/// +/// 太短的、开头哪条前缀都对不上的直接回去：[`classify_prefixed`] 认得出的 token 都过得了 +/// 这两关（[`KEY_MIN`]、[`KEY_FIRST`]）。 +fn classify_token(tok: &str, set: &RuleSet, openai: bool) -> Option<&'static str> { + if tok.len() < KEY_MIN || !KEY_FIRST[usize::from(tok.as_bytes()[0])] { + return None; + } + classify_prefixed(tok, set, openai) +} + +/// [`classify_token`] 的判据本身。 +fn classify_prefixed(tok: &str, set: &RuleSet, openai: bool) -> Option<&'static str> { for b in BUILTINS { if let Matcher::Prefix { prefix, min_tail } = b.matcher && let Some(tail) = tok.strip_prefix(prefix) @@ -539,7 +595,7 @@ fn classify_token(tok: &str, set: &RuleSet) -> Option<&'static str> { return set.is_on(b.id).then_some(b.id); } } - if set.is_on("openai-api-key") + if openai && let Some(tail) = tok.strip_prefix("sk-") && tok.len() >= OPENAI_MIN && !tail.starts_with("ant-") @@ -569,34 +625,43 @@ fn classify_token(tok: &str, set: &RuleSet) -> Option<&'static str> { /// /// 回调的第三个参数说这个 token 在不在一个 JSON 字符串里面（数没转义的引号； /// 文本是 JSON 时才有意义）。 +/// +/// **按字节走，不按字符解码**：token、反斜杠、引号都是 ASCII（见 [`is_tok`]），多字节字符 +/// 的字节在这里只是分隔。反斜杠后面跟着的是一个多字节字符时，跳过的只是它的头一个字节， +/// 剩下的几个照样只是分隔 —— 切出来的和按字符切的一样（测试拿两种切法对过）。token 的两头 +/// 都落在 ASCII 字节上，切片一定在字符边界上。 fn for_each_token(text: &str, mut f: impl FnMut(&str, Range, bool)) { + let b = text.as_bytes(); let mut start: Option = None; let mut quoted = false; - let mut chars = text.char_indices().peekable(); - while let Some((i, c)) = chars.next() { - if c == '\\' { - if let Some(s) = start.take() { - f(&text[s..i], s..i, quoted); - } - // 反斜杠后面那个字符是转义的一部分；`\u` 再带四位十六进制 - if let Some((_, 'u')) = chars.next() { - for _ in 0..4 { - if chars.next_if(|(_, h)| h.is_ascii_hexdigit()).is_none() { - break; - } - } - } - continue; - } - if is_tok(c) { + let mut i = 0; + while i < b.len() { + let c = b[i]; + if TOK[usize::from(c)] { start.get_or_insert(i); + i += 1; continue; } if let Some(s) = start.take() { f(&text[s..i], s..i, quoted); } - if c == '"' { - quoted = !quoted; + i += 1; + match c { + // 反斜杠后面那个字符是转义的一部分；`\u` 再带四位十六进制 + b'\\' => { + if b.get(i) == Some(&b'u') { + i += 1; + let mut n = 0; + while n < 4 && b.get(i).is_some_and(u8::is_ascii_hexdigit) { + i += 1; + n += 1; + } + } else if i < b.len() { + i += 1; + } + } + b'"' => quoted = !quoted, + _ => {} } } if let Some(s) = start { @@ -615,7 +680,8 @@ const PEM_KEY: &str = "PRIVATE KEY-----"; /// 空壳 PEM，它会以为文件坏了然后建议你重新生成一把 —— 那是在帮倒忙。 fn private_keys(text: &str, out: &mut Vec) { let mut from = 0; - while let Some(b) = text[from..].find(PEM_BEGIN) { + // 按字节找（SIMD）：整个请求体都要过一遍，而它几乎从来不在里面 + while let Some(b) = memchr::memmem::find(&text.as_bytes()[from..], PEM_BEGIN.as_bytes()) { let begin = from + b; // **不能按换行找那一行的结尾。**请求体是 JSON，里面的换行是 // `\n` 两个字符，不是一个换行符 —— 按换行找的话整个 PEM 会落在 @@ -663,6 +729,10 @@ fn is_b64url(s: &str) -> bool { /// 那个解码是必需的：光看「三段用点分开的 base64」会把版本号、文件路径、 /// 甚至 `a.b.c` 这种普通标识符全算上。 fn looks_like_jwt(tok: &str) -> bool { + // 首段至少 8 个字符，后两段各至少一个，加两个点：短于 12 的不可能是 + if tok.len() < 12 { + return false; + } // 先数点，不把三段收集起来：每个 token 都要过这一关，绝大多数在这里就回去了 if tok.bytes().filter(|&b| b == b'.').count() != 2 || !tok.split('.').all(is_b64url) { return false; @@ -695,7 +765,7 @@ fn looks_like_jwt(tok: &str) -> bool { fn conn_strings(text: &str, out: &mut Vec) { let bytes = text.as_bytes(); let mut from = 0; - while let Some(i) = text[from..].find("://") { + while let Some(i) = memchr::memmem::find(&bytes[from..], b"://") { let sep = from + i; let after = sep + 3; // userinfo 到 `@` 为止；`@` 必须在同一段里（不能跨空白或斜杠） @@ -752,14 +822,16 @@ fn is_rfc1918(tok: &str) -> bool { /// **不含回环。**`127.0.0.1` 和 `localhost` 是这台机器自己 —— 而用户问 /// 的很可能正是「我本地这个服务为什么连不上」，把它换成占位符等于把问题 /// 本身藏起来了。何况我们的网关自己就住在那儿。 -fn internal_token(tok: &str, set: &RuleSet) -> Option<&'static str> { - if set.is_on("internal-ip") && is_rfc1918(tok) { +/// +/// `ip`、`domain`：那两条开没开（扫之前问一次，不在每个 token 上查一遍表）。 +fn internal_token(tok: &str, ip: bool, domain: bool) -> Option<&'static str> { + if ip && is_rfc1918(tok) { return Some("internal-ip"); } // 后缀不分大小写地比，**不先转成小写**：每个 token 都要过这一关，转一次就是 // 一次分配 let t = tok.as_bytes(); - if set.is_on("internal-domain") + if domain && t.len() > 7 && t.contains(&b'.') && INTERNAL_SUFFIXES.iter().any(|s| { @@ -1264,6 +1336,11 @@ pub fn scan(text: &str, set: &RuleSet) -> Vec { if set.is_on("email") { emails(text, &mut out); } + // 每个 token 都要问的几条开没开，扫之前问一次：在每个 token 上查一遍表，光哈希就占了 + // 整遍扫描的两成 + let want_openai = set.is_on("openai-api-key"); + let want_ip = set.is_on("internal-ip"); + let want_domain = set.is_on("internal-domain"); let want_jwt = set.is_on("jwt"); let want_id = set.is_on("cn-resident-id"); let want_card = set.is_on("bank-card"); @@ -1291,7 +1368,7 @@ pub fn scan(text: &str, set: &RuleSet) -> Vec { // 字符串外面的东西把在接的那一串截断 spaced.finish(text, &mut out); } - if let Some(id) = classify_token(tok, set) { + if let Some(id) = classify_token(tok, set, want_openai) { out.push(builtin_hit(id, span)); return; } @@ -1299,7 +1376,7 @@ pub fn scan(text: &str, set: &RuleSet) -> Vec { out.push(builtin_hit("jwt", span)); return; } - if let Some(id) = internal_token(tok, set) { + if let Some(id) = internal_token(tok, want_ip, want_domain) { out.push(builtin_hit(id, span)); } }); @@ -2672,4 +2749,187 @@ mod tests { ); assert_eq!(crate::redact::replace::restore(&r.text, &r.ledger), body); } + + /// 以前的切法：按字符解码着走。按字节走的 [`for_each_token`] 切出来的要和它一模一样 + fn tokens_by_char(text: &str) -> Vec<(String, Range, bool)> { + let mut out = Vec::new(); + let mut f = |t: &str, r: Range, q: bool| out.push((t.to_string(), r, q)); + let mut start: Option = None; + let mut quoted = false; + let mut chars = text.char_indices().peekable(); + while let Some((i, c)) = chars.next() { + if c == '\\' { + if let Some(s) = start.take() { + f(&text[s..i], s..i, quoted); + } + if let Some((_, 'u')) = chars.next() { + for _ in 0..4 { + if chars.next_if(|(_, h)| h.is_ascii_hexdigit()).is_none() { + break; + } + } + } + continue; + } + if c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.' { + start.get_or_insert(i); + continue; + } + if let Some(s) = start.take() { + f(&text[s..i], s..i, quoted); + } + if c == '"' { + quoted = !quoted; + } + } + if let Some(s) = start { + f(&text[s..], s..text.len(), quoted); + } + out + } + + fn tokens_by_byte(text: &str) -> Vec<(String, Range, bool)> { + let mut out = Vec::new(); + for_each_token(text, |t, r, q| out.push((t.to_string(), r, q))); + out + } + + /// 随机拼出来的一段：ASCII 的字母数字和标点、引号、反斜杠和各种转义（含不完整的 + /// `\u`、反斜杠后面跟着中文和 emoji、结尾一个落单的反斜杠）、中文、emoji、组合字符 + fn random_text(next: &mut impl FnMut(usize) -> usize) -> String { + const PIECES: &[&str] = &[ + "a", + "Z", + "9", + "-", + "_", + ".", + " ", + ",", + ":", + "{", + "}", + "[", + "]", + "\"", + "\\", + "\\\\", + "\\n", + "\\t", + "\\\"", + "\\u", + "\\u00e9", + "\\u5bc6", + "\\uD83D\\uDE00", + "\\u12", + "\\uzz", + "\\é", + "\\密", + "\\😀", + "u", + "00e9", + "sk-ant-", + "AKIA", + "eyJ", + "密钥", + "：", + "😀", + "👩‍💻", + "e\u{301}", + "\u{200B}", + "\u{7f}", + "\t", + "\n", + "é", + "→", + "ghp_", + "@", + "/", + "+", + "=", + ]; + let n = next(40); + (0..n).map(|_| PIECES[next(PIECES.len())]).collect() + } + + #[test] + fn cutting_by_byte_cuts_exactly_where_cutting_by_char_did() { + let mut seed = 0x2545_f491_4f6c_dd1d_u64; + let mut next = |n: usize| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + (seed % n as u64) as usize + }; + for _ in 0..20_000 { + let t = random_text(&mut next); + assert_eq!(tokens_by_byte(&t), tokens_by_char(&t), "{t:?}"); + } + for t in [ + "", + "\\", + "a\\", + "\\u", + "\\u0", + "\\😀x", + "\"a\\\"b\"c", + "密\\钥sk", + "x\u{1F600}", + ] { + assert_eq!(tokens_by_byte(t), tokens_by_char(t), "{t:?}"); + } + } + + /// 太短的、开头哪条前缀都对不上的 token 直接回去：哪一条规则开着、关着，结论都和逐条 + /// 比前缀的一样 + #[test] + fn the_quick_rejections_never_drop_a_key() { + let mut seed = 0x9e37_79b9_7f4a_7c15_u64; + let mut next = |n: usize| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + (seed % n as u64) as usize + }; + const TAIL: &[u8] = b"abcXYZ0189-_."; + let mut heads: Vec<&str> = BUILTINS + .iter() + .filter_map(|b| match b.matcher { + Matcher::Prefix { prefix, .. } => Some(prefix), + _ => None, + }) + .collect(); + heads.extend(["sk-", "s", "x", "", "Sk-", "ak"]); + let sets = [ + all(), + RuleSet::defaults(), + RuleSet::none(), + RuleSet::only(&["openai-api-key"]), + RuleSet::only(&["anthropic-api-key", "aws-access-key-id"]), + ]; + for _ in 0..20_000 { + let mut tok = heads[next(heads.len())].to_string(); + for _ in 0..next(50) { + tok.push(TAIL[next(TAIL.len())] as char); + } + if tok.is_empty() { + continue; + } + for set in &sets { + for openai in [true, false] { + assert_eq!( + classify_token(&tok, set, openai), + classify_prefixed(&tok, set, openai), + "{tok}" + ); + } + } + } + // 最短的那种刚好够长时认得出 + assert_eq!(KEY_MIN, "AKIA".len() + 12); + assert_eq!( + classify_token("AKIAABCDEFGHIJKL", &all(), true), + Some("aws-access-key-id") + ); + } } diff --git a/crates/tw-observe/src/bus.rs b/crates/tw-observe/src/bus.rs index d0dd2f61..d51119b4 100644 --- a/crates/tw-observe/src/bus.rs +++ b/crates/tw-observe/src/bus.rs @@ -7,16 +7,23 @@ //! //! 注意这条豁免**不适用于成本记账**：那类数据丢了账单永久对不 //! 上，必须走别的路径。这里只走「丢了只是图上少个点」的东西。 +//! +//! **存储层不订阅广播**，它有自己的一条有界通道（[`EventBus::record_feed`]）：广播上 +//! 掉队丢掉的是界面上的几个点，记录丢掉的是那几个请求，永远不在库里了。 use std::collections::{BTreeMap, HashMap, VecDeque}; use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, OnceLock}; use std::time::{Duration, Instant}; -use tokio::sync::broadcast; +use tokio::sync::{broadcast, mpsc}; const CAPACITY: usize = 1024; +/// 存储层那条通道有多少格（见 [`EventBus::record_feed`]）。一个请求五六条事件，这么多格 +/// 够存储层落后两千多个请求 +pub const RECORD_CAPACITY: usize = 16 * 1024; + /// 生成速率看最近这么久里跑完的请求 pub const RATE_WINDOW: Duration = Duration::from_secs(60); @@ -26,6 +33,39 @@ pub struct EventBus { next_id: Arc, /// 在跑的和最近跑完的。见 [`EventBus::in_flight`]、[`EventBus::live`] tally: Arc>, + /// 存储层自己的那条通道，接上了才有。见 [`EventBus::record_feed`] + record: Arc>, +} + +/// 往存储层那条通道里送的一头。 +struct RecordSink { + tx: mpsc::Sender, + /// 通道满着、没送进去的条数，存储层拿去记一行日志 + dropped: Arc, +} + +/// 存储层收事件的那一头（[`EventBus::record_feed`]）。 +pub struct RecordFeed { + rx: mpsc::Receiver, + dropped: Arc, +} + +impl RecordFeed { + /// 等下一条。总线没了是 None。**会挡住当前线程**：给存储层自己的记录线程用，不在 + /// 异步线程上调 + pub fn blocking_recv(&mut self) -> Option { + self.rx.blocking_recv() + } + + /// 已经到了的下一条，不等 + pub fn try_recv(&mut self) -> Option { + self.rx.try_recv().ok() + } + + /// 上次问到现在，通道满着丢掉了几条 + pub fn take_dropped(&self) -> u64 { + self.dropped.swap(0, Ordering::Relaxed) + } } /// 按事件数出来的现状。 @@ -124,6 +164,7 @@ impl EventBus { tx, next_id: Arc::new(AtomicU64::new(1)), tally: Arc::default(), + record: Arc::default(), } } @@ -149,11 +190,38 @@ impl EventBus { /// /// **先记账再发**（见 [`EventBus::in_flight`]）：订阅者收到一个结局的 /// 时候，快照里已经没有它了。 + /// + /// 存储层那条通道满着时**丢，不等**：转发不等观测。 pub fn emit(&self, ev: tw_api::Event) { self.track(&ev); + if let Some(sink) = self.record.get() + && let Err(mpsc::error::TrySendError::Full(_)) = sink.tx.try_send(ev.clone()) + { + sink.dropped.fetch_add(1, Ordering::Relaxed); + } let _ = self.tx.send(ev); } + /// 存储层收事件的那条通道：**有界，和界面的广播分开**，之后发的每一条都进来。 + /// + /// 广播满了丢最老的，掉队的订阅者跳过一截。对界面那是少画几个点，它还会收到 + /// `EventsDropped` 自己对账；对记录那是几个请求永远不在库里。以前存储层就订阅广播， + /// 和界面挤同一个 [`CAPACITY`] 格的环。走自己这条的话，[`RECORD_CAPACITY`] 格平常 + /// 一条都不丢；真满了照样丢（转发不等观测），存储层说一声丢了几条。 + /// + /// **只给一次**，第二次是 None：两个记录器各分到一半事件，两边都缺。 + pub fn record_feed(&self) -> Option { + let (tx, rx) = mpsc::channel(RECORD_CAPACITY); + let dropped = Arc::new(AtomicU64::new(0)); + self.record + .set(RecordSink { + tx, + dropped: dropped.clone(), + }) + .ok()?; + Some(RecordFeed { rx, dropped }) + } + /// 开始的记下，有了结局的划掉。 /// /// **这里什么都不能 panic。**`Ending` 在 Drop 里也会走到这儿，那时可能 @@ -753,6 +821,41 @@ mod tests { assert!(c > a); } + /// **存储层不和界面挤一个环。**界面那个订阅者落后到被跳过一截时，存储层那条通道里 + /// 一条不少、顺序不乱 —— 以前它也订阅广播，界面落后的时候它多半也落后，丢掉的那几个 + /// 请求永远不在库里 + #[test] + fn the_store_gets_every_event_while_a_lagging_subscriber_skips_some() { + let b = EventBus::new(); + let mut ui = b.subscribe(); + let mut feed = b.record_feed().unwrap(); + let n = CAPACITY as u64 * 4; + for id in 0..n { + b.emit(finished(id, 0, None)); + } + assert!(matches!( + ui.try_recv(), + Err(broadcast::error::TryRecvError::Lagged(_)) + )); + let got: Vec = std::iter::from_fn(|| feed.try_recv().map(|e| e.id())).collect(); + assert_eq!(got, (0..n).collect::>()); + assert_eq!(feed.take_dropped(), 0); + } + + /// 存储层那条通道满了：**丢，不挡**，丢了几条说得出来。只给一次 + #[test] + fn a_full_store_feed_drops_and_counts_rather_than_blocking() { + let b = EventBus::new(); + let mut feed = b.record_feed().unwrap(); + assert!(b.record_feed().is_none(), "给了第二个记录器"); + for id in 0..(RECORD_CAPACITY as u64 + 10) { + b.emit(finished(id, 0, None)); + } + assert_eq!(feed.take_dropped(), 10); + assert_eq!(feed.take_dropped(), 0, "问过一次就清零"); + assert_eq!(feed.try_recv().map(|e| e.id()), Some(0)); + } + #[tokio::test] async fn a_slow_subscriber_lags_but_does_not_block_the_producer() { // 这是关键性质：观测端跟不上，数据面**不能**因此变慢。 diff --git a/crates/tw-observe/src/lib.rs b/crates/tw-observe/src/lib.rs index 0c4bd591..dd312d29 100644 --- a/crates/tw-observe/src/lib.rs +++ b/crates/tw-observe/src/lib.rs @@ -6,4 +6,4 @@ pub mod bus; -pub use bus::EventBus; +pub use bus::{EventBus, RecordFeed}; diff --git a/crates/tw-store/src/blobs.rs b/crates/tw-store/src/blobs.rs index 4644a9ee..f1bae15b 100644 --- a/crates/tw-store/src/blobs.rs +++ b/crates/tw-store/src/blobs.rs @@ -3,6 +3,10 @@ //! **不进 SQLite。**body 可能几百 KB（长上下文），塞进库里会让它膨胀到 //! GB 级 —— VACUUM 慢、备份慢、回收难。存文件系统按天分目录，过期直接 //! 删掉整个目录，**回收成本 O(1)**。 +//! +//! **一份正文要么整个在、要么不在**：先写到旁边一个临时文件，写完再换上去（见 +//! [`Blobs::put`]）。读的一方（详情、对话记录、按正文找）和写的一方不排队，以前读到一半 +//! 写着的文件会当成一份读不懂的正文。 use std::path::{Path, PathBuf}; @@ -99,15 +103,8 @@ impl Blobs { // 截断而不是跳过：**开头那几 KB 是最有用的部分**（模型名、system // prompt、工具定义都在前面），而完整存下来会挤掉别人的。 let slice = &body[..body.len().min(MAX_ONE)]; - match write_private(&p, slice) { - Ok(()) => { - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let _ = std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o600)); - } - true - } + match replace_private(&p, slice) { + Ok(()) => true, Err(e) => { tracing::debug!(path = %p.display(), "the request body could not be written: {e}"); false @@ -131,6 +128,9 @@ impl Blobs { /// `original_len` 比**真正存下的**长时才另记一个 `.len` —— 截断可能发生在交来之前 /// （网关只攒了开头），也可能发生在这里（`body` 比 [`MAX_ONE`] 长）。以前只看前一种： /// 一个 5 MB 的请求体存下 4 MB，却没有一处说它被截过，重放照样把半截 JSON 发了出去。 + /// + /// **`.len` 先写**：正文一出现，读的一方就该知道它是不是截过的。反过来的话，中间有一刻 + /// 读到的是一份看起来完整的半截正文。 pub fn put_with_len( &self, at_ms: i64, @@ -139,22 +139,17 @@ impl Blobs { body: &[u8], original_len: usize, ) -> bool { - if !self.put(at_ms, id, which, body) { - return false; - } if original_len > body.len().min(MAX_ONE) { let p = self .path_for(at_ms, id, which) .with_extension(format!("{}.len", which.suffix())); - if write_private(&p, original_len.to_string().as_bytes()).is_ok() { - #[cfg(unix)] - { - use std::os::unix::fs::PermissionsExt; - let _ = std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o600)); - } + if let Some(dir) = p.parent() + && create_private_dirs(dir).is_ok() + { + let _ = replace_private(&p, original_len.to_string().as_bytes()); } } - true + self.put(at_ms, id, which, body) } /// 原始长度。**没有这个文件说明没截断** —— 那时 body 自己的长度 @@ -186,13 +181,18 @@ impl Blobs { /// 盘上最老的那一天从哪一刻起（那天零点，UTC）。一天都没有是 `None`。 /// /// **按正文找时用它划界**：比它早的请求一份正文都不会有，不必一条一条去盘上问 —— - /// 三个月的记录里，正文只占最近几天。回收是整天整天地删（先删最老的），所以划在 - /// 最老的那一天上是准的。 + /// 三个月的记录里，正文只占最近几天。回收先删最老的，整天整天地删；只剩最新的一天时 + /// 才从那一天里最早的请求删起。所以比它早的一份都没有，它之后的照样要去盘上问。 pub fn oldest_ms(&self) -> Option { self.days().first().and_then(|(name, _)| start_of_day(name)) } /// 回收：先按天数删，再按总量删。返回删掉的字节数。 + /// + /// **按总量删时，最新的那一天不整个删。**从最老的那天整天整天地删；删到只剩最新的一天 + /// 还超，就在这一天里从最早的请求删起（[`trim_oldest`]）。以前它也整个删：一天的正文就 + /// 超过上限时（一个请求最多存 8 MB 多），每小时的回收都把当天的目录清空一次，连同正在 + /// 看的那次会话。 pub fn gc(&self, now_ms: i64, keep_days: u64, max_bytes: u64) -> u64 { let cutoff = day_of(now_ms - (keep_days as i64) * 86_400_000); let mut freed = 0; @@ -207,10 +207,13 @@ impl Blobs { remaining.push((name, path, size)); } } - // 还超总量的话，继续从最旧的开始删。**整目录删，不删单个文件** —— - // 半天的记录比没有记录更难解释（「为什么上午的请求点开是空的」）。 let mut total: u64 = remaining.iter().map(|(_, _, s)| *s).sum(); - for (_, path, size) in &remaining { + let Some(((_, newest, _), older)) = remaining.split_last() else { + return freed; + }; + // 还超总量的话，继续从最旧的开始删。**旧的整目录删，不删单个文件** —— + // 半天的记录比没有记录更难解释（「为什么上午的请求点开是空的」）。 + for (_, path, size) in older { if total <= max_bytes { break; } @@ -219,6 +222,9 @@ impl Blobs { freed += size; } } + if total > max_bytes { + freed += trim_oldest(newest, total - max_bytes); + } freed } @@ -254,6 +260,60 @@ fn write_private(p: &Path, bytes: &[u8]) -> std::io::Result<()> { opts.open(p)?.write_all(bytes) } +/// 换上一份只给自己看的文件：**先写到旁边的临时文件，写完再改名换上去**。改名是一步完成 +/// 的，读的一方要么读到旧的（或者没有），要么读到写完的那一份，不会读到一半。 +fn replace_private(p: &Path, bytes: &[u8]) -> std::io::Result<()> { + let mut tmp = p.as_os_str().to_owned(); + tmp.push(".tmp"); + let tmp = PathBuf::from(tmp); + let done = write_private(&tmp, bytes).and_then(|()| std::fs::rename(&tmp, p)); + if done.is_err() { + let _ = std::fs::remove_file(&tmp); + } + #[cfg(unix)] + if done.is_ok() { + // 临时文件早就在（上次写到一半断了）的话，新建时给的 0600 不生效，这里补上 + use std::os::unix::fs::PermissionsExt; + let _ = std::fs::set_permissions(p, std::fs::Permissions::from_mode(0o600)); + } + done +} + +/// 在一天的目录里从最早的请求删起，删够 `excess` 字节为止。返回删掉的字节数。 +/// +/// **一个请求的几份一起删**（请求体、回答、插件改过的、`.len`）：按文件名开头的请求号归在 +/// 一起，号小的先删 —— 号是按请求开始的先后发的。认不出请求号的文件不动。 +fn trim_oldest(day: &Path, excess: u64) -> u64 { + let Ok(rd) = std::fs::read_dir(day) else { + return 0; + }; + let mut by_id: std::collections::BTreeMap> = Default::default(); + for e in rd.flatten() { + let name = e.file_name(); + let Some(id) = name + .to_str() + .and_then(|n| n.split_once('.')) + .and_then(|(id, _)| id.parse::().ok()) + else { + continue; + }; + let size = e.metadata().map(|m| m.len()).unwrap_or(0); + by_id.entry(id).or_default().push((e.path(), size)); + } + let mut freed = 0; + for files in by_id.into_values() { + if freed >= excess { + break; + } + for (path, size) in files { + if std::fs::remove_file(&path).is_ok() { + freed += size; + } + } + } + freed +} + fn looks_like_a_day(n: &str) -> bool { let b = n.as_bytes(); b.len() == 10 @@ -463,6 +523,58 @@ mod tests { assert_eq!(left[0].0, day_of(now)); } + /// **一天的正文就超过上限时，那一天不整个删。**以前每小时的回收都把当天的目录清空： + /// 正在看的那次会话、刚结束的那个请求，点开都是空的。现在删到只剩最新的一天，就在这一天 + /// 里从最早的请求删起，一个请求的几份一起删 + #[test] + fn gc_trims_the_newest_day_from_its_oldest_requests_instead_of_emptying_it() { + let (_d, b) = setup(); + let now = 100 * DAY; + // 前一天的整个删 + b.put(now - DAY, 1, Which::Request, &vec![b'x'; 1000]); + // 当天五个请求，各有请求体和回答，2000 字节一个；第 10 个还记着原本多长 + for id in 10..15 { + b.put(now + id, id, Which::Request, &vec![b'x'; 1000]); + b.put(now + id, id, Which::Response, &vec![b'x'; 1000]); + } + b.put_with_len(now, 10, Which::Request, &vec![b'x'; 1000], 9999); + let today = b.root().join(day_of(now)); + let before = b.total_bytes(); + + let freed = b.gc(now, 30, 3500); + + assert_eq!(b.days().len(), 1, "{:?}", b.days()); + assert!(today.is_dir(), "当天的目录被整个删了"); + assert!(b.total_bytes() <= 3500, "{}", b.total_bytes()); + assert_eq!(freed, before - b.total_bytes()); + // 最早的四个连同 `.len` 一起没了，最新的那个还在 + for id in 10..14 { + assert!(b.get(now, id, Which::Request).is_none(), "{id}"); + assert!(b.get(now, id, Which::Response).is_none(), "{id}"); + } + assert_eq!(b.original_len(now, 10, Which::Request), None); + assert!(b.get(now, 14, Which::Request).is_some()); + assert!(b.get(now, 14, Which::Response).is_some()); + + // 下一小时没有新写入：已经在上限之内，什么都不再删 + assert_eq!(b.gc(now + 3_600_000, 30, 3500), 0); + assert!(b.get(now, 14, Which::Response).is_some()); + } + + /// 写到一半的正文读不到：先写临时文件再换上去，换完不留临时文件 + #[test] + fn a_body_is_swapped_in_whole_and_leaves_no_temporary_file_behind() { + let (_d, b) = setup(); + assert!(b.put(0, 1, Which::Response, b"first")); + assert!(b.put(0, 1, Which::Response, b"second, longer")); + assert_eq!(b.get(0, 1, Which::Response).unwrap(), b"second, longer"); + let names: Vec = std::fs::read_dir(b.root().join(day_of(0))) + .unwrap() + .map(|e| e.unwrap().file_name().to_string_lossy().to_string()) + .collect(); + assert_eq!(names, ["1.res"]); + } + #[test] fn a_stray_directory_the_user_put_there_is_left_alone() { // 这是用户自己的磁盘。**不认识的东西不删。** diff --git a/crates/tw-store/src/db.rs b/crates/tw-store/src/db.rs index b809c39c..b55c36db 100644 --- a/crates/tw-store/src/db.rs +++ b/crates/tw-store/src/db.rs @@ -50,6 +50,20 @@ const NO_PRICE: &str = "(cost_micros IS NULL AND input_tokens IS NOT NULL \ const NO_USAGE: &str = "(cost_micros IS NULL AND input_tokens IS NULL AND error IS NULL \ AND billing = 'per-token' AND (cancelled = 1 OR status < 300))"; +/// 数路由命中要的那几项，从路由那一列（`tw_api::RoutingView` 的 JSON）里取（见 +/// [`Db::route_hits`]）。**建索引的和查询的是同一串**：表达式一字不差，SQLite 才用索引里 +/// 存好的值，不回表去解 JSON。 +/// +/// 用 `json_extract` 不用 `->>`：索引写在库文件里，读这个库的 SQLite 都得认得它。 +const ROUTE: &str = "json_extract(routing, '$.route')"; +const RULE: &str = "json_extract(routing, '$.rule')"; +/// 一个 JSON 数组的原文（`[]`、`["关思考"]`） +const REWRITTEN_BY: &str = "json_extract(routing, '$.rewritten_by')"; +const DENIED_BY: &str = "json_extract(routing, '$.denied_by')"; +/// 经过路由、路由那一列解得开的行。**解不开的不进索引**：一行坏掉的记录不该让整张表出不来， +/// 也不该让它自己写不进去 +const ROUTED: &str = "local = 0 AND json_valid(routing)"; + #[derive(Debug, thiserror::Error)] pub enum DbError { #[error("{path} could not be opened: {source}")] @@ -272,9 +286,42 @@ impl Db { if found == 0 { db.create()?; } + db.indexes(); Ok(db) } + /// 表之外另加的索引，**每次打开都走一遍**（`IF NOT EXISTS`，建过了就是空操作）。 + /// + /// **它们不算 schema，加它们不加 [`SCHEMA`]**：索引不改变一行长什么样，少了它查询照样 + /// 对、只是慢，多了它别的版本照样读写（SQLite 自己维护）。加 [`SCHEMA`] 却要清掉每个人 + /// 的请求历史 —— 为了快一点不值得。已有的库第一次打开时要建一遍（二十七万行一秒上下）。 + /// + /// **建不成只记一行日志**：那时查询照样对，只是慢。 + fn indexes(&self) { + // - `requests_client`：每把密钥最后一次用在什么时候，一把一把地跳着取（见 + // `last_seen_by_client`）。**只给那一条查询用**：条件里要写上 `client > ''` 才用得上 + // 它。按时间窗、按密钥分组的那些（密钥的用量、按密钥的费用）用了它，会把整个索引按 + // 密钥扫一遍，比沿着时刻走慢得多 + // - `requests_session_at`：最近活动过的会话，按时间倒着走，只读索引（见 `sessions`） + // - `requests_routed`：各条路由、规则命中了多少要的那几项，落库时从路由那一列里取出来 + // 存在索引里。数的时候只读索引，不再把一周几 MB 的 JSON 读出来解一遍（见 + // `route_hits`） + let sql = format!( + "CREATE INDEX IF NOT EXISTS requests_client ON requests (client, at_ms) + WHERE client > ''; + CREATE INDEX IF NOT EXISTS requests_session_at ON requests (at_ms, session) + WHERE session IS NOT NULL AND local = 0; + CREATE INDEX IF NOT EXISTS requests_routed + ON requests (at_ms, {ROUTE}, {RULE}, {REWRITTEN_BY}, {DENIED_BY}, error IS NOT NULL) + WHERE {ROUTED};" + ); + if let Err(e) = self.conn.execute_batch(&sql) { + tracing::warn!( + "the request history could not be indexed, so some views load slowly: {e}" + ); + } + } + /// 建表。**只有这一份，没有迁移**：表的样子一变，改这里、[`SCHEMA`] 加一， /// 版本对不上的旧库整个重建。 /// @@ -535,6 +582,52 @@ pub struct SessionRow { pub errors: i64, } +/// 一次会话汇总的那几列，和 [`session_row`] 一列对一列。`GROUP BY session` 之后用 +fn session_columns() -> String { + format!( + "session, + client, + MIN(at_ms), MAX(at_ms), COUNT(*), + COALESCE(SUM(cost_micros), 0), + COALESCE(SUM({NO_PRICE}), 0), + COALESCE(SUM(input_tokens), 0), + COALESCE(SUM(output_tokens), 0), + COALESCE(SUM(cache_read_tokens), 0), + COALESCE(SUM(cache_write_tokens), 0), + COALESCE(SUM(cache_saved_micros), 0), + COALESCE(MAX(input_tokens), 0), + GROUP_CONCAT(DISTINCT model), + SUM(CASE WHEN error IS NOT NULL THEN 1 ELSE 0 END), + COALESCE(SUM(CASE WHEN cost_estimated = 1 THEN cost_micros ELSE 0 END), 0), + COUNT(cost_micros), + COALESCE(SUM({NO_USAGE}), 0)" + ) +} + +/// 读出 [`session_columns`] 的一行 +fn session_row(r: &rusqlite::Row) -> rusqlite::Result { + Ok(SessionRow { + id: r.get(0)?, + client: r.get(1)?, + started_ms: r.get(2)?, + ended_ms: r.get(3)?, + turns: r.get(4)?, + cost_micros: r.get(5)?, + unpriced_turns: r.get(6)?, + input_tokens: r.get(7)?, + output_tokens: r.get(8)?, + cache_read_tokens: r.get(9)?, + cache_write_tokens: r.get(10)?, + cache_saved_micros: r.get(11)?, + peak_input_tokens: r.get(12)?, + models: r.get::<_, Option>(13)?.unwrap_or_default(), + errors: r.get(14)?, + cost_micros_estimated: r.get(15)?, + priced_turns: r.get(16)?, + no_usage_turns: r.get(17)?, + }) +} + /// 会话里的一轮。上下文增长曲线和成本瀑布画的就是它。 #[derive(Debug, Clone, PartialEq)] pub struct TurnRow { @@ -568,63 +661,93 @@ impl Db { /// /// **筛的是会话，不是轮次。**把轮次按时间筛掉再聚合的话，一次跨过 /// 窗口边界的任务会少算几轮、少算一截钱 —— 而「那次重构花了多少」 - /// 问的是整次任务，不是它落在某个窗口里的那一段。所以先整体聚合， - /// 再按「有没有任何一轮落在窗口里」留下整条。 + /// 问的是整次任务，不是它落在某个窗口里的那一段。所以留下的会话整条聚合， + /// 按「首尾之间和窗口有没有重叠」留。 + /// + /// **先挑会话，再聚合挑中的那几次**（[`Db::recent_sessions`]）。以前把三个月里每一行按 + /// 会话分组、排序之后再截前几百条：二十七万行要两百毫秒，而界面每结束一个请求就来要一次。 pub fn sessions( &self, within: Option<(i64, i64)>, limit: usize, ) -> Result, DbError> { + let picked = self.recent_sessions(within, limit)?; + if picked.is_empty() { + return Ok(Vec::new()); + } + // 挑中的会话号作为一个 JSON 数组交进去：几百个占位符拼不出一条好读的 SQL + let picked = serde_json::Value::from(picked).to_string(); let mut st = self.conn.prepare(&format!( - "SELECT session, - client, - MIN(at_ms), MAX(at_ms), COUNT(*), - COALESCE(SUM(cost_micros), 0), - COALESCE(SUM({NO_PRICE}), 0), - COALESCE(SUM(input_tokens), 0), - COALESCE(SUM(output_tokens), 0), - COALESCE(SUM(cache_read_tokens), 0), - COALESCE(SUM(cache_write_tokens), 0), - COALESCE(SUM(cache_saved_micros), 0), - COALESCE(MAX(input_tokens), 0), - GROUP_CONCAT(DISTINCT model), - SUM(CASE WHEN error IS NOT NULL THEN 1 ELSE 0 END), - COALESCE(SUM(CASE WHEN cost_estimated = 1 THEN cost_micros ELSE 0 END), 0), - COUNT(cost_micros), - COALESCE(SUM({NO_USAGE}), 0) - FROM requests - WHERE session IS NOT NULL AND local = 0 + "SELECT {} FROM requests + WHERE session IN (SELECT value FROM json_each(?1)) AND local = 0 GROUP BY session - HAVING MAX(at_ms) >= ?1 AND MIN(at_ms) <= ?2 - ORDER BY MAX(at_ms) DESC - LIMIT ?3" + ORDER BY MAX(at_ms) DESC, session DESC", + session_columns() ))?; - let (from, to) = within.unwrap_or((i64::MIN, i64::MAX)); - let rows = st.query_map([from, to, limit as i64], |r| { - Ok(SessionRow { - id: r.get(0)?, - client: r.get(1)?, - started_ms: r.get(2)?, - ended_ms: r.get(3)?, - turns: r.get(4)?, - cost_micros: r.get(5)?, - unpriced_turns: r.get(6)?, - input_tokens: r.get(7)?, - output_tokens: r.get(8)?, - cache_read_tokens: r.get(9)?, - cache_write_tokens: r.get(10)?, - cache_saved_micros: r.get(11)?, - peak_input_tokens: r.get(12)?, - models: r.get::<_, Option>(13)?.unwrap_or_default(), - errors: r.get(14)?, - cost_micros_estimated: r.get(15)?, - priced_turns: r.get(16)?, - no_usage_turns: r.get(17)?, - }) - })?; + let rows = st.query_map([picked], session_row)?; Ok(rows.collect::, _>>()?) } + /// 一次会话的汇总。没有这次会话（或者它的每一轮都是本地应答的）是 None。 + /// + /// **按会话号直接取**，不从列表里找：以前取最近五百次会话再从里面挑，更早的会话点开是 404。 + pub fn session(&self, id: &str) -> Result, DbError> { + Ok(self + .conn + .prepare(&format!( + "SELECT {} FROM requests WHERE session = ?1 AND local = 0 GROUP BY session", + session_columns() + ))? + .query_row([id], session_row) + .optional()?) + } + + /// 最近活动过的 `limit` 次会话，最后一次活动最近的在前；`within` 给了就只要首尾之间和它 + /// 重叠的。 + /// + /// **只走索引，不按会话分组整张表**：按时间倒着走（`requests_session_at`，只有带会话、 + /// 不是本地应答的行），一次会话第一次碰到的那一行就是它最后一次活动，碰够了就停。窗口的 + /// 起点之前不必再走：最后一次活动比它早的会话不和窗口重叠。最后一次活动在窗口终点之后的， + /// 再按会话号取它的第一轮（`requests_session`）看是不是不晚于终点。 + fn recent_sessions( + &self, + within: Option<(i64, i64)>, + limit: usize, + ) -> Result, DbError> { + let (from, to) = within.unwrap_or((i64::MIN, i64::MAX)); + let mut out = Vec::new(); + if limit == 0 { + return Ok(out); + } + let mut walk = self.conn.prepare( + "SELECT session, at_ms FROM requests + WHERE session IS NOT NULL AND local = 0 AND at_ms >= ?1 + ORDER BY at_ms DESC, session DESC", + )?; + let mut first = self.conn.prepare( + "SELECT at_ms FROM requests WHERE session = ?1 AND local = 0 + ORDER BY at_ms LIMIT 1", + )?; + let mut seen = std::collections::HashSet::new(); + let mut rows = walk.query([from])?; + while let Some(r) = rows.next()? { + let session = r.get_ref(0)?.as_str().map_err(rusqlite::Error::from)?; + if seen.contains(session) { + continue; + } + let last: i64 = r.get(1)?; + let overlaps = last <= to || first.query_row([session], |r| r.get::<_, i64>(0))? <= to; + seen.insert(session.to_string()); + if overlaps { + out.push(session.to_string()); + if out.len() == limit { + break; + } + } + } + Ok(out) + } + /// 一次会话里的每一轮，**按时间正序** —— 曲线是从左往右画的。 pub fn turns(&self, session: &str) -> Result, DbError> { let mut st = self.conn.prepare( @@ -672,10 +795,22 @@ impl Db { /// 只有带着为那个客户端生成的密钥的请求能证明。按请求头里自报的客户端 /// 标识分组的话，那个标识谁都能写 —— 「这把钥匙还有没有人在用」「接好了 /// 没有」都只能按密钥问。 + /// + /// **一把一把地跳着取**（`requests_client`）：取下一把比上一把大的密钥名，再取它最大的 + /// 时刻，各是一次索引查找。按密钥分组的话要把三个月里每一行过一遍，而密钥只有几把。 + /// 每个子查询都写着 `client > ''`：那个索引只收这些行，条件里没有它就用不上。 pub fn last_seen_by_client(&self) -> Result, DbError> { let mut st = self.conn.prepare( - "SELECT client, MAX(at_ms) FROM requests - WHERE client <> '' GROUP BY client", + "WITH RECURSIVE keys(client) AS ( + SELECT MIN(client) FROM requests WHERE client > '' + UNION ALL + SELECT (SELECT MIN(client) FROM requests WHERE client > keys.client AND client > '') + FROM keys WHERE keys.client IS NOT NULL + ) + SELECT client, + (SELECT MAX(at_ms) FROM requests + WHERE requests.client = keys.client AND client > '') + FROM keys WHERE client IS NOT NULL", )?; let rows = st.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?; Ok(rows.collect::, _>>()?) @@ -1311,15 +1446,10 @@ impl Db { /// **按每一行记下的路由算**，不按现在的配置推：请求走的是它那一刻的路由和 /// 规则。一个请求算在这几条规则上，每条只算一次：决定去向的那一条、附加了 /// 改写的每一条、选定上游之后拒绝了它的那一条。本地应答的没有路由，不算。 + /// + /// **只读索引**（`requests_routed`）：要的那几项落库时就从路由那一列里取出来存在索引里， + /// 路由图开着时每十秒来一次，不再把一周几 MB 的 JSON 读出来一行一行地解。 fn route_hits(&self, since_ms: i64, until_ms: i64) -> Result, DbError> { - /// 数命中要的那几项。尝试链不用解 - #[derive(serde::Deserialize)] - struct Routed { - route: String, - rule: String, - rewritten_by: Vec, - denied_by: Option, - } #[derive(Default)] struct Tally { requests: i64, @@ -1334,44 +1464,52 @@ impl Db { self.last_ms = self.last_ms.max(at_ms); } } - let mut st = self.conn.prepare( - "SELECT at_ms, error IS NOT NULL, routing FROM requests - WHERE at_ms >= ?1 AND at_ms < ?2 AND local = 0 AND routing IS NOT NULL", - )?; - let rows = st.query_map(params![since_ms, until_ms], |r| { - Ok(( - r.get::<_, i64>(0)?, - r.get::<_, bool>(1)?, - r.get::<_, String>(2)?, - )) - })?; - let mut routes: std::collections::BTreeMap< - String, - (Tally, std::collections::BTreeMap), - > = Default::default(); - for row in rows { - let (at_ms, failed, json) = row?; - // 解不开的那一行不算。**一条坏掉的记录不该让整张表出不来** - let Ok(r) = serde_json::from_str::(&json) else { + let mut st = self.conn.prepare(&format!( + "SELECT at_ms, error IS NOT NULL, {ROUTE}, {RULE}, {REWRITTEN_BY}, {DENIED_BY} + FROM requests + WHERE at_ms >= ?1 AND at_ms < ?2 AND {ROUTED}" + ))?; + let mut routes: Named<(Tally, Named)> = Named::default(); + let mut rows = st.query(params![since_ms, until_ms])?; + while let Some(r) = rows.next()? { + let at_ms: i64 = r.get(0)?; + let failed: bool = r.get(1)?; + // 缺了哪一项的那一行不算。**一条坏掉的记录不该让整张表出不来** + let (Ok(route), Ok(rule), Ok(rewritten_by), Ok(denied_by)) = ( + r.get_ref(2)?.as_str(), + r.get_ref(3)?.as_str(), + r.get_ref(4)?.as_str(), + r.get_ref(5)?.as_str_or_null(), + ) else { continue; }; - let (total, rules) = routes.entry(r.route).or_default(); + // 绝大多数请求没有附加改写，不必解 + let rewritten_by: Vec = match rewritten_by { + "[]" => Vec::new(), + j => match serde_json::from_str(j) { + Ok(names) => names, + Err(_) => continue, + }, + }; + let (total, rules) = routes.get(route); total.add(at_ms, failed); - let decider = rules.entry(r.rule.clone()).or_default(); + let decider = rules.get(rule); decider.add(at_ms, failed); decider.decided += 1; - let mut counted = vec![r.rule]; - for name in r.rewritten_by.into_iter().chain(r.denied_by) { + let mut counted = vec![rule]; + for name in rewritten_by.iter().map(String::as_str).chain(denied_by) { if !counted.contains(&name) { - rules.entry(name.clone()).or_default().add(at_ms, failed); + rules.get(name).add(at_ms, failed); counted.push(name); } } } let mut out: Vec = routes + .items .into_iter() .map(|(route, (total, rules))| { let mut rules: Vec = rules + .items .into_iter() .map(|(rule, t)| tw_api::RuleHits { rule, @@ -1412,6 +1550,36 @@ impl Db { } } +/// 按名字数的一组东西。**名字只在第一次见到时拷一份**：数路由命中时一周两万行，每行都按 +/// 名字找一遍（见 [`Db::route_hits`]） +struct Named { + at: std::collections::HashMap, + items: Vec<(String, T)>, +} + +impl Default for Named { + fn default() -> Self { + Self { + at: Default::default(), + items: Vec::new(), + } + } +} + +impl Named { + fn get(&mut self, name: &str) -> &mut T { + let i = match self.at.get(name) { + Some(&i) => i, + None => { + self.at.insert(name.to_string(), self.items.len()); + self.items.push((name.to_string(), T::default())); + self.items.len() - 1 + } + }; + &mut self.items[i].1 + } +} + /// 一个插件在一个请求上的一次运行，落库的样子（`plugin_runs` 一行，少了 `seq`： /// 它在写入时按这个请求已有的行数定）。 #[derive(Debug, Clone, PartialEq)] @@ -2014,6 +2182,87 @@ pub(crate) mod tests { assert!(db.sessions(Some((20_000, 30_000)), 10).unwrap().is_empty()); } + /// 列表是最后一次活动最近的那几次会话，每一次都整条聚合 —— 挑会话时只看了最近的那几行， + /// 聚合的却是它的每一轮。窗口按首尾和它重叠不重叠留：最后一次活动在窗口之后、开始在窗口 + /// 之内的留，整个在窗口之后、之前的不留 + #[test] + fn the_list_is_the_most_recently_active_sessions_each_counted_whole() { + let db = Db::in_memory().unwrap(); + // （会话，时刻）：a 最早开始、最后结束；b 整个在中间；c 在 a 的两头之间断断续续 + for (id, (session, at)) in [ + ("a", 1_000), + ("c", 2_000), + ("b", 3_000), + ("b", 4_000), + ("c", 5_000), + ("c", 6_000), + ("a", 9_000), + ] + .into_iter() + .enumerate() + { + let mut r = row(id as i64 + 1, at); + r.session = Some(session.into()); + db.insert(&r).unwrap(); + } + // 本地应答的不算，哪怕它最新 + let mut probe = row(99, 99_000); + probe.session = Some("b".into()); + probe.local = true; + db.insert(&probe).unwrap(); + + let ids = |got: Vec| got.into_iter().map(|s| s.id).collect::>(); + let two = db.sessions(None, 2).unwrap(); + assert_eq!( + two.iter() + .map(|s| (s.id.as_str(), s.turns)) + .collect::>(), + [("a", 2), ("c", 3)] + ); + assert_eq!(two[0].started_ms, 1_000); + assert_eq!(ids(db.sessions(None, 10).unwrap()), ["a", "c", "b"]); + assert!(db.sessions(None, 0).unwrap().is_empty()); + + // 窗口 [3500, 4500]：三次都和它重叠（a、c 跨过它，b 在里面） + assert_eq!( + ids(db.sessions(Some((3_500, 4_500)), 10).unwrap()), + ["a", "c", "b"] + ); + // [0, 1500]：只有 a 开始得那么早 + assert_eq!(ids(db.sessions(Some((0, 1_500)), 10).unwrap()), ["a"]); + // [6500, 8000]：b、c 那时都结束了，a 还没结束 + assert_eq!(ids(db.sessions(Some((6_500, 8_000)), 10).unwrap()), ["a"]); + // 窗口在所有会话之后 + assert!(db.sessions(Some((10_000, 20_000)), 10).unwrap().is_empty()); + } + + /// 会话详情按会话号直接取：多老、前面压着多少次会话都取得到。以前从最近五百次里找， + /// 更早的点开是 404 + #[test] + fn one_session_is_found_by_its_id_however_old() { + let db = Db::in_memory().unwrap(); + for i in 0..600i64 { + let mut r = row(i + 1, 1_000 + i); + r.session = Some(format!("s{i}")); + db.insert(&r).unwrap(); + } + let mut second = row(1_000, 5_000); + second.session = Some("s0".into()); + db.insert(&second).unwrap(); + + let s0 = db.session("s0").unwrap().unwrap(); + assert_eq!((s0.turns, s0.started_ms, s0.ended_ms), (2, 1_000, 5_000)); + assert_eq!(db.session("s599").unwrap().unwrap().turns, 1); + assert_eq!(db.session("nope").unwrap(), None); + + // 只有本地应答的不算一次会话，和列表一样 + let mut probe = row(2_000, 6_000); + probe.session = Some("probes".into()); + probe.local = true; + db.insert(&probe).unwrap(); + assert_eq!(db.session("probes").unwrap(), None); + } + #[test] fn the_last_request_id_is_where_the_next_run_has_to_start() { let db = Db::in_memory().unwrap(); @@ -2489,6 +2738,10 @@ pub(crate) mod tests { r.client_hint = hint.map(|s| s.to_string()); db.insert(&r).unwrap(); } + // 没有密钥的（本地来的、旧记录）不算一把 + let mut keyless = row(9, 900); + keyless.client = String::new(); + db.insert(&keyless).unwrap(); let mut got = db.last_seen_by_client().unwrap(); got.sort(); assert_eq!( @@ -2497,6 +2750,37 @@ pub(crate) mod tests { ); } + /// 另加的索引不算 schema：同一版本、还没有这些索引的库照样打开，行都在，缺的索引补上 + /// —— 加 SCHEMA 的话，每个人的请求历史都要清掉 + #[test] + fn a_database_without_the_extra_indexes_keeps_its_rows_and_gains_them() { + const EXTRA: [&str; 3] = ["requests_client", "requests_session_at", "requests_routed"]; + let d = tempfile::tempdir().unwrap(); + let path = d.path().join("data.db"); + { + let db = Db::open(&path).unwrap(); + db.insert(&row(1, 100)).unwrap(); + for name in EXTRA { + db.conn.execute(&format!("DROP INDEX {name}"), []).unwrap(); + } + } + let db = Db::open(&path).unwrap(); + assert_eq!(db.count().unwrap(), 1); + let names: Vec = db + .conn + .prepare( + "SELECT name FROM sqlite_master WHERE type = 'index' AND tbl_name = 'requests'", + ) + .unwrap() + .query_map([], |r| r.get(0)) + .unwrap() + .collect::>() + .unwrap(); + for name in EXTRA { + assert!(names.iter().any(|n| n == name), "{name}: {names:?}"); + } + } + #[test] fn a_database_from_another_version_is_refused_rather_than_read() { // 不迁移，也不硬读：版本对不上就说出来，怎么处理由 `crate::open` 定 @@ -3193,6 +3477,27 @@ mod route_stats_tests { } /// 窗口外的、本地应答的、没有路由的不算；**一条解不开的记录不挡住别的** + /// 路由那一列解不开的那一行照样写得进去（索引不收它），数命中时不算。**一条坏掉的记录 + /// 不该让整张表出不来，也不该让它自己写不进去** + #[test] + fn a_row_whose_routing_does_not_parse_is_kept_and_left_out_of_the_count() { + let db = Db::in_memory().unwrap(); + let mut broken = row(1, 100); + broken.routing = Some("{\"route\": ".into()); + db.insert(&broken).unwrap(); + // 解得开、却缺了必有的一项的也不算 + let mut partial = row(2, 150); + partial.routing = Some(r#"{"route":"默认"}"#.into()); + db.insert(&partial).unwrap(); + db.insert(&routed(3, 200, "默认", "兜底", &[], None)) + .unwrap(); + assert_eq!(db.count().unwrap(), 3); + let got = db.route_stats(0, 1_000).unwrap(); + assert_eq!(got.routes.len(), 1, "{got:?}"); + assert_eq!(got.routes[0].requests, 1); + assert_eq!(hits(&got.routes[0], "兜底").decided, 1); + } + #[test] fn only_routed_requests_inside_the_window_count() { let db = Db::in_memory().unwrap(); @@ -3286,3 +3591,256 @@ mod route_stats_tests { assert_eq!(db.route_stats(700, 700).unwrap().covered_since_ms, None); } } + +/// 一份和真实库一样大的请求库上，界面常问的那几条查询各花多久。 +/// +/// 默认不跑：要先写二十七万行。发布构建下跑（`TW_BENCH_DB` 给一个路径的话库留在那里， +/// 下次接着用）： +/// `cargo test --release -p tw-store history_cost -- --ignored --nocapture` +#[cfg(test)] +mod cost { + use super::tests::row; + use super::*; + + /// 九十天、每天三千条，和留得最久的记录一样长 + const ROWS: i64 = 270_000; + const NOW: i64 = 1_790_000_000_000; + const DAY: i64 = 86_400_000; + + /// 不引随机数的库：一个线性同余就够造数据 + struct Lcg(u64); + impl Lcg { + fn next(&mut self) -> u64 { + self.0 = self + .0 + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + self.0 >> 33 + } + fn below(&mut self, n: u64) -> u64 { + self.next() % n + } + } + + const CLIENTS: &[&str] = &["claude-code", "codex", "cursor", "opencode", "gemini", "pi"]; + const ROUTES: &[&str] = &["默认", "工作", "便宜"]; + const RULES: &[&str] = &[ + "兜底", + "opus 走中转", + "haiku 本地", + "长上下文", + "关思考", + "拒绝 gpt-4", + ]; + const PROVIDERS: &[&str] = &["官方", "中转 A", "中转 B", "OpenRouter", "Bedrock"]; + const MODELS: &[&str] = &[ + "claude-sonnet-4-5", + "claude-opus-4-1", + "gpt-5-codex", + "gemini-2.5-pro", + ]; + + fn routing(g: &mut Lcg) -> String { + let hops = if g.below(10) < 4 { 2 + g.below(2) } else { 1 }; + let attempts = (0..hops) + .map(|i| { + let last = i + 1 == hops; + tw_api::AttemptView { + provider: PROVIDERS[g.below(PROVIDERS.len() as u64) as usize].into(), + model: (g.below(4) == 0).then(|| "claude-sonnet-4-5-20250929".into()), + outcome: if last { + tw_api::AttemptOutcome::Served + } else { + tw_api::AttemptOutcome::Error + }, + status: last.then_some(200), + error: (!last).then(|| Msg { + code: "gw.upstream.timeout".into(), + args: Default::default(), + text: "The upstream answered 529 Overloaded: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"},\"request_id\":\"req_011CTxxxxxxxxxxxxxxxxxxx\"}; trying the next one." + .into(), + }), + ms: 200 + g.below(30_000), + usage: None, + queued_ms: None, + skipped: None, + } + }) + .collect(); + serde_json::to_string(&tw_api::RoutingView { + route: ROUTES[g.below(ROUTES.len() as u64) as usize].into(), + rule: RULES[g.below(RULES.len() as u64) as usize].into(), + group: Some("__all__".into()), + rewritten_by: (0..g.below(3)) + .map(|_| RULES[g.below(RULES.len() as u64) as usize].to_string()) + .collect(), + denied_by: (g.below(50) == 0).then(|| RULES[5].into()), + affinity: Some(tw_api::AffinityView { + held_route: g.below(2) == 0, + stayed: None, + }), + attempts, + }) + .unwrap() + } + + /// 造一份：会话两三个交错着进行，一次几轮到几百轮；少数请求认不出会话（WebSocket） + /// 或者是网关自己答的 + fn fill(db: &Db, rows: i64) { + let mut g = Lcg(7); + // （会话号，密钥，还剩几轮） + let mut open: Vec<(String, &str, u64)> = Vec::new(); + db.conn.execute_batch("BEGIN").unwrap(); + for id in 1..=rows { + let at = NOW - 90 * DAY + id * (90 * DAY / rows); + let mut r = row(id, at); + r.price_source = Some(r#"{"kind":"default","date":"2026-09-30"}"#.into()); + r.model = MODELS[g.below(MODELS.len() as u64) as usize].into(); + r.sent_model = r.model.clone(); + r.provider = PROVIDERS[g.below(PROVIDERS.len() as u64) as usize].into(); + match g.below(100) { + 0..=2 => { + r.local = true; + r.path = "count_tokens".into(); + r.client = CLIENTS[g.below(CLIENTS.len() as u64) as usize].into(); + } + 3..=9 => { + r.client = CLIENTS[g.below(CLIENTS.len() as u64) as usize].into(); + r.routing = Some(routing(&mut g)); + } + _ => { + while open.len() < 3 { + let len = if g.below(20) == 0 { + 300 + g.below(400) + } else { + 1 + g.below(120) + }; + open.push(( + format!("{:012x}-{at}", g.next()), + CLIENTS[g.below(CLIENTS.len() as u64) as usize], + len, + )); + } + let i = g.below(open.len() as u64) as usize; + r.session = Some(open[i].0.clone()); + r.client = open[i].1.into(); + r.routing = Some(routing(&mut g)); + open[i].2 -= 1; + if open[i].2 == 0 { + open.swap_remove(i); + } + } + } + if g.below(30) == 0 { + r.error = Some(super::tests::upstream_failed("upstream returned 529")); + } + db.insert(&r).unwrap(); + } + db.conn.execute_batch("COMMIT").unwrap(); + } + + /// 以前的会话列表：整张表按会话分组，排序之后截前几条。新的那条要和它一模一样 + fn grouping_everything(db: &Db, within: Option<(i64, i64)>, limit: usize) -> Vec { + let (from, to) = within.unwrap_or((i64::MIN, i64::MAX)); + db.conn + .prepare(&format!( + "SELECT {} FROM requests + WHERE session IS NOT NULL AND local = 0 + GROUP BY session + HAVING MAX(at_ms) >= ?1 AND MIN(at_ms) <= ?2 + ORDER BY MAX(at_ms) DESC, session DESC + LIMIT ?3", + session_columns() + )) + .unwrap() + .query_map([from, to, limit as i64], session_row) + .unwrap() + .collect::>() + .unwrap() + } + + /// 挑会话再聚合，和整张表分组再截前几条，在各种窗口、各种条数下给的是同一张表 + #[test] + fn the_list_is_what_grouping_every_row_gives() { + let db = Db::in_memory().unwrap(); + fill(&db, 6_000); + for limit in [1, 7, 50, 2_000] { + for within in [ + None, + Some((NOW - 7 * DAY, NOW)), + Some((NOW - 31 * DAY, NOW - 30 * DAY)), + Some((NOW - 60 * DAY, NOW - 59 * DAY + 1)), + Some((NOW + DAY, NOW + 2 * DAY)), + Some((i64::MIN, NOW - 89 * DAY)), + ] { + assert_eq!( + db.sessions(within, limit).unwrap(), + grouping_everything(&db, within, limit), + "{within:?} {limit}" + ); + } + } + } + + fn time(what: &str, mut f: impl FnMut() -> T) -> T { + let mut out = f(); + let mut best = std::time::Duration::MAX; + for _ in 0..5 { + let t0 = std::time::Instant::now(); + out = f(); + best = best.min(t0.elapsed()); + } + eprintln!("{what:<48} {best:>12.2?}"); + out + } + + #[test] + #[ignore] + fn history_cost() { + let dir = tempfile::tempdir().unwrap(); + let path = std::env::var_os("TW_BENCH_DB") + .map(std::path::PathBuf::from) + .unwrap_or_else(|| dir.path().join("data.db")); + let t0 = std::time::Instant::now(); + let db = Db::open(&path).unwrap(); + eprintln!( + "{:<48} {:>12.2?}", + "open (builds missing indexes)", + t0.elapsed() + ); + if db.count().unwrap() < ROWS { + fill(&db, ROWS); + } + let size = std::fs::metadata(&path).map(|m| m.len()).unwrap_or(0); + eprintln!( + "{} rows, {} MB on disk", + db.count().unwrap(), + size / 1_000_000 + ); + for (what, within) in [ + ("sessions (no window, 200)", None), + ("sessions (last 7 days, 200)", Some((NOW - 7 * DAY, NOW))), + ( + "sessions (a day a month ago, 200)", + Some((NOW - 31 * DAY, NOW - 30 * DAY)), + ), + ] { + let got = time(what, || db.sessions(within, 200).unwrap()); + let grouped = time(" the same by grouping every row", || { + grouping_everything(&db, within, 200) + }); + assert_eq!(got, grouped, "{what}"); + } + // 会话详情要的那一次：最老的一次会话 + let oldest = db.sessions(None, 2000).unwrap().pop().unwrap(); + let one = time("one session (oldest of 2000)", || { + db.session(&oldest.id).unwrap() + }); + assert_eq!(one, Some(oldest.clone())); + time("one session's turns", || db.turns(&oldest.id).unwrap()); + time("last_seen_by_client", || db.last_seen_by_client().unwrap()); + time("route_stats (last 7 days)", || { + db.route_stats(NOW - 7 * DAY, NOW).unwrap() + }); + } +} diff --git a/crates/tw-store/src/recorder.rs b/crates/tw-store/src/recorder.rs index 42f97333..bf077d62 100644 --- a/crates/tw-store/src/recorder.rs +++ b/crates/tw-store/src/recorder.rs @@ -93,6 +93,9 @@ pub struct Recorder { /// /// `None` 表示没人要听（测试、以及不带总线的调用方）。 bus: Option, + /// 最近几次会话读成对话读到了哪儿（见 [`crate::transcript::Cache`]）。正文回收删了东西 + /// 就作废（[`crate::task::gc`]） + transcripts: std::sync::Arc, /// 每记下一行请求就交一份 [`Settled`] 出去：密钥的用量上限拿它把预留换成实数。 /// /// **直接调，不走总线。**总线上丢了事件的话，这一行也就不在库里，重启之后从库里 @@ -153,6 +156,7 @@ impl Recorder { pricing, inflight: HashMap::new(), bus: None, + transcripts: Default::default(), settled: None, } } @@ -230,6 +234,12 @@ impl Recorder { &self.blobs } + /// 读对话记录用的那一份记忆（见 [`crate::transcript::Cache`]）。读的时候不拿这把锁： + /// 几百轮的会话要读几百份正文 + pub fn transcripts(&self) -> std::sync::Arc { + self.transcripts.clone() + } + /// 记一次请求体或响应体。 /// /// **在事件之外单独走** —— body 不进事件流：那是个广播通道，每个 @@ -808,16 +818,15 @@ impl Recorder { } } - /// 定期回收。**返回删了多少字节**，调用方记一行日志就够了。 - pub fn gc(&self, now_ms: i64, keep_days: u64, metadata_keep_days: u64, max_bytes: u64) -> u64 { - let freed = self.blobs.gc(now_ms, keep_days, max_bytes); - let cutoff = now_ms - (metadata_keep_days as i64) * 86_400_000; + /// 删掉过期的记录行（`keep_days` 天以前的）。正文的回收不在这里：它不需要这把锁，见 + /// [`crate::task::gc`]。 + pub fn prune(&self, now_ms: i64, keep_days: u64) { + let cutoff = now_ms - (keep_days as i64) * 86_400_000; match self.db.prune_before(cutoff) { Ok(n) if n > 0 => tracing::info!(rows = n, "cleaned up expired request rows"), Err(e) => tracing::debug!("the request rows could not be cleaned up: {e}"), _ => {} } - freed } } @@ -2327,8 +2336,13 @@ mod cancellation_tests { r.db().insert(&crate::db::tests::row(1, old)).unwrap(); r.record_body(old, 1, Which::Request, b"an old body", 11); + let shared = tokio::sync::Mutex::new(r); + let gc = |keep_days, metadata_keep_days| { + crate::task::gc(&shared, now, keep_days, metadata_keep_days, u64::MAX) + }; // 正文留 7 天、记录留 90 天：正文该没了，记录还在 - r.gc(now, 7, 90, u64::MAX); + gc(7, 90); + let r = shared.blocking_lock(); assert_eq!( r.db().recent(None, 10).unwrap().len(), 1, @@ -2339,9 +2353,17 @@ mod cancellation_tests { "正文过了它自己的期限却还在" ); + drop(r); // 记录的期限也到了才删行 - r.gc(now, 7, 7, u64::MAX); - assert!(r.db().recent(None, 10).unwrap().is_empty()); + gc(7, 7); + assert!( + shared + .blocking_lock() + .db() + .recent(None, 10) + .unwrap() + .is_empty() + ); } /// 总量上限是给突发准备的：按天算出来的占用取决于用量，而用量会有 @@ -2358,9 +2380,10 @@ mod cancellation_tests { let before = r.blobs().total_bytes(); assert!(before > 0); // 一天都没过期，但总量超了 —— 从最旧的整天开始删 - let freed = r.gc(now, 30, 30, 4096); + let shared = tokio::sync::Mutex::new(r); + let freed = crate::task::gc(&shared, now, 30, 30, 4096); assert!(freed > 0, "总量超了却什么都没删"); - assert!(r.blobs().total_bytes() < before); + assert!(shared.blocking_lock().blobs().total_bytes() < before); } /// 会话里的那一轮也要看得出是取消的。否则在每轮花费里，它就是一轮 diff --git a/crates/tw-store/src/task.rs b/crates/tw-store/src/task.rs index bfd6d2fd..e59fbd96 100644 --- a/crates/tw-store/src/task.rs +++ b/crates/tw-store/src/task.rs @@ -1,15 +1,22 @@ -//! 起一个后台任务，把事件流写进库里。 +//! 起几个后台线程，把事件流写进库里、把正文写到盘上。 //! -//! **它订阅事件总线，而不是被数据面直接调用。**这条边界是刻意的：写库 +//! **它收事件总线上的事件，而不是被数据面直接调用。**这条边界是刻意的：写库 //! 慢了、库锁住了、磁盘满了，转发这条路上一行代码都不会等它。 +//! +//! **跑在自己的线程上，不在异步运行时里。**写库、写几 MB 的正文、`chmod` 都是会挡住线程的 +//! 调用：放在运行时的工作线程上，同一个线程上排着的转发和控制面请求都得等它。 use std::sync::Arc; use tokio::sync::Mutex; -use crate::blobs::Which; +use crate::blobs::{Blobs, Which}; use crate::recorder::Recorder; +/// 记录线程拿着锁一次最多吃几条事件。**拿着锁的时候查询在等**：一条几十微秒，一批也就 +/// 几毫秒；一条一拿又白白多出许多次抢锁 +const BATCH: usize = 64; + /// 一份等着落盘的 body。 /// /// **这个类型属于 store，不属于网关。**让 store 依赖网关会把「观测」 @@ -29,14 +36,14 @@ pub struct StoredBody { } /// 把插件的运行记录写进库里，**和正文同一个道理**：数据面只管交出去，写库在这条 -/// 任务上。记录和请求那一行各走各的路，落库不分先后（见 `Db::insert_plugin_run`）。 +/// 线程上。记录和请求那一行各走各的路，落库不分先后（见 `Db::insert_plugin_run`）。 pub fn record_plugin_runs( recorder: Arc>, mut runs: tokio::sync::mpsc::Receiver, ) { - tokio::spawn(async move { - while let Some(r) = runs.recv().await { - recorder.lock().await.record_plugin_run(&r); + thread("tw-plugin-runs", move || { + while let Some(r) = runs.blocking_recv() { + recorder.blocking_lock().record_plugin_run(&r); } }); } @@ -44,38 +51,143 @@ pub fn record_plugin_runs( /// 起来。返回的 handle 给别的地方查历史用 —— **同一个 Recorder**， /// 不是第二个连接：两个连接会让「刚写进去的还查不到」变成可能。 /// -/// **回收不在这里。**它按配置里的 `retention` 走（`twcore` 每小时读一次配置，调 -/// [`Recorder::gc`]）。这里以前另跑着一个按写死的 7 天、90 天、2 GB 回收的循环，和按 -/// 配置的那个同时跑：配得比出厂更宽的期限和上限，每小时被它削回出厂值一次。 +/// 事件从总线给存储层的那条通道来（[`tw_observe::EventBus::record_feed`]），**不订阅界面的 +/// 广播**：那是个一千来格的环，界面落后时存储层多半也落后，丢掉的那几个请求永远不在库里。 +/// +/// **正文不拿锁**：正文目录只是一个路径，写一份 4 MB 的正文时，记录和查询不必等它。 +/// +/// **回收不在这里。**它按配置里的 `retention` 走（`twcore` 每小时读一次配置，调 [`gc`]）。 +/// 这里以前另跑着一个按写死的 7 天、90 天、2 GB 回收的循环，和按配置的那个同时跑：配得比出厂 +/// 更宽的期限和上限，每小时被它削回出厂值一次。 pub fn spawn( recorder: Recorder, - mut rx: tokio::sync::broadcast::Receiver, + mut events: tw_observe::RecordFeed, mut bodies: tokio::sync::mpsc::Receiver, ) -> Arc> { + let blobs = Blobs::new(recorder.blobs().root().to_path_buf()); let shared = Arc::new(Mutex::new(recorder)); - let r = shared.clone(); - tokio::spawn(async move { - while let Some(b) = bodies.recv().await { - // **写盘在这条任务上，不在转发那条路上。**慢磁盘只会让 - // 通道积压然后丢，不会让请求变慢。 - r.lock() - .await - .record_body(b.at_ms, b.id, b.which, &b.body, b.original_len); + thread("tw-bodies", move || { + // 慢磁盘只会让通道积压然后丢，不会让请求变慢 + while let Some(b) = bodies.blocking_recv() { + blobs.put_with_len(b.at_ms, b.id as i64, b.which, &b.body, b.original_len); } }); let r = shared.clone(); - tokio::spawn(async move { - loop { - match rx.recv().await { - Ok(ev) => r.lock().await.on_event(&ev), - // **落后了就跳过，不要退出。**广播通道满的时候丢的是 - // 观测记录，而退出会让之后的记录全部丢失。 - Err(tokio::sync::broadcast::error::RecvError::Lagged(n)) => { - tracing::warn!("recording fell behind; {n} requests were not recorded"); + thread("tw-recorder", move || { + while let Some(ev) = events.blocking_recv() { + { + let mut g = r.blocking_lock(); + g.on_event(&ev); + for _ in 1..BATCH { + let Some(ev) = events.try_recv() else { break }; + g.on_event(&ev); } - Err(tokio::sync::broadcast::error::RecvError::Closed) => return, + } + // **落后了就丢，但要说出来。**通道满了丢的是观测记录，而等着它会拖慢转发 + let dropped = events.take_dropped(); + if dropped > 0 { + tracing::warn!( + dropped, + "recording fell behind; {dropped} events were not recorded" + ); } } }); shared } + +/// 回收一轮：过期的正文、超出总量的正文、过期的记录行。返回删掉了多少字节的正文。 +/// +/// **删正文的时候不拿锁**：正文目录只是一个路径，一次删几个 GB 可以要好几秒，而那期间记录 +/// 和查询都要这把锁。删行要拿：行在库里，一小时一轮删的只是这一小时过期的那些。 +/// +/// 会挡住当前线程：在阻塞线程上调。 +pub fn gc( + shared: &Mutex, + now_ms: i64, + keep_days: u64, + metadata_keep_days: u64, + max_bytes: u64, +) -> u64 { + let (blobs, transcripts) = { + let g = shared.blocking_lock(); + (Blobs::new(g.blobs().root().to_path_buf()), g.transcripts()) + }; + let freed = blobs.gc(now_ms, keep_days, max_bytes); + // 删掉的正文可能在记着的对话里：记着的那些作废，下次重读 + if freed > 0 { + transcripts.invalidate(); + } + shared.blocking_lock().prune(now_ms, metadata_keep_days); + freed +} + +/// 起一条有名字的线程。**起不来只记一行**：那时这一路什么都不记，转发照常 +fn thread(name: &str, f: impl FnOnce() + Send + 'static) { + if let Err(e) = std::thread::Builder::new().name(name.into()).spawn(f) { + tracing::warn!("the {name} thread could not start, so it records nothing this run: {e}"); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::transcript::Ask; + + const NOW: i64 = 1_790_000_000_000; + const DAY: i64 = 86_400_000; + + /// 回收删了正文，记着的对话跟着作废：之后读出来的那一轮说出缺了什么，而不是删掉之前的 + /// 样子 + #[test] + fn reclaiming_bodies_forgets_the_transcripts_read_from_them() { + let d = tempfile::tempdir().unwrap(); + let rec = Recorder::new( + crate::Db::in_memory().unwrap(), + Blobs::new(d.path().join("blobs")), + tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), + ); + let mut row = crate::db::tests::row(1, NOW); + row.session = Some("s".into()); + rec.db().insert(&row).unwrap(); + let request = + br#"{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hi"}]}"#; + let answer = br#"{"type":"message","content":[{"type":"text","text":"hello"}]}"#; + rec.record_body(NOW, 1, Which::Request, request, request.len()); + rec.record_body(NOW, 1, Which::Response, answer, answer.len()); + let shared = Mutex::new(rec); + let read = || { + let (rows, blobs, transcripts) = { + let g = shared.blocking_lock(); + ( + g.db().session_requests("s").unwrap(), + Blobs::new(g.blobs().root().to_path_buf()), + g.transcripts(), + ) + }; + transcripts.read( + &blobs, + &Ask { + session: "s", + rows: &rows, + from_turn: 0, + running: None, + now_ms: NOW + 3 * DAY, + }, + ) + }; + let first = read(); + assert!(first.turns[0].gaps.is_empty(), "{first:?}"); + assert!(!first.turns[0].output.is_empty()); + + assert!(gc(&shared, NOW + 3 * DAY, 1, 90, u64::MAX) > 0); + let again = read(); + assert_eq!( + again.turns[0].gaps, + [ + tw_api::TranscriptGap::RequestMissing, + tw_api::TranscriptGap::ResponseMissing + ] + ); + } +} diff --git a/crates/tw-store/src/transcript/mod.rs b/crates/tw-store/src/transcript/mod.rs index d5a8e676..0ad55314 100644 --- a/crates/tw-store/src/transcript/mod.rs +++ b/crates/tw-store/src/transcript/mod.rs @@ -18,11 +18,16 @@ //! //! 几百轮的会话，每个请求几 MB。每份正文读一次、解析一次，读完这一轮就丢；从一个请求留到 //! 下一个的只有每条消息 16 字节的指纹和系统提示。 +//! +//! **读过的不再读**（[`Cache`]）：界面开着一次会话时，每来一轮就要一次，从头读的话，两百轮 +//! 的会话每来一轮就要读、解析几百 MB。记着最近几次会话读到了哪儿（读完时的 [`Chain`] 和读出 +//! 来的几轮），下一次只读新来的请求。只记定下来的那一截：正文还可能变的那几轮每次都重读。 mod answer; mod read; use std::collections::HashSet; +use std::sync::Arc; use serde_json::Value; use tw_api::{TranscriptGap, TranscriptMessage, TranscriptPart, TranscriptRole, TranscriptTurn}; @@ -33,29 +38,207 @@ use crate::db::RequestRow; use crate::search::text::{client_dialect, dialect_of}; use read::{Input, Piece, Role}; -/// 读成一段对话。`rows` 是这次会话的请求，和会话详情同样的顺序 +/// 一轮结束这么久之后就定下来了：它的正文要么已经在盘上，要么不会再来。正文和那一行各走 +/// 各的路落盘（见 `crate::task`），先后差不了几毫秒，排队排得再久也到不了这么久 +const SETTLE_MS: i64 = 120_000; + +/// 记着最近几次会话读到了哪儿。一次几百轮的会话记下的是几 MB 的文字 +const CACHED: usize = 4; + +/// 整段读一遍，不记。`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); + Cache::default().read( + blobs, + &Ask { + session, + rows, + from_turn: 0, + running: None, + now_ms: i64::MAX, + }, + ) +} + +/// 要读哪一次会话、交出去哪几轮（见 [`Cache::read`]）。 +pub struct Ask<'a> { + pub session: &'a str, + /// 这次会话的请求，和会话详情同样的顺序（[`crate::Db::session_requests`]） + pub rows: &'a [RequestRow], + /// 只交出从这一轮起的那些（从 0 数）。0 是整段 + pub from_turn: usize, + /// 这次会话里还在跑的请求里最早开始的那个：（开始的时刻，请求号）。它落库时排在开始的 + /// 那一刻，不一定排在最后：排在它后面的几轮还没定下来（见 + /// [`tw_api::Transcript::settled_turns`]） + pub running: Option<(i64, i64)>, + /// 此刻，Unix 毫秒 + pub now_ms: i64, +} + +/// 最近几次会话读到了哪儿（见 [`Cache::read`]）。存储层一份（[`crate::Recorder::transcripts`]）。 +#[derive(Default)] +pub struct Cache { + kept: std::sync::Mutex, +} + +#[derive(Default)] +struct Kept { + /// 正文回收每删掉一次东西就加一。读到一半它变了的，读出来的就不记 + generation: u64, + /// 最近用过的在最后 + sessions: Vec<(String, Arc)>, +} + +/// 一次会话读到了哪儿：读过的那几个请求，读完它们时的 [`Chain`]，读出来的几轮（没打码）。 +#[derive(Clone, Default)] +struct Progress { + ids: Vec, + chain: Chain, + doc: Doc, +} + +/// 读出来的东西：系统提示和每一轮。**没打码**：前后比对、换工具调用的号都要原文 +#[derive(Clone, Default)] +struct Doc { + system: Option, + turns: Vec, +} + +impl Cache { + /// 正文回收删掉了东西：记着的都不作数了 + pub fn invalidate(&self) { + let mut k = self.kept(); + k.generation += 1; + k.sessions.clear(); + } + + /// 锁中毒了照样用里面的：记着的东西坏不了谁，最多重读一遍 + fn kept(&self) -> std::sync::MutexGuard<'_, Kept> { + self.kept.lock().unwrap_or_else(|p| p.into_inner()) + } + + /// 读成一段对话。**已脱敏**，和请求详情里的正文同一套打码（`tw_secret::mask_body`）。 + /// + /// **接着上次读到的地方往下读**：记着的那几个请求还是这次会话开头的那几个，就只读后面 + /// 新来的；中间插进来一个（开始得早、结束得晚的请求）、开头的过期删掉了，都从头读。 + /// 记下的是定下来的那一截：最后几轮的正文可能还在路上，它们下一次还要重读。 + pub fn read(&self, blobs: &Blobs, ask: &Ask) -> tw_api::Transcript { + let (generation, kept) = { + let mut k = self.kept(); + let kept = k + .sessions + .iter() + .position(|(s, _)| s == ask.session) + .map(|i| { + let hit = k.sessions.remove(i); + let p = hit.1.clone(); + k.sessions.push(hit); + p + }); + (k.generation, kept) + }; + let base = kept + .filter(|p| { + p.ids.len() <= ask.rows.len() + && p.ids.iter().zip(ask.rows).all(|(id, row)| *id == row.id) + }) + .unwrap_or_default(); + let fresh = &ask.rows[base.ids.len()..]; + // 定下来的那一截有多长。记着的那一截本来就是定下来的 + let (whole, settled) = if fresh.is_empty() { + let n = base.ids.len(); + (base, n) + } else { + let (whole, keep) = continue_reading(&base, fresh, blobs, ask.now_ms); + let n = keep.ids.len(); + self.keep(ask.session, generation, keep); + (whole, n) + }; + + let total = whole.doc.turns.len(); + // 后面的请求还会改写的那一轮：最后一个读得懂的请求的回答（见 `same_calls`） + let rewritable = match whole.chain.last { + Some((at, true)) => at, + _ => total, + }; + let before_running = ask.running.map_or(total, |r| { + ask.rows.partition_point(|row| (row.at_ms, row.id) < r) + }); + let from = ask.from_turn.min(total); + let mut t = tw_api::Transcript { + session: ask.session.to_string(), + system: whole.doc.system.clone(), + total_turns: total as u32, + settled_turns: settled.min(rewritable).min(before_running) as u32, + turns: whole.doc.turns[from..].to_vec(), + }; + mask(&mut t); + t + } + + /// 记下一次会话读到的地方。正文回收在这期间删过东西的不记：读的时候它们可能还在 + fn keep(&self, session: &str, generation: u64, p: Arc) { + let mut k = self.kept(); + if k.generation != generation { + return; + } + k.sessions.retain(|(s, _)| s != session); + k.sessions.push((session.to_string(), p)); + if k.sessions.len() > CACHED { + k.sessions.remove(0); + } } - mask(&mut t); - t +} + +/// 从 `base` 读到的地方接着读 `fresh`。返回读完的样子，和其中定下来的那一截（要记下的）。 +/// +/// 一轮定下来了：它结束得够久（[`SETTLE_MS`]），或者它什么都不缺。还缺着正文、又刚结束的 +/// 那一轮，正文可能还在路上；从它起都不算定下来。 +fn continue_reading( + base: &Progress, + fresh: &[RequestRow], + blobs: &Blobs, + now_ms: i64, +) -> (Arc, Arc) { + let mut p = base.clone(); + let mut settled: Option = None; + for row in fresh { + let ended = row.at_ms.saturating_add(row.duration_ms.unwrap_or(0)); + let recent = now_ms.saturating_sub(ended) < SETTLE_MS; + // 这一个可能定不下来：先留着读它之前的样子。读它只会改它自己和上一个读得懂的那一轮 + let before = (settled.is_none() && recent).then(|| { + let last = p.chain.last.map(|(at, _)| (at, p.doc.turns[at].clone())); + (p.chain.clone(), p.doc.system.clone(), last) + }); + p.chain.turn(row, blobs, &mut p.doc); + p.ids.push(row.id); + if let Some((chain, system, last)) = before + && p.doc.turns.last().is_some_and(|t| !t.gaps.is_empty()) + { + let n = p.ids.len() - 1; + let mut turns = p.doc.turns[..n].to_vec(); + if let Some((at, turn)) = last { + turns[at] = turn; + } + settled = Some(Progress { + ids: p.ids[..n].to_vec(), + chain, + doc: Doc { system, turns }, + }); + } + } + let p = Arc::new(p); + let settled = settled.map_or_else(|| p.clone(), Arc::new); + (p, settled) } /// 一条消息比对用的指纹 type Fp = [u8; 16]; /// 读到哪儿了：上一个读得懂的请求留下来的东西。 -#[derive(Default)] +#[derive(Clone, Default)] struct Chain { /// 它的每条消息的指纹（只有推理的消息不比，不在里面）。还没有读得懂的请求是 None prev: Option>, @@ -68,7 +251,7 @@ struct Chain { } impl Chain { - fn turn(&mut self, row: &RequestRow, blobs: &Blobs, t: &mut tw_api::Transcript) { + fn turn(&mut self, row: &RequestRow, blobs: &Blobs, t: &mut Doc) { let mut turn = TranscriptTurn { id: row.id.to_string(), restart: false, @@ -119,7 +302,7 @@ impl Chain { client: Dialect, v: &Value, turn: &mut TranscriptTurn, - t: &mut tw_api::Transcript, + t: &mut Doc, freeform: &mut HashSet, ) { let read::Body { @@ -1844,6 +2027,239 @@ mod tests { assert!(t.turns[0].output.is_empty() && t.turns[0].gaps.is_empty()); } + // ───────────────────────────────────────────────── 读过的不再读 + + /// 很久以后的此刻：每一轮都早就结束了 + const LATER: i64 = NOW + 86_400_000; + + impl Disk { + /// 用 `cache` 读：从 `from_turn` 起，此刻是 `now_ms` + fn read(&self, cache: &Cache, from_turn: usize, now_ms: i64) -> Transcript { + self.read_running(cache, from_turn, None, now_ms) + } + + fn read_running( + &self, + cache: &Cache, + from_turn: usize, + running: Option<(i64, i64)>, + now_ms: i64, + ) -> Transcript { + let rows = self.db.session_requests("s").unwrap(); + cache.read( + &self.blobs, + &Ask { + session: "s", + rows: &rows, + from_turn, + running, + now_ms, + }, + ) + } + + /// 第 `id` 轮结束之后过了 `ms` 毫秒（`row` 给的每一轮都跑了 4 秒） + fn after(&self, id: i64, ms: i64) -> i64 { + NOW + id * 1000 + 4000 + ms + } + } + + /// 一段带工具调用的对话：第 `k` 轮的请求和回答。回答里工具调用的号和客户端记下的不一样 + /// （转换过格式的上游就是这样）：读下一轮时，上一轮回答里的号换成客户端的那个 + fn claude_turn(k: usize) -> (Value, Vec) { + let mut messages = vec![user("把 bug 修了")]; + for i in 0..k { + messages.push(json!({"role": "assistant", "content": [ + {"type": "text", "text": format!("读 {i}")}, + {"type": "tool_use", "id": format!("t{i}"), "name": "Read", "input": {"n": i}}]})); + messages.push(json!({"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": format!("t{i}"), "content": format!("文件 {i}")}]})); + } + let answer = anthropic_stream(&[ + json!({"type": "text", "text": format!("读 {k}")}), + json!({"type": "tool_use", "id": format!("up{k}"), "name": "Read", "input": {"n": k}}), + ]); + (anthropic(&messages), answer) + } + + /// **接着读和从头读是同一段对话**：每来一轮读一次，每一次都和整段从头读的一样；只要 + /// 后面几轮的，就是整段的那一截。中间夹着一个数 token 的调用 + #[test] + fn reading_on_from_where_it_stopped_gives_what_reading_it_all_gives() { + let mut d = Disk::new(); + let cache = Cache::default(); + for k in 0..8 { + let (request, answer) = claude_turn(k); + d.turn("/v1/messages", &request, &answer); + if k == 3 { + d.turn("/v1/messages/count_tokens", &request, b"{}"); + } + let full = d.transcript(); + let got = d.read(&cache, 0, LATER); + assert_eq!(got.turns, full.turns, "第 {k} 轮"); + assert_eq!(got.system, full.system); + assert_eq!(got.total_turns as usize, full.turns.len()); + assert!(got.turns.iter().skip(1).all(|t| !t.restart), "{got:?}"); + // 上一轮回答里的号换成了客户端记下的，最后一轮的还是上游给的 + if k > 0 { + let ids = |t: &TranscriptTurn| -> Vec { + t.output + .iter() + .filter_map(|p| match p { + P::ToolCall { id, .. } => Some(id.clone()), + _ => None, + }) + .collect() + }; + let answered: Vec> = got + .turns + .iter() + .map(ids) + .filter(|x| !x.is_empty()) + .collect(); + assert_eq!(answered[k - 1], [format!("t{}", k - 1)]); + assert_eq!(answered[k], [format!("up{k}")]); + } + for from in [0, 1, full.turns.len() - 1, full.turns.len(), 99] { + let part = d.read(&cache, from, LATER); + assert_eq!( + part.turns, + full.turns[from.min(full.turns.len())..], + "{from}" + ); + assert_eq!(part.total_turns, got.total_turns); + } + } + } + + /// **读过的不再读，正文回收之后才重读。**读过一遍之后正文没了（这里直接删掉），记着的 + /// 照样交出去 —— 说明它真的没再去读；作废之后重读，那几轮说出缺了什么 + #[test] + fn what_was_read_is_not_read_again_until_bodies_are_reclaimed() { + let mut d = Disk::new(); + let cache = Cache::default(); + for k in 0..2 { + let (request, answer) = claude_turn(k); + d.turn("/v1/messages", &request, &answer); + } + let first = d.read(&cache, 0, LATER); + assert!(first.turns.iter().all(|t| t.gaps.is_empty()), "{first:?}"); + // 整天的正文都删掉：两天之后、只留一天 + d.blobs.gc(NOW + 2 * 86_400_000, 1, u64::MAX); + assert!(d.blobs.get(NOW + 1000, 1, Which::Request).is_none()); + assert_eq!(d.read(&cache, 0, LATER), first, "记着的没用上"); + + cache.invalidate(); + let again = d.read(&cache, 0, LATER); + for t in &again.turns { + assert_eq!( + t.gaps, + [Gap::RequestMissing, Gap::ResponseMissing], + "{again:?}" + ); + } + } + + /// 刚结束、回答还没落盘的那一轮**不记下**：回答一落盘，下一次就读得到，后面那一轮跟着 + /// 对得上（上一轮的回答交出去了，它的输入里不再重复那条助手消息）。结束得够久还缺着的， + /// 就是真的缺了 + #[test] + fn a_turn_whose_answer_is_still_on_its_way_is_read_again() { + let mut d = Disk::new(); + let cache = Cache::default(); + let (r0, a0) = claude_turn(0); + d.turn("/v1/messages", &r0, &a0); + let (r1, a1) = claude_turn(1); + let id = d.put( + "/v1/messages", + Some(r1.to_string().as_bytes()), + None, + |_| {}, + ); + let now = d.after(id, 500); + + let t = d.read(&cache, 0, now); + assert_eq!(t.turns[1].gaps, [Gap::ResponseMissing]); + assert_eq!(t.settled_turns, 1, "回答还在路上的那一轮算定下来了"); + // 下一轮来了，回答还没到 + let (r2, a2) = claude_turn(2); + d.turn("/v1/messages", &r2, &a2); + let t = d.read(&cache, 0, d.after(id + 1, 500)); + assert_eq!(t.turns[1].gaps, [Gap::ResponseMissing]); + assert_eq!(t.settled_turns, 1); + + // 回答落盘了：重读那一轮，和从头读的一样 + assert!(d.blobs.put(NOW + id * 1000, id, Which::Response, &a1)); + let t = d.read(&cache, 0, d.after(id + 1, 600)); + assert_eq!(t.turns, d.transcript().turns); + assert!(t.turns.iter().all(|x| x.gaps.is_empty()), "{t:?}"); + assert_eq!(t.total_turns, 3); + // 最后一轮的回答下一轮还会改写（工具调用的号），它还不算定下来 + assert_eq!(t.settled_turns, 2); + + // 回答一直没来：结束得够久之后它就是缺了，算定下来 + let (r3, _) = claude_turn(3); + let late = d.put( + "/v1/messages", + Some(r3.to_string().as_bytes()), + None, + |_| {}, + ); + let t = d.read(&cache, 0, d.after(late, 1000)); + assert_eq!(t.settled_turns, 3); + let t = d.read(&cache, 0, d.after(late, SETTLE_MS)); + assert_eq!(t.turns[3].gaps, [Gap::ResponseMissing]); + assert_eq!(t.settled_turns, 4, "缺着回答的最后一轮不会再被改写"); + } + + /// 定下来的几轮也要排在还在跑的请求前面：一个开始得早、结束得晚的请求落库时插在它开始的 + /// 那一刻，它后面的几轮都要往后挪一格 + #[test] + fn turns_after_a_request_still_running_are_not_settled() { + let mut d = Disk::new(); + let cache = Cache::default(); + for k in 0..3 { + let (request, answer) = claude_turn(k); + d.turn("/v1/messages", &request, &answer); + } + let t = d.read(&cache, 0, LATER); + assert_eq!((t.total_turns, t.settled_turns), (3, 2)); + // 第二轮之前开始的一个请求还在跑 + let t = d.read_running(&cache, 0, Some((NOW + 2 * 1000 - 1, 99)), LATER); + assert_eq!(t.settled_turns, 1); + // 在最后一轮之后开始的不碍事 + let t = d.read_running(&cache, 0, Some((NOW + 10_000, 99)), LATER); + assert_eq!(t.settled_turns, 2); + } + + /// 一个请求插进了读过的几轮中间（开始得早、结束得晚）：记着的那一截对不上了，从头读 + #[test] + fn a_request_landing_between_turns_already_read_is_read_from_the_start() { + let mut d = Disk::new(); + let cache = Cache::default(); + let (r0, a0) = claude_turn(0); + let (r1, a1) = claude_turn(1); + let (r2, a2) = claude_turn(2); + d.turn("/v1/messages", &r0, &a0); + d.turn("/v1/messages", &r2, &a2); + d.read(&cache, 0, LATER); + // 第二轮开始在前两轮之间，结束在它们之后 + d.put( + "/v1/messages", + Some(r1.to_string().as_bytes()), + Some(&a1), + |r| r.at_ms = NOW + 1500, + ); + let got = d.read(&cache, 0, LATER); + let full = d.transcript(); + assert_eq!(got.turns, full.turns); + assert_eq!( + got.turns.iter().map(|t| t.id.as_str()).collect::>(), + ["1", "3", "2"] + ); + assert!(got.turns.iter().all(|t| !t.restart), "{got:?}"); + } + // ───────────────────────────────────────────────── 量一量 /// 读一次几百轮的会话要多久：照 Claude Code 的样子造一段 300 轮的对话，每一轮的请求带着 @@ -1928,5 +2344,34 @@ mod tests { .all(|x| x.gaps.is_empty() && x.output.len() == 3) ); } + + // 界面开着这次会话：读过前 299 轮，第 300 轮来了 + let cache = Cache::default(); + let read = |rows: &[RequestRow], from_turn| { + let t0 = std::time::Instant::now(); + let t = cache.read( + &d.blobs, + &Ask { + session: "s", + rows, + from_turn, + running: None, + now_ms: i64::MAX, + }, + ); + (t0.elapsed(), t) + }; + let (took, before) = read(&rows[..rows.len() - 1], 0); + eprintln!("first read of {} turns: {took:?}", before.total_turns); + let (took, whole) = read(&rows, 0); + eprintln!("one more turn, whole transcript again: {took:?}"); + assert_eq!(whole.turns, build("s", &rows, &d.blobs).turns); + let (took, tail) = read(&rows, before.settled_turns as usize); + eprintln!( + "one more turn, from settled_turns ({}): {took:?}, {} turns, {} KB of JSON", + before.settled_turns, + tail.turns.len(), + serde_json::to_string(&tail).unwrap().len() / 1000 + ); } } diff --git a/crates/tw-watch/Cargo.toml b/crates/tw-watch/Cargo.toml index 0390cfe0..c5ee0818 100644 --- a/crates/tw-watch/Cargo.toml +++ b/crates/tw-watch/Cargo.toml @@ -15,6 +15,10 @@ notify = { workspace = true } tokio = { workspace = true } thiserror = { workspace = true } +# macOS 上不递归的监听自己用 kqueue 写（见 src/kqueue.rs） +[target.'cfg(target_os = "macos")'.dependencies] +libc = { workspace = true } + [dev-dependencies] tempfile = "3" tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } diff --git a/crates/tw-watch/src/kqueue.rs b/crates/tw-watch/src/kqueue.rs new file mode 100644 index 00000000..0e165465 --- /dev/null +++ b/crates/tw-watch/src/kqueue.rs @@ -0,0 +1,917 @@ +//! macOS 上的不递归监听：kqueue。 +//! +//! **为什么不用 notify 在 macOS 上默认的 FSEvents。**FSEvents 只会盯一整棵树： +//! `RecursiveMode::NonRecursive` 是 notify 在我们进程里把深处的事件丢掉，而内核照样把 +//! 底下每一次写都送过来（零延迟、按文件报），每一次都叫醒我们。桌面端要盯 `~` 本身 +//! （`~/.claude.json` 在那儿），实测空闲时 `~` 底下每秒七到十四个事件，没有一个相关 —— +//! 「空闲时接近零」就这么没了。 +//! +//! **也不用 notify 的 kqueue 后端**（`macos_kqueue`）：它不递归地盯一个目录时，目录里 +//! 一有变化，就把目录里第一个它没盯过的条目当成「新建的」，并且**递归地**盯上它 —— +//! 盯着 `~` 的话，第一次有人在 `~` 里建个文件，它就去把 `~/Library` 整棵树挨个打开。 +//! +//! 这里的做法： +//! +//! - **目录**：打开（`O_EVTONLY`：只为了收事件，不妨碍卸载）挂一个 vnode 事件。只有目录 +//! 自己的条目增、删、改名（包括 rename 进来的原子保存）才叫醒我们，子目录里再怎么变 +//! 都不会。醒了之后列一遍目录，只看算数的条目：多了、少了、换了一个（inode 变了）才算 +//! 一次。 +//! - **算数的普通文件**：也各挂一个。原地写（不经过 rename 的保存、追加、截断）不动目录， +//! 只有盯着文件本身才看得见。文件被换掉之后，手上那个 fd 指的还是旧的那个 inode —— +//! 它报一个删除或改名，目录那边重新列一遍，就地挂到新的上。 +//! - **只看内容和存在性**：属性变了（`NOTE_ATTRIB`）只有大小或修改时间跟着变才算。截断 +//! 只报这一个；改权限不算，读文件本来就不报。 +//! +//! **fd 有名额。**每盯一个目录或文件占一个 fd，而从访达打开的应用默认只有 256 个。整个 +//! 进程里所有监听合起来最多占软上限的一半（[`budget`]）；超了，这个监听整个换回 FSEvents —— +//! 费电，但一个事件都不漏。只盯一部分文件、悄悄漏掉它们的原地写，比费电更糟。 +//! +//! **盯的是路径，不是那个 inode**，和 FSEvents 一样：盯着的目录自己被删掉、挪走时算一次 +//! （`relevant` 问的是目录自己的路径），然后改盯它的上一层，等这个名字再出现 —— 整个目录 +//! 被换掉（rename 一个新的过来）、删了又建，都接着盯新的那个，出现时也算一次。上一层也 +//! 没了，才不再盯。 + +use std::collections::{BTreeSet, HashMap}; +use std::ffi::{CString, OsString}; +use std::io; +use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd}; +use std::os::unix::ffi::OsStrExt; +use std::os::unix::fs::{DirEntryExt, MetadataExt}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::mpsc::Sender; +use std::sync::{Arc, Mutex, PoisonError}; + +use crate::{Relevant, WatchError}; + +/// 目录上要的事件：条目变了（`NOTE_WRITE`），和它自己没了、被挪走 +const DIR_EVENTS: u32 = libc::NOTE_WRITE | GONE; + +/// 文件上要的事件：内容变了、属性变了（[`Loop::stamp_moved`] 再筛），和它自己没了、被挪走 +const FILE_EVENTS: u32 = libc::NOTE_WRITE | libc::NOTE_EXTEND | libc::NOTE_ATTRIB | GONE; + +/// 手上这个 vnode 不再是那个路径上的东西了：删了、改名挪走了、文件系统卸载了 +const GONE: u32 = libc::NOTE_DELETE | libc::NOTE_RENAME | libc::NOTE_REVOKE; + +/// 叫事件循环停下的那个用户事件（`EVFILT_USER`；和 vnode 事件不在一个名字空间里） +const STOP: usize = 0; + +/// fd 名额的硬顶：软上限再高，所有监听合起来也不占更多 +const MAX_HELD: usize = 1024; + +/// 整个进程里所有监听此刻占着的 fd +static HELD: AtomicUsize = AtomicUsize::new(0); + +/// 所有监听合起来最多占多少个 fd：软上限的一半，最多 [`MAX_HELD`]。 +/// +/// **按进程算，不按监听算**：一个进程里可能同时有好几个（桌面端换一份监听时，新的起来了 +/// 旧的才放掉）。每次都现问软上限：进程中途调高了，名额跟着多。 +fn budget() -> usize { + let mut r = libc::rlimit { + rlim_cur: 0, + rlim_max: 0, + }; + // SAFETY: getrlimit 只往传进去的结构体里写 + let soft = if unsafe { libc::getrlimit(libc::RLIMIT_NOFILE, &mut r) } == 0 { + r.rlim_cur + } else { + 256 + }; + usize::try_from(soft / 2).unwrap_or(MAX_HELD).min(MAX_HELD) +} + +/// 一个占着名额的、只为了收事件打开的 fd。放掉它，fd 关上（挂在上面的事件跟着摘掉）， +/// 名额还回去。 +struct Held(OwnedFd); + +impl Held { + /// 打开 `path`。名额用完了是 `Ok(None)`。 + fn open(path: &Path) -> io::Result> { + // 先占名额再开 fd。手写比较交换：`fetch_update` 在新版标准库里改了名 + let mut n = HELD.load(Ordering::Acquire); + loop { + if n >= budget() { + return Ok(None); + } + match HELD.compare_exchange_weak(n, n + 1, Ordering::AcqRel, Ordering::Acquire) { + Ok(_) => break, + Err(now) => n = now, + } + } + let release = || HELD.fetch_sub(1, Ordering::AcqRel); + let Ok(c) = CString::new(path.as_os_str().as_bytes()) else { + release(); + return Err(io::ErrorKind::InvalidInput.into()); + }; + // O_NONBLOCK：万一打开的是个管道（条目刚被换成了别的东西），也不会卡在这里 + // SAFETY: c 是以 NUL 结尾的路径；返回的 fd 由下面的 OwnedFd 接管 + let fd = unsafe { + libc::open( + c.as_ptr(), + libc::O_EVTONLY | libc::O_CLOEXEC | libc::O_NONBLOCK, + ) + }; + if fd < 0 { + let e = io::Error::last_os_error(); + release(); + return Err(e); + } + // SAFETY: fd 是刚打开的、只有这里拿着 + Ok(Some(Held(unsafe { OwnedFd::from_raw_fd(fd) }))) + } + + fn raw(&self) -> RawFd { + self.0.as_raw_fd() + } + + /// fstat：用的是它的类型、(dev, ino)、大小和修改时间 + fn stat(&self) -> io::Result { + // SAFETY: stat 是纯输出参数，全零是合法的初值 + let mut st: libc::stat = unsafe { std::mem::zeroed() }; + // SAFETY: fd 活着；fstat 只往 st 里写 + if unsafe { libc::fstat(self.raw(), &mut st) } != 0 { + return Err(io::Error::last_os_error()); + } + Ok(st) + } +} + +impl Drop for Held { + fn drop(&mut self) { + HELD.fetch_sub(1, Ordering::AcqRel); + } +} + +/// 一个活着的 kqueue 监听；名额不够时是退回去的 FSEvents。 +pub struct Watch { + _backend: Backend, +} + +/// 拿着它就是在盯；扔掉就停 +enum Backend { + Kqueue { + _running: Running, + }, + Notify { + _fsevents: notify::RecommendedWatcher, + }, +} + +/// 跑着的事件循环：kqueue 本身，和跑循环的线程 +struct Running { + kq: Arc, + thread: Option>, + /// 事件循环醒过几次（测试数它：深处的改动不该叫醒它） + #[cfg(test)] + wakeups: Arc, +} + +impl Drop for Running { + /// **等线程收完再返回**：它手上的 fd 这时都已经关了、名额都还了，调用方紧接着再起一个 + /// 监听不会撞上名额 + fn drop(&mut self) { + let _ = control( + self.kq.as_raw_fd(), + &[event(STOP, libc::EVFILT_USER, 0, libc::NOTE_TRIGGER)], + ); + if let Some(t) = self.thread.take() { + let _ = t.join(); + } + } +} + +impl Watch { + /// 事件循环到现在醒过几次。退回 FSEvents 的是 `None` + #[cfg(test)] + pub(crate) fn wakeups(&self) -> Option { + match &self._backend { + Backend::Kqueue { _running: r } => Some(r.wakeups.load(Ordering::Relaxed)), + Backend::Notify { .. } => None, + } + } +} + +/// 盯住 `dirs`：见模块文档。每有一处算数的改动，往 `raw_tx` 里发一个。 +/// +/// **返回时就已经盯上了**：目录和文件都挂好、列过一遍，之后的改动一个都不会漏。 +pub(crate) fn watch( + dirs: &[PathBuf], + relevant: Relevant, + raw_tx: Sender<()>, +) -> Result { + let start_error = |path: &Path, e: io::Error| WatchError::Start { + path: path.to_path_buf(), + source: notify::Error::io(e), + }; + let first = dirs.first().map(PathBuf::as_path).unwrap_or(Path::new("")); + // SAFETY: kqueue() 不碰内存；返回的 fd 由 OwnedFd 接管 + let kq = unsafe { libc::kqueue() }; + if kq < 0 { + return Err(start_error(first, io::Error::last_os_error())); + } + // SAFETY: kq 是刚建的、只有这里拿着 + let kq = Arc::new(unsafe { OwnedFd::from_raw_fd(kq) }); + // fork 出去的子进程不该拿着它 + // SAFETY: kq 活着；FD_CLOEXEC 只改这个 fd 自己的标志 + unsafe { libc::fcntl(kq.as_raw_fd(), libc::F_SETFD, libc::FD_CLOEXEC) }; + control( + kq.as_raw_fd(), + &[event( + STOP, + libc::EVFILT_USER, + libc::EV_ADD | libc::EV_CLEAR, + 0, + )], + ) + .map_err(|e| start_error(first, e))?; + + let mut l = Loop { + kq: kq.clone(), + dirs: dirs + .iter() + .map(|p| Dir { + path: p.clone(), + fd: None, + parent: None, + seen: HashMap::new(), + }) + .collect(), + files: HashMap::new(), + by_fd: HashMap::new(), + relevant: Arc::new(Mutex::new(relevant)), + capped: false, + #[cfg(test)] + wakeups: Arc::new(AtomicUsize::new(0)), + }; + let mut over = false; + for (i, path) in dirs.iter().enumerate() { + match l.open_dir(i) { + Ok(true) => {} + Ok(false) => { + over = true; + break; + } + Err(e) => return Err(start_error(path, e)), + } + } + // 先挂好再列：列的时候已经有改动的话，事件也已经排在那儿了 + if !over { + over = (0..l.dirs.len()).any(|i| l.rescan(i).is_err()); + } + if over { + // 一开始就不够：手上的放掉，整个用 FSEvents + let relevant = l.handover(); + drop(l); + return Ok(Watch { + _backend: Backend::Notify { + _fsevents: crate::by_notify(dirs, relevant, raw_tx)?, + }, + }); + } + + #[cfg(test)] + let wakeups = l.wakeups.clone(); + let thread = std::thread::Builder::new() + .name("tw-watch kqueue".into()) + .spawn(move || l.run(&raw_tx)) + .map_err(|e| start_error(first, e))?; + Ok(Watch { + _backend: Backend::Kqueue { + _running: Running { + kq, + thread: Some(thread), + #[cfg(test)] + wakeups, + }, + }, + }) +} + +/// 一个盯着的目录 +struct Dir { + path: PathBuf, + /// 没了、挪走了，那里还没有新的时是 None：这时盯着 `parent` + fd: Option, + /// 它没了的时候盯着的上一层：等这个名字再出现 + parent: Option, + /// 上次列出来的、算数的条目：名字 → (inode, 是什么) + seen: HashMap, +} + +/// 一个条目是什么（readdir 给的，不另外 stat） +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Kind { + File, + Symlink, + Other, +} + +/// 一个盯着的文件（目录里算数的普通文件，或者指向普通文件的符号链接） +struct File { + fd: Held, + /// 它在哪个目录里（[`Loop::dirs`] 的下标） + dir: usize, + /// 挂上时它是哪个文件（跟完符号链接） + dev: u64, + ino: u64, + /// (大小, 修改时间的秒, 纳秒)：属性变了时拿来比，见 [`Loop::stamp_moved`] + stamp: (i64, i64, i64), +} + +/// 一个 fd 挂的是谁 +enum Target { + Dir(usize), + /// 第几个目录没了之后盯着的上一层 + Parent(usize), + File(PathBuf), +} + +/// 名额用完了：整个换回 FSEvents +struct OverBudget; + +/// 事件循环为什么停下 +enum Exit { + /// 监听被扔掉了，或者去抖那一头没了 + Stop, + /// 名额用完了 + OverBudget, +} + +struct Loop { + kq: Arc, + dirs: Vec, + files: HashMap, + by_fd: HashMap, + /// 退回 FSEvents 时要交给它，交不成还得自己接着用：所以是共用的 + relevant: Arc>, + /// 退不回 FSEvents（它起不来）：名额不够的就不盯，能盯多少盯多少 + capped: bool, + #[cfg(test)] + wakeups: Arc, +} + +impl Loop { + fn run(mut self, raw_tx: &Sender<()>) { + loop { + if let Exit::Stop = self.serve(raw_tx) { + return; + } + // **先起 FSEvents，再放手上的**：中间那一下不会有改动没人看。此刻没了的目录不盯 + // （FSEvents 盯不存在的路径会报错，整个起不来） + let alive: Vec = self + .dirs + .iter() + .map(|d| d.path.clone()) + .filter(|p| p.is_dir()) + .collect(); + // 名额用完的那一批改动没算完：当作有过一次 + let _ = raw_tx.send(()); + match crate::by_notify(&alive, self.handover(), raw_tx.clone()) { + Ok(fsevents) => { + self.files.clear(); + self.by_fd.clear(); + self.dirs.clear(); + self.wait_for_stop(); + drop(fsevents); + return; + } + // 起不来就接着用 kqueue，盯得住多少算多少 + Err(_) => self.capped = true, + } + } + } + + /// 交给 FSEvents 的那一份 `relevant`：和这里用的是同一个 + fn handover(&self) -> Relevant { + let r = self.relevant.clone(); + Box::new(move |p| (r.lock().unwrap_or_else(PoisonError::into_inner))(p)) + } + + /// `p` 算不算数 + fn relevant(&self, p: &Path) -> bool { + (self.relevant.lock().unwrap_or_else(PoisonError::into_inner))(p) + } + + /// 收事件，直到被叫停或者名额用完 + fn serve(&mut self, raw_tx: &Sender<()>) -> Exit { + let mut buf = [event(0, 0, 0, 0); 64]; + loop { + let n = match wait(self.kq.as_raw_fd(), &mut buf) { + Ok(n) => n, + Err(e) if e.kind() == io::ErrorKind::Interrupted => continue, + Err(_) => return Exit::Stop, + }; + #[cfg(test)] + self.wakeups.fetch_add(1, Ordering::Relaxed); + match self.handle(&buf[..n]) { + Ok(false) => {} + Ok(true) => { + if raw_tx.send(()).is_err() { + return Exit::Stop; + } + } + Err(exit) => return exit, + } + } + } + + /// 处理一批事件，有算数的改动是 `Ok(true)`。 + /// + /// **分两遍**：第一遍只记下要做什么，第二遍才关 fd、重新打开。关掉的 fd 号马上会被 + /// 新打开的复用，一边处理一边开关的话，同一批里排在后面的、关于旧文件的事件会被 + /// 算到新文件头上。 + fn handle(&mut self, events: &[libc::kevent]) -> Result { + let mut changed = false; + let mut rescan = BTreeSet::new(); + let mut gone_dirs = Vec::new(); + let mut gone_parents = Vec::new(); + let mut gone_files = Vec::new(); + for ev in events { + if ev.filter == libc::EVFILT_USER { + return Err(Exit::Stop); + } + let Ok(fd) = RawFd::try_from(ev.ident) else { + continue; + }; + match self.by_fd.get(&fd) { + Some(Target::Dir(i)) => { + let i = *i; + if ev.fflags & GONE != 0 { + changed |= self.relevant(&self.dirs[i].path); + gone_dirs.push(i); + } + rescan.insert(i); + } + // 上一层的条目变了：等的那个名字可能出现了 + Some(Target::Parent(i)) => { + let i = *i; + if ev.fflags & GONE != 0 { + gone_parents.push(i); + } + rescan.insert(i); + } + Some(Target::File(path)) => { + let path = path.clone(); + let Some(f) = self.files.get_mut(&path) else { + continue; + }; + // 写过就算；大小和修改时间每次都记下新的，之后单报属性时才比得准 + let wrote = ev.fflags & (libc::NOTE_WRITE | libc::NOTE_EXTEND) != 0; + let moved = + ev.fflags & (libc::NOTE_WRITE | libc::NOTE_EXTEND | libc::NOTE_ATTRIB) != 0 + && Self::stamp_moved(f); + changed |= wrote || moved; + if ev.fflags & GONE != 0 { + changed = true; + rescan.insert(f.dir); + gone_files.push(path); + } + } + None => {} + } + } + for path in gone_files { + self.drop_file(&path); + } + for i in gone_dirs { + self.drop_dir(i); + } + for i in gone_parents { + self.drop_parent(i); + } + for i in rescan { + changed |= self.rescan(i).map_err(|OverBudget| Exit::OverBudget)?; + } + Ok(changed) + } + + /// 属性变了的文件，大小或修改时间是不是也变了（变了就记下新的）。截断只报属性； + /// 改权限、改属主也报属性，但它们不算 + fn stamp_moved(f: &mut File) -> bool { + let Ok(st) = f.fd.stat() else { + return false; + }; + let now = stamp(&st); + let moved = now != f.stamp; + f.stamp = now; + moved + } + + /// 打开第 `i` 个目录、挂上事件。名额用完了是 `Ok(false)`。 + fn open_dir(&mut self, i: usize) -> io::Result { + let Some(fd) = Held::open(&self.dirs[i].path)? else { + return Ok(false); + }; + if fd.stat()?.st_mode & libc::S_IFMT != libc::S_IFDIR { + return Err(io::Error::from_raw_os_error(libc::ENOTDIR)); + } + arm(self.kq.as_raw_fd(), fd.raw(), DIR_EVENTS)?; + self.by_fd.insert(fd.raw(), Target::Dir(i)); + self.dirs[i].fd = Some(fd); + Ok(true) + } + + /// 第 `i` 个目录自己没了、被挪走：手上那个 fd 不再是那个路径上的东西，放掉 + fn drop_dir(&mut self, i: usize) { + if let Some(fd) = self.dirs[i].fd.take() { + self.by_fd.remove(&fd.raw()); + } + } + + /// 第 `i` 个目录的上一层也没了、被挪走：不再等 + fn drop_parent(&mut self, i: usize) { + if let Some(fd) = self.dirs[i].parent.take() { + self.by_fd.remove(&fd.raw()); + } + } + + /// 没了的第 `i` 个目录：按原来的路径再打开。**先盯上一层，再试着打开** —— 反过来的话， + /// 两步之间换上来的那个就没人看见了。打开了就不用再盯上一层；打不开就等上一层报信 + fn revive(&mut self, i: usize) -> Result<(), OverBudget> { + let over = |capped: bool| if capped { Ok(()) } else { Err(OverBudget) }; + if self.dirs[i].parent.is_none() + && let Some(up) = self.dirs[i].path.parent().map(Path::to_path_buf) + && !up.as_os_str().is_empty() + { + match Held::open(&up) { + Ok(Some(fd)) => { + if arm(self.kq.as_raw_fd(), fd.raw(), DIR_EVENTS).is_ok() { + self.by_fd.insert(fd.raw(), Target::Parent(i)); + self.dirs[i].parent = Some(fd); + } + } + Ok(None) => over(self.capped)?, + // 上一层也没了：不再盯 + Err(_) => {} + } + } + match self.open_dir(i) { + Ok(true) => self.drop_parent(i), + Ok(false) => over(self.capped)?, + Err(_) => {} + } + Ok(()) + } + + /// 重新列第 `i` 个目录：算数的条目多了、少了、换了，是 `Ok(true)`。顺带把要盯的文件 + /// 对上：新的挂上，换掉的重新挂，不在了的放掉。 + /// + /// 目录自己没了的，先按原来的路径再打开（[`Self::revive`]）：又有了就接着盯它，它的 + /// 出现也算一次；没有的话，算数的条目一个都不剩了。 + fn rescan(&mut self, i: usize) -> Result { + let mut changed = false; + if self.dirs[i].fd.is_none() { + self.revive(i)?; + if self.dirs[i].fd.is_some() { + changed |= self.relevant(&self.dirs[i].path); + } + } + let now = match &self.dirs[i].fd { + Some(_) => self.list(&self.dirs[i].path), + None => HashMap::new(), + }; + changed |= now != self.dirs[i].seen; + self.dirs[i].seen = now; + self.sync_files(i)?; + Ok(changed) + } + + /// 目录里算数的条目 + fn list(&self, dir: &Path) -> HashMap { + let Ok(rd) = std::fs::read_dir(dir) else { + return HashMap::new(); + }; + rd.flatten() + .filter(|e| self.relevant(&dir.join(e.file_name()))) + .map(|e| { + // readdir 给的类型；给不出来时 std 自己补一次 lstat + let kind = match e.file_type() { + Ok(t) if t.is_file() => Kind::File, + Ok(t) if t.is_symlink() => Kind::Symlink, + _ => Kind::Other, + }; + (e.file_name(), (e.ino(), kind)) + }) + .collect() + } + + /// 把第 `i` 个目录里要盯的文件对上它此刻的样子（[`Dir::seen`]） + fn sync_files(&mut self, i: usize) -> Result<(), OverBudget> { + let dir = &self.dirs[i]; + let want: Vec<(PathBuf, u64, Kind)> = dir + .seen + .iter() + .filter(|(_, (_, kind))| *kind != Kind::Other) + .map(|(name, (ino, kind))| (dir.path.join(name), *ino, *kind)) + .collect(); + let stale: Vec = self + .files + .iter() + .filter(|(p, f)| f.dir == i && !want.iter().any(|(w, ..)| w == *p)) + .map(|(p, _)| p.clone()) + .collect(); + for p in stale { + self.drop_file(&p); + } + for (path, ino, kind) in want { + if let Some(f) = self.files.get(&path) { + // 还是挂着的那一个就不动。符号链接按它指向的那个文件比 + let same = match kind { + Kind::Symlink => { + std::fs::metadata(&path).is_ok_and(|m| (m.dev(), m.ino()) == (f.dev, f.ino)) + } + _ => ino == f.ino, + }; + if same { + continue; + } + self.drop_file(&path); + } + // 符号链接指向的不是普通文件（目录、设备、不存在）就不盯 + if kind == Kind::Symlink && !std::fs::metadata(&path).is_ok_and(|m| m.is_file()) { + continue; + } + self.add_file(i, path)?; + } + Ok(()) + } + + /// 挂上一个文件。打不开（刚好又没了、没有权限）就算了：它的增删改名目录那边照样看得见 + fn add_file(&mut self, dir: usize, path: PathBuf) -> Result<(), OverBudget> { + let fd = match Held::open(&path) { + Ok(Some(fd)) => fd, + Ok(None) if self.capped => return Ok(()), + Ok(None) => return Err(OverBudget), + Err(_) => return Ok(()), + }; + let Ok(st) = fd.stat() else { return Ok(()) }; + if st.st_mode & libc::S_IFMT != libc::S_IFREG { + return Ok(()); + } + if arm(self.kq.as_raw_fd(), fd.raw(), FILE_EVENTS).is_err() { + return Ok(()); + } + self.by_fd.insert(fd.raw(), Target::File(path.clone())); + self.files.insert( + path, + File { + dir, + dev: st.st_dev as u64, + ino: st.st_ino, + stamp: stamp(&st), + fd, + }, + ); + Ok(()) + } + + fn drop_file(&mut self, path: &Path) { + if let Some(f) = self.files.remove(path) { + self.by_fd.remove(&f.fd.raw()); + } + } + + /// 只等叫停的那个事件 + fn wait_for_stop(&self) { + let mut buf = [event(0, 0, 0, 0); 4]; + loop { + match wait(self.kq.as_raw_fd(), &mut buf) { + Ok(n) if buf[..n].iter().any(|e| e.filter == libc::EVFILT_USER) => return, + Ok(_) => {} + Err(e) if e.kind() == io::ErrorKind::Interrupted => {} + Err(_) => return, + } + } + } +} + +/// (大小, 修改时间的秒, 纳秒) +fn stamp(st: &libc::stat) -> (i64, i64, i64) { + (st.st_size, st.st_mtime, st.st_mtime_nsec) +} + +fn event(ident: usize, filter: i16, flags: u16, fflags: u32) -> libc::kevent { + libc::kevent { + ident, + filter, + flags, + fflags, + data: 0, + udata: std::ptr::null_mut(), + } +} + +/// 在 `fd` 上挂 vnode 事件。`EV_CLEAR`：报过一次就清掉，不会同一件事一直报 +fn arm(kq: RawFd, fd: RawFd, fflags: u32) -> io::Result<()> { + let ident = usize::try_from(fd).map_err(|_| io::Error::from_raw_os_error(libc::EBADF))?; + control( + kq, + &[event( + ident, + libc::EVFILT_VNODE, + libc::EV_ADD | libc::EV_CLEAR, + fflags, + )], + ) +} + +/// 改 kqueue 里挂着的事件，不收 +fn control(kq: RawFd, changes: &[libc::kevent]) -> io::Result<()> { + let n = libc::c_int::try_from(changes.len()).unwrap_or(libc::c_int::MAX); + // SAFETY: changes 活着、长度如实；不收事件，输出缓冲区是空的 + let r = unsafe { + libc::kevent( + kq, + changes.as_ptr(), + n, + std::ptr::null_mut(), + 0, + std::ptr::null(), + ) + }; + if r < 0 { + return Err(io::Error::last_os_error()); + } + Ok(()) +} + +/// 等到有事件为止（不设超时），收进 `buf`，交回收到几个 +fn wait(kq: RawFd, buf: &mut [libc::kevent]) -> io::Result { + let cap = libc::c_int::try_from(buf.len()).unwrap_or(libc::c_int::MAX); + // SAFETY: buf 活着、可写、长度如实；不改挂着的事件 + let n = unsafe { + libc::kevent( + kq, + std::ptr::null(), + 0, + buf.as_mut_ptr(), + cap, + std::ptr::null(), + ) + }; + usize::try_from(n).map_err(|_| io::Error::last_os_error()) +} + +#[cfg(test)] +mod tests { + use std::path::Path; + use std::time::Duration; + + use crate::tests::{Signal, WINDOW, append, only_md, recv, save_by_rename, scratch, settle}; + + /// 事件循环醒过几次 + fn wakeups(w: &crate::Watch) -> usize { + w._inner.wakeups().expect("退回了 FSEvents") + } + + /// **这一条就是用 kqueue 的原因。**子目录里再怎么写，事件循环一次都不醒 —— FSEvents + /// 在这里每一次写都会叫醒我们，再在进程里丢掉。而盯着的目录自己的条目变了，照样醒 + #[tokio::test] + async fn changes_deeper_down_do_not_even_wake_it() { + let d = scratch(); + let deep = d.path().join("history/2026/10"); + std::fs::create_dir_all(&deep).unwrap(); + std::fs::write(d.path().join("a.md"), "0\n").unwrap(); + let (w, mut rx) = crate::watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + for i in 0..200 { + std::fs::write(deep.join(format!("{i}.md")), "x\n").unwrap(); + } + std::fs::create_dir_all(deep.join("more/and/more")).unwrap(); + std::fs::remove_dir_all(d.path().join("history/2026")).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_millis(600)).await, + Signal::Quiet + ); + assert_eq!(wakeups(&w), 0, "深处的改动叫醒了事件循环"); + + // 这一层的条目变了才醒；不算数的条目醒了也不发信号 + std::fs::write(d.path().join("x.log"), "x\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_millis(600)).await, + Signal::Quiet + ); + assert!(wakeups(&w) > 0, "这一层的改动没叫醒它"); + } + + /// 属性变了只有大小、修改时间跟着变才算：改权限不算，`touch` 算 + #[tokio::test] + async fn only_size_and_modification_time_count_among_attributes() { + use std::os::unix::fs::PermissionsExt as _; + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = crate::watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + std::fs::set_permissions(&p, std::fs::Permissions::from_mode(0o600)).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_millis(600)).await, + Signal::Quiet, + "改权限" + ); + let later = std::time::SystemTime::now() + Duration::from_secs(5); + std::fs::File::options() + .write(true) + .open(&p) + .unwrap() + .set_modified(later) + .unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "touch" + ); + } + + /// 符号链接指过去的文件也盯：原地写算数；它被换掉之后，重新挂到换上来的那一个 + /// （`~/.claude.json` 指进 dotfiles 仓库的那种布局） + #[tokio::test] + async fn a_symlinked_file_is_followed_to_where_it_points() { + let d = scratch(); + let (watched, elsewhere) = (d.path().join("home"), d.path().join("dotfiles")); + std::fs::create_dir_all(&watched).unwrap(); + std::fs::create_dir_all(&elsewhere).unwrap(); + let real = elsewhere.join("a.md"); + std::fs::write(&real, "0\n").unwrap(); + std::os::unix::fs::symlink(&real, watched.join("a.md")).unwrap(); + let (_w, mut rx) = crate::watch(std::slice::from_ref(&watched), WINDOW, only_md).unwrap(); + + append(&real, "1\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "追加" + ); + settle(&mut rx).await; + save_by_rename(&real, "2\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换文件" + ); + settle(&mut rx).await; + append(&real, "3\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换过之后的追加" + ); + } + + /// 盯着的目录被整个换掉（rename 一个新目录过来）、删了过一阵又建出来：目录自己算数的话 + /// 各算一次，之后接着盯新的那个。盯的是路径，和 FSEvents 一样 + #[tokio::test] + async fn a_watched_directory_that_is_replaced_or_recreated_is_followed() { + let d = scratch(); + let dir = d.path().join("skills"); + std::fs::create_dir_all(&dir).unwrap(); + let me = dir.clone(); + let (_w, mut rx) = crate::watch(std::slice::from_ref(&dir), WINDOW, move |p: &Path| { + p == me || only_md(p) + }) + .unwrap(); + + let next = d.path().join("skills.new"); + std::fs::create_dir_all(&next).unwrap(); + std::fs::rename(&dir, d.path().join("skills.old")).unwrap(); + std::fs::rename(&next, &dir).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "目录被挪走" + ); + settle(&mut rx).await; + std::fs::write(dir.join("a.md"), "0\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换上来的目录没盯上" + ); + settle(&mut rx).await; + // 挪走的那个不再盯 + std::fs::write(d.path().join("skills.old/b.md"), "0\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_millis(600)).await, + Signal::Quiet + ); + + std::fs::remove_dir_all(&dir).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "删除" + ); + settle(&mut rx).await; + std::fs::create_dir(&dir).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "再建" + ); + settle(&mut rx).await; + std::fs::write(dir.join("c.md"), "0\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "再建出来的目录没盯上" + ); + } +} diff --git a/crates/tw-watch/src/lib.rs b/crates/tw-watch/src/lib.rs index 3ca8b18f..1a99809e 100644 --- a/crates/tw-watch/src/lib.rs +++ b/crates/tw-watch/src/lib.rs @@ -9,10 +9,14 @@ //! //! **一、盯目录，不盯文件。**编辑器保存不是「往原文件写」，而是「写一个临时文件 //! 再 rename 过去」。rename 之后原来那个 inode 就没人指了 —— 盯着文件的监听会在 -//! 第一次保存之后**永久失效**，而且不报任何错。 +//! 第一次保存之后**永久失效**，而且不报任何错。（macOS 的 kqueue 看不见目录里文件的 +//! 原地写，所以那边算数的文件**另外**也盯着，换掉之后由目录那边重新挂上，见 `kqueue` 模块。） //! //! **二、不递归。**配置旁边就是历史目录；`~/.claude/` 底下有每分钟都在变的东西。 -//! 递归盯它们等于给自己找一个永不停歇的事件源。 +//! 递归盯它们等于给自己找一个永不停歇的事件源。**不递归要做到内核那一层**：macOS 上 +//! notify 默认的 FSEvents 只会盯整棵树，「不递归」是在我们进程里把深处的事件丢掉 —— +//! 丢之前已经被叫醒了。所以 macOS 上用 kqueue（`kqueue` 模块），Linux 的 inotify、 +//! Windows 的 ReadDirectoryChangesW 本来就只报这一层，照旧用 notify。 //! //! **三、去抖。**一次保存常常触发好几个事件（建临时文件、rename、改属性、改 //! mtime）。不聚合的话，一次保存会触发好几轮重读。 @@ -20,6 +24,9 @@ use std::path::{Path, PathBuf}; use std::time::Duration; +#[cfg(target_os = "macos")] +mod kqueue; + #[derive(Debug, thiserror::Error)] pub enum WatchError { #[error("changes to {path} could not be watched: {source}")] @@ -31,9 +38,17 @@ pub enum WatchError { /// 一个活着的监听。**扔掉它就停止监听** —— 所以调用方必须留着它。 pub struct Watch { - _inner: notify::RecommendedWatcher, + _inner: Inner, } +#[cfg(target_os = "macos")] +type Inner = kqueue::Watch; +#[cfg(not(target_os = "macos"))] +type Inner = notify::RecommendedWatcher; + +/// 哪些路径算数（[`watch`] 的 `relevant`） +type Relevant = Box bool + Send + 'static>; + /// 盯住 `dirs`（不递归），改动涉及的路径里有一个让 `relevant` 说是的，就算一次； /// 一次之后 `debounce` 之内的都并进去，然后往通道里发**一个**信号。 /// @@ -43,6 +58,11 @@ pub struct Watch { /// 只看内容和存在性的变化（建、改、删）。**访问时间之类的不算** —— 备份和索引 /// 工具会大量产生它们。 /// +/// `relevant` 拿到的是改动的那个条目的路径（盯着的目录自己没了、被挪走时是目录自己）。 +/// 写法各平台不完全一样：inotify、kqueue 是 `dirs` 里那个目录接上条目名，FSEvents（macOS +/// 上名额不够时退回去用它）是跟完符号链接的真实路径 —— 要按整条路径比的，盯的目录先 +/// 换成真实路径，两种写法就一样了。 +/// /// 通道里最多攒一个信号：满了就丢。信号的意思是「去看一眼」，攒两个和攒一个是 /// 一回事。 pub fn watch( @@ -50,10 +70,29 @@ pub fn watch( debounce: Duration, relevant: impl Fn(&Path) -> bool + Send + 'static, ) -> Result<(Watch, tokio::sync::mpsc::Receiver<()>), WatchError> { - use notify::{EventKind, RecursiveMode, Watcher as _}; - let (tx, rx) = tokio::sync::mpsc::channel(1); let (raw_tx, raw_rx) = std::sync::mpsc::channel::<()>(); + #[cfg(target_os = "macos")] + let inner = kqueue::watch(dirs, Box::new(relevant), raw_tx)?; + #[cfg(not(target_os = "macos"))] + let inner = by_notify(dirs, Box::new(relevant), raw_tx)?; + + // 监听的回调跑在它自己的线程上，这里把它接到 tokio 上并去抖。 + std::thread::spawn(move || debounce_loop(&raw_rx, debounce, &tx)); + + Ok((Watch { _inner: inner }, rx)) +} + +/// 用 notify 给这个平台推荐的那一种盯：Linux 的 inotify、Windows 的 +/// ReadDirectoryChangesW。macOS 上是 FSEvents，只在 kqueue 的 fd 名额用完时退回来用它 +/// （见 `kqueue` 模块）。每有一处算数的改动，往 `raw_tx` 里发一个。 +fn by_notify( + dirs: &[PathBuf], + relevant: Relevant, + raw_tx: std::sync::mpsc::Sender<()>, +) -> Result { + use notify::{EventKind, RecursiveMode, Watcher as _}; + let mut w = notify::recommended_watcher(move |ev: notify::Result| { let Ok(ev) = ev else { return }; if !matches!( @@ -77,11 +116,7 @@ pub fn watch( source, })?; } - - // notify 的回调跑在它自己的线程上，这里把它接到 tokio 上并去抖。 - std::thread::spawn(move || debounce_loop(&raw_rx, debounce, &tx)); - - Ok((Watch { _inner: w }, rx)) + Ok(w) } /// 去抖：收到一个事件之后，把间隔不到 `window` 的后续事件全部并进来，静下来 @@ -110,21 +145,21 @@ fn debounce_loop( } #[cfg(test)] -mod tests { +pub(crate) mod tests { use super::*; - const WINDOW: Duration = Duration::from_millis(200); + pub(crate) const WINDOW: Duration = Duration::from_millis(200); /// 等一个信号，**区分三种结果**：收到了、窗口内没动静、通道关了（监听线程 /// 没了之后永远不会再有信号 —— 不分开的话它和「收到了」分不清）。 #[derive(Debug, PartialEq)] - enum Signal { + pub(crate) enum Signal { Got, Quiet, Closed, } - async fn recv(rx: &mut tokio::sync::mpsc::Receiver<()>, within: Duration) -> Signal { + pub(crate) async fn recv(rx: &mut tokio::sync::mpsc::Receiver<()>, within: Duration) -> Signal { match tokio::time::timeout(within, rx.recv()).await { Ok(Some(())) => Signal::Got, Ok(None) => Signal::Closed, @@ -132,7 +167,7 @@ mod tests { } } - fn only_md(p: &Path) -> bool { + pub(crate) fn only_md(p: &Path) -> bool { p.extension().is_some_and(|e| e == "md") } @@ -145,7 +180,7 @@ mod tests { /// 是对的，错的是测试。内存盘上的写不碰磁盘，前提才真的成立。 /// /// macOS 没有 `/dev/shm`，用系统的临时目录 —— 那边从来没有因此失败过。 - fn scratch() -> tempfile::TempDir { + pub(crate) fn scratch() -> tempfile::TempDir { let shm = Path::new("/dev/shm"); if cfg!(target_os = "linux") && shm.is_dir() { tempfile::tempdir_in(shm).unwrap() @@ -244,4 +279,158 @@ mod tests { }; assert_eq!(path, missing); } + + /// 调用方把它放在跨线程共用的地方（桌面端放在应用状态的 Mutex 里） + #[test] + fn a_watch_can_move_between_threads() { + fn send() {} + send::(); + } + + /// 等前面那些改动的信号都收完、静下来：之后每一步要的是它自己的那个信号 + pub(crate) async fn settle(rx: &mut tokio::sync::mpsc::Receiver<()>) { + while recv(rx, WINDOW * 3).await == Signal::Got {} + } + + /// 往文件末尾追加：原地写，不经过 rename + pub(crate) fn append(p: &Path, text: &str) { + use std::io::Write as _; + let mut f = std::fs::OpenOptions::new().append(true).open(p).unwrap(); + f.write_all(text.as_bytes()).unwrap(); + } + + /// 编辑器那种保存：写一个临时文件，rename 过去 + pub(crate) fn save_by_rename(p: &Path, text: &str) { + let tmp = p.with_extension("swp"); + std::fs::write(&tmp, text).unwrap(); + std::fs::rename(&tmp, p).unwrap(); + } + + #[tokio::test] + async fn every_save_by_rename_is_seen() { + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + for i in 1..=3 { + save_by_rename(&p, &format!("{i}\n")); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "第 {i} 次 rename 保存没被发现" + ); + settle(&mut rx).await; + } + } + + /// 原地写也算：追加、截断都不经过 rename。**换过一次文件之后还看得见** —— 盯着文件本身 + /// 的实现（macOS）手上的是换掉之前的那个 inode，要重新挂到新的上 + #[tokio::test] + async fn writing_in_place_is_seen_also_after_the_file_was_replaced() { + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + + append(&p, "1\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "追加" + ); + settle(&mut rx).await; + + save_by_rename(&p, "2\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换文件" + ); + settle(&mut rx).await; + + append(&p, "3\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "换过之后的追加" + ); + settle(&mut rx).await; + + // 只截断、不写：大小变了，别的什么都没动 + std::fs::File::create(&p).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "截断" + ); + } + + /// 后来才有的文件：建出来算一次，之后原地写也算 + #[tokio::test] + async fn a_file_created_later_is_seen_and_so_are_its_writes() { + let d = scratch(); + let p = d.path().join("a.md"); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + std::fs::write(&p, "0\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "新建" + ); + settle(&mut rx).await; + append(&p, "1\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "追加" + ); + } + + /// 删掉算一次，再建出来又算一次，再建出来的那个原地写也算 + #[tokio::test] + async fn a_deleted_file_that_comes_back_is_seen_again() { + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + std::fs::remove_file(&p).unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "删除" + ); + settle(&mut rx).await; + std::fs::write(&p, "1\n").unwrap(); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "再建" + ); + settle(&mut rx).await; + append(&p, "2\n"); + assert_eq!( + recv(&mut rx, Duration::from_secs(3)).await, + Signal::Got, + "追加" + ); + } + + /// **读文件不算。**调用方收到信号就去读：读也算的话，读一次就又叫醒自己一次 + #[tokio::test] + async fn reading_a_watched_file_stays_quiet() { + let d = scratch(); + let p = d.path().join("a.md"); + std::fs::write(&p, "0\n").unwrap(); + let (_w, mut rx) = watch(&[d.path().to_path_buf()], WINDOW, only_md).unwrap(); + for _ in 0..20 { + std::fs::read_to_string(&p).unwrap(); + std::fs::metadata(&p).unwrap(); + } + std::fs::read_dir(d.path()).unwrap().for_each(drop); + assert_eq!( + recv(&mut rx, Duration::from_millis(800)).await, + Signal::Quiet, + "读文件把它叫醒了" + ); + } } diff --git a/crates/tw-watch/tests/fd_budget.rs b/crates/tw-watch/tests/fd_budget.rs new file mode 100644 index 00000000..9b0244f6 --- /dev/null +++ b/crates/tw-watch/tests/fd_budget.rs @@ -0,0 +1,122 @@ +//! macOS 上 fd 名额紧的时候：kqueue 每盯一个目录、一份文件占一个 fd，而从访达打开的应用 +//! 只有 256 个。名额不够就整个换回 FSEvents —— 费电，但一个改动都不漏。 +//! +//! **单独一个测试进程、只有一个测试**：要调低整个进程的 fd 软上限、数整个进程开着的 fd， +//! 和别的测试放在一起会互相连累。 +#![cfg(target_os = "macos")] + +use std::path::{Path, PathBuf}; +use std::time::Duration; + +const WINDOW: Duration = Duration::from_millis(200); + +fn only_md(p: &Path) -> bool { + p.extension().is_some_and(|e| e == "md") +} + +/// 此刻整个进程开着几个 fd（列 `/dev/fd` 自己也占一个，每次都一样） +fn open_fds() -> usize { + std::fs::read_dir("/dev/fd").unwrap().count() +} + +/// 把 fd 软上限调到 `n`：监听最多能占一半 +fn set_soft_limit(n: libc::rlim_t) { + let mut r = libc::rlimit { + rlim_cur: 0, + rlim_max: 0, + }; + // SAFETY: 两个调用都只读写传进去的结构体 + unsafe { + assert_eq!(libc::getrlimit(libc::RLIMIT_NOFILE, &mut r), 0); + r.rlim_cur = n; + assert_eq!(libc::setrlimit(libc::RLIMIT_NOFILE, &r), 0); + } +} + +fn append(p: &Path, text: &str) { + use std::io::Write as _; + let mut f = std::fs::OpenOptions::new().append(true).open(p).unwrap(); + f.write_all(text.as_bytes()).unwrap(); +} + +/// `n` 个目录，各有 `files` 份算数的文件 +fn dirs(root: &Path, prefix: &str, n: usize, files: usize) -> Vec { + (0..n) + .map(|i| { + let d = root.join(format!("{prefix}{i}")); + std::fs::create_dir_all(&d).unwrap(); + for j in 0..files { + std::fs::write(d.join(format!("{j}.md")), "0\n").unwrap(); + } + d + }) + .collect() +} + +async fn got(rx: &mut tokio::sync::mpsc::Receiver<()>) -> bool { + matches!( + tokio::time::timeout(Duration::from_secs(5), rx.recv()).await, + Ok(Some(())) + ) +} + +/// 等前面那些改动的信号都收完、静下来 +async fn settle(rx: &mut tokio::sync::mpsc::Receiver<()>) { + while let Ok(Some(())) = tokio::time::timeout(WINDOW * 3, rx.recv()).await {} +} + +#[test] +fn with_too_few_fds_it_falls_back_to_fsevents_and_misses_nothing() { + // 名额 32 + set_soft_limit(64); + let rt = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + rt.block_on(async { + let root = tempfile::tempdir().unwrap(); + let before = open_fds(); + + // 8 个目录、16 份文件，名额够：都用 kqueue 盯，用完都还回去 + let some = dirs(root.path(), "some", 8, 2); + let (w, mut rx) = tw_watch::watch(&some, WINDOW, only_md).unwrap(); + assert_eq!( + open_fds(), + before + 1 + 8 + 16, + "kqueue、8 个目录、16 份文件" + ); + append(&some[3].join("1.md"), "1\n"); + assert!(got(&mut rx).await, "原地写"); + settle(&mut rx).await; + + // 又多出 10 份：超了名额，换回 FSEvents，手上的都放掉（只剩 kqueue 自己） + for j in 2..12 { + std::fs::write(some[0].join(format!("{j}.md")), "0\n").unwrap(); + } + assert!(got(&mut rx).await, "超名额的那一批新文件"); + settle(&mut rx).await; + assert_eq!(open_fds(), before + 1, "换回 FSEvents 之后还拿着 fd"); + // 换过去之后照样一个不漏：新文件的原地写、rename 保存 + append(&some[0].join("7.md"), "1\n"); + assert!(got(&mut rx).await, "换回 FSEvents 之后的原地写"); + settle(&mut rx).await; + std::fs::write(some[5].join("x.tmp"), "2\n").unwrap(); + std::fs::rename(some[5].join("x.tmp"), some[5].join("0.md")).unwrap(); + assert!(got(&mut rx).await, "换回 FSEvents 之后的 rename 保存"); + drop(w); + assert_eq!(open_fds(), before, "扔掉之后还拿着 fd"); + + // 一开始就不够：40 个目录，直接用 FSEvents + let many = dirs(root.path(), "many", 40, 1); + let (w, mut rx) = tw_watch::watch(&many, WINDOW, only_md).unwrap(); + assert_eq!(open_fds(), before, "名额不够还占着 fd"); + settle(&mut rx).await; + append(&many[39].join("0.md"), "1\n"); + assert!(got(&mut rx).await, "FSEvents 盯着的原地写"); + drop(w); + + // 名额还回来了：又能用 kqueue 盯 + let (_w, _rx) = tw_watch::watch(&some[..2], WINDOW, only_md).unwrap(); + assert!(open_fds() > before + 2, "名额没还回来"); + }); +}