From 2ecd1c4fec0a01048b8f39b88aaa6949449a92c0 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 10 Oct 2026 18:27:45 +0800 Subject: [PATCH] feat(control): per-turn context window and what fills it A session's turns only carried their measured usage. The desktop app wants to show how full a turn's context window was and what was occupying it, so: - `TurnView` gains `context_window`: the window of the model that answered the turn, resolved now through the one spec resolution core already has (`tw_config::model_specs`: the upstream's manual spec first, then the price table), looked up by the model name the upstream actually received (`sent_model`, now carried by `TurnRow`). - New endpoint `GET /sessions/{id}/turns/{turn}/context` answering `TurnContext`: the stored request body decoded with the client's dialect and estimated in one pass as four parts (system prompt, tool definitions, every message before the last user message, the last user message with the tool results it carries) by the same estimator `count_tokens` answers with, next to the measured input and cache-read tokens from the history row. A body that was never kept, has expired, was truncated or cannot be decoded answers `kept: false` with `parts: null`. An unknown session is `control.session_not_found`; a turn that is not in that session is `control.request_not_found`. - `tw_engine::estimate_parts` is the per-part twin of `estimate_tokens`; each part rounds on its own, so the sum differs from the whole-body estimate by at most a few tokens. - Protocol version 45. tw-control now depends on tw-dialect directly to decode stored bodies. Co-Authored-By: Claude Fable 5.1 --- Cargo.lock | 1 + crates/tw-api/src/ep.rs | 3 + crates/tw-api/src/lib.rs | 53 +++- crates/tw-api/src/ts.rs | 31 +++ crates/tw-control/Cargo.toml | 2 + crates/tw-control/src/lib.rs | 100 +++++++- crates/tw-control/tests/turn_context.rs | 308 ++++++++++++++++++++++++ crates/tw-engine/src/facts.rs | 101 +++++++- crates/tw-engine/src/lib.rs | 2 +- crates/tw-store/src/db.rs | 5 +- 10 files changed, 600 insertions(+), 6 deletions(-) create mode 100644 crates/tw-control/tests/turn_context.rs diff --git a/Cargo.lock b/Cargo.lock index 30c8835a..2d231186 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3235,6 +3235,7 @@ dependencies = [ "tw-api", "tw-bedrock", "tw-config", + "tw-dialect", "tw-engine", "tw-gateway", "tw-guard", diff --git a/crates/tw-api/src/ep.rs b/crates/tw-api/src/ep.rs index eddf78ac..0d3dcbc8 100644 --- a/crates/tw-api/src/ep.rs +++ b/crates/tw-api/src/ep.rs @@ -76,6 +76,9 @@ endpoints! { /// 一次会话读成一段对话：每一轮新说的话、回答、工具调用和结果（已脱敏）。可以只要 /// 从某一轮起的那些 SessionTranscript: GET "/sessions/{id}/transcript" [id], api::TranscriptQuery => api::Transcript; + /// 一次会话里的一轮，上下文由什么占着：系统提示、工具、历史、最后一条用户消息各估多少 + /// token，连同实测的输入和缓存读。会话不存在、那一轮不在这次会话里都是 404 + SessionTurnContext: GET "/sessions/{id}/turns/{turn}/context" [id, turn], () => api::TurnContext; /// 中止一次会话里所有在跑的请求。一个都没有是 404 AbortSession: POST "/sessions/{id}/abort" [id], () => api::Aborted; diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index 093507e4..ac6ead23 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -868,7 +868,16 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// 客户端收到的话、删掉了几段，[`OutcomeDetail`]）。[`RequestDetail`] 的 `security` 是同一个 /// 类型。事件跟着改：[`SecretItem`]、[`Event::ContentMatched`]、[`Event::ToolCallFlagged`] /// 多了 `detail`（[`SecurityHitDetail`]）。细节存在请求记录的库里，和请求一起过期。 -pub const CONTROL_API_VERSION: u32 = 44; +/// +/// **45 起会话里的每一轮说得出上下文窗口有多大、被什么占着**：[`TurnView`] 多了 +/// `context_window`（答这一轮的那家那个模型的上下文窗口，按**此刻**知道的规格：手写的 +/// 优先，其次价目表；不知道是 null）。新端点 `GET /sessions/{id}/turns/{turn}/context` +/// （→ [`TurnContext`]）：这一轮存下来的请求体按系统提示、工具定义、历史、最后一条用户消息 +/// 四块各估多少 token（和数 token 由网关自己答的是同一个估算），连同实测的输入和缓存读。 +/// 请求体没存下来、已经清掉、只存了开头或读不懂的，`kept` 是 false、`parts` 是 null。 +/// 会话不存在是 404（`control.session_not_found`），那一轮不在这次会话里也是 404 +/// （`control.request_not_found`）。照 44 写的界面读不到 `context_window`。 +pub const CONTROL_API_VERSION: u32 = 45; #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] @@ -5019,6 +5028,10 @@ pub struct TurnView { /// 服务它的那家怎么收钱，和 `HistoryRow::billing` 同一套词：`per-token` / /// `free` pub billing: Billing, + /// 答这一轮的那家那个模型的上下文窗口（token），按**此刻**知道的规格：这一家手写的 + /// `model_specs` 优先，其次价目表。不知道是 None。**不是那一轮当时的数** —— 规格改了 + /// 这里跟着变 + pub context_window: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -5028,6 +5041,44 @@ pub struct SessionDetail { pub turns: Vec, } +/// 一轮的上下文由什么占着（`GET /sessions/{id}/turns/{turn}/context`）。 +/// +/// `parts` 是**本地的估算**，在存下来的请求体上算，和数 token 由网关自己答的那个数是 +/// 同一个算法（`tw_engine::estimate_tokens`）；`input_tokens`、`cache_read_tokens` 是上游 +/// 实测报的。两个数差一两成是正常的，界面别把估算画成实测。 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct TurnContext { + /// 请求体还在不在。没存下来、已经清掉（`retention.body_days`）、只存了开头、读不懂 + /// 的都是 false，那时 `parts` 是 null + pub kept: bool, + /// 和 [`TurnView::context_window`] 同一个数 + pub window: Option, + /// 这一轮实测的输入 token（上游按它收费），记录里的那个 + pub input_tokens: Option, + /// 这一轮实测的缓存读 token，记录里的那个 + pub cache_read_tokens: Option, + /// 四块各估多少。请求体不在了是 null + pub parts: Option, +} + +/// 请求体的四块各估多少 token。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "ts", derive(ts_rs::TS))] +pub struct ContextParts { + /// 系统提示：Anthropic 的 `system`、Chat 和 Responses 的 system、developer 消息、 + /// Responses 的 `instructions`、Gemini 的 `systemInstruction` + pub system: u64, + /// 工具定义 + pub tools: u64, + /// 最后一条用户消息之前的每一条消息 + pub history: u64, + /// 最后一条用户消息，连同它带着的工具结果 + pub last_user: u64, + /// 四块之和 + pub total: u64, +} + /// 一次会话读成一段对话（`GET /sessions/{id}/transcript`）：每一轮新说的话、回答、推理、 /// 工具调用和工具结果。 /// diff --git a/crates/tw-api/src/ts.rs b/crates/tw-api/src/ts.rs index a1196fc6..da0c5c73 100644 --- a/crates/tw-api/src/ts.rs +++ b/crates/tw-api/src/ts.rs @@ -296,6 +296,37 @@ mod tests { assert!(event.contains("answered_model?: string"), "{event}"); } + /// 一轮的上下文：两个路径参数都在，`parts` 可以是 null，四块和合计都是数 + #[test] + fn a_turn_context_has_its_parts_or_null() { + let ts = typescript(); + assert!(ts.contains(" SessionTurnContext: { req: null; res: TurnContext };")); + assert!(ts.contains( + " SessionTurnContext: { method: \"GET\", path: \"/sessions/{id}/turns/{turn}/context\", params: [\"id\", \"turn\"], format: \"json\" }," + )); + let context = decl_of(&ts, "TurnContext"); + for field in [ + "kept: boolean", + "window: number | null", + "input_tokens: number | null", + "cache_read_tokens: number | null", + "parts: ContextParts | null", + ] { + assert!(context.contains(field), "{field} in {context}"); + } + let parts = decl_of(&ts, "ContextParts"); + for field in [ + "system: number", + "tools: number", + "history: number", + "last_user: number", + "total: number", + ] { + assert!(parts.contains(field), "{field} in {parts}"); + } + assert!(decl_of(&ts, "TurnView").contains("context_window: number | null")); + } + /// 对话记录：请求号是字符串；可以为空的是 null（必有的字段，不是省掉）；块是按 `kind` /// 分的联合，角色和缺口是字面量 #[test] diff --git a/crates/tw-control/Cargo.toml b/crates/tw-control/Cargo.toml index 151a148c..eb339a88 100644 --- a/crates/tw-control/Cargo.toml +++ b/crates/tw-control/Cargo.toml @@ -23,6 +23,8 @@ tw-config = { workspace = true } tw-guard = { workspace = true } bytes = { workspace = true } tw-engine = { workspace = true } +# 一轮的上下文按块估：存下来的请求体按客户端的格式解码 +tw-dialect = { workspace = true } tw-gateway = { workspace = true } # 重放发给 Bedrock 的请求要签名、拆帧 tw-bedrock = { workspace = true } diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index 6fe26133..c54a6fcc 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -140,6 +140,7 @@ pub fn router(state: ControlState) -> Router { .at(ep::Sessions, sessions) .at(ep::SessionDetail, session_detail) .at(ep::SessionTranscript, session_transcript) + .at(ep::SessionTurnContext, session_turn_context) .at(ep::AbortSession, abort_session) .at(ep::DryRun, dryrun::dry_run) // 为客户端发专用密钥。接管本身在桌面端做 @@ -1735,8 +1736,24 @@ fn session_view(s: &tw_store::db::SessionRow) -> tw_api::SessionView { } } -fn turn_view(t: &tw_store::db::TurnRow) -> tw_api::TurnView { +/// 答这一轮的那家那个模型，**此刻**的上下文窗口：和 `/v1/models`、上游页、别名同一个查法 +/// （`tw_config::model_specs`），按发给上游的名字查 —— 上游收到的是它 +fn context_window( + cfg: &tw_config::Config, + book: &tw_pricing::PriceBook, + provider: &str, + sent_model: &str, +) -> Option { + cfg.model_spec(book, provider, sent_model).context_window() +} + +fn turn_view( + cfg: &tw_config::Config, + book: &tw_pricing::PriceBook, + t: &tw_store::db::TurnRow, +) -> tw_api::TurnView { tw_api::TurnView { + context_window: context_window(cfg, book, &t.provider, &t.sent_model), id: t.id, at_ms: t.at_ms as u64, model: t.model.clone(), @@ -1794,12 +1811,91 @@ async fn session_detail( .map_err(records)? }; let session = session.ok_or_else(|| no_such_session(&id))?; + let (cfg, book) = (s.config(), s.gateway.pricing.load()); Ok(Json(tw_api::SessionDetail { session: session_view(&session), - turns: turns.iter().map(turn_view).collect(), + turns: turns.iter().map(|t| turn_view(&cfg, &book, t)).collect(), })) } +/// 一轮的上下文由什么占着（见 [`tw_api::TurnContext`]）。 +/// +/// 请求体按客户端的格式解码成中间表示，再按块估 token（`tw_engine::estimate_parts`）： +/// 四种格式的系统提示、工具、消息都读到同一个地方，和数 token 由网关自己答的是同一个 +/// 估算。**正文读和解析在阻塞线程上**，锁只在取行的时候拿：一份请求体可以有几 MB。 +async fn session_turn_context( + State(s): State, + axum::extract::Path((id, turn)): axum::extract::Path<(String, i64)>, +) -> Result, Fail> { + let store = need_store(&s)?; + let (row, blobs) = { + let id = id.clone(); + on_store(store, move |g| { + let row = g + .db() + .get(turn)? + .filter(|r| !r.local && r.session.as_deref() == Some(id.as_str())); + let Some(row) = row else { + // 会话不存在和那一轮不在这次会话里是两个 404 + return Ok(Err(match g.db().session(&id)? { + Some(_) => fail( + StatusCode::NOT_FOUND, + msg!("control.request_not_found", id = turn => "There is no request {id}."), + ), + None => no_such_session(&id), + })); + }; + // 正文目录只是一个路径，读它不需要锁 + let blobs = tw_store::Blobs::new(g.blobs().root().to_path_buf()); + Ok::<_, tw_store::DbError>(Ok((row, blobs))) + }) + .await? + .map_err(records)?? + }; + let window = { + let (cfg, book) = (s.config(), s.gateway.pricing.load()); + context_window(&cfg, &book, &row.provider, &row.sent_model) + }; + let (input_tokens, cache_read_tokens) = (row.input_tokens, row.cache_read_tokens); + let parts = tokio::task::spawn_blocking(move || stored_parts(&row, &blobs)) + .await + .map_err(internal)?; + Ok(Json(tw_api::TurnContext { + kept: parts.is_some(), + window, + input_tokens, + cache_read_tokens, + parts, + })) +} + +/// 存下来的请求体按块估。没存、清掉了、只存了开头、解析不了、解码不了的都是 None。 +fn stored_parts( + row: &tw_store::db::RequestRow, + blobs: &tw_store::Blobs, +) -> Option { + let raw = blobs.get(row.at_ms, row.id, tw_store::Which::Request)?; + if blobs + .original_len(row.at_ms, row.id, tw_store::Which::Request) + .is_some_and(|n| n > raw.len()) + { + return None; + } + let v: serde_json::Value = serde_json::from_slice(&raw).ok()?; + let client = tw_store::search::text::client_dialect(&row.path)?; + let r = tw_dialect::convert::decode(client, &v, &row.path, None) + .ok()? + .request; + let p = tw_engine::estimate_parts(&r); + Some(tw_api::ContextParts { + system: p.system, + tools: p.tools, + history: p.history, + last_user: p.last_user, + total: p.total(), + }) +} + /// 中止一个在跑的请求（见 `tw_gateway::abort`）。**叫停是立刻的**：和上游的连接在请求自己 /// 那边断开，结局（`RequestFailed`，来源 `aborted`）随后从事件流上到。已经结束了的、从来 /// 没有过的、跑在 WebSocket 连接上的，都是不在跑：404 diff --git a/crates/tw-control/tests/turn_context.rs b/crates/tw-control/tests/turn_context.rs new file mode 100644 index 00000000..feb388d7 --- /dev/null +++ b/crates/tw-control/tests/turn_context.rs @@ -0,0 +1,308 @@ +//! 一轮的上下文端点（`GET /sessions/{id}/turns/{turn}/context`）和会话详情里每一轮的 +//! `context_window`。 +//! +//! 按块怎么估在 tw-engine 里测；这里盯的是**界面拿到的那份 JSON**：四种格式的请求体都读得 +//! 出四块、请求体不在了是什么样、两种 404 各是哪个码。 + +use std::sync::Arc; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use serde_json::{Value, json}; +use tower::ServiceExt; +use tw_control::{ConfigManager, ControlState}; +use tw_store::Which; + +const BASE: &str = "version: 1\nlisten:\n control:\n key: c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00c0ffee00\nclients:\n - name: 我\n key: tw-一把钥匙就够\nproviders:\n - name: 官方\n base_url: https://api.anthropic.com\n key: sk-x\n model_specs:\n handwritten: { context_window: 4096 }\n"; + +/// 一轮：`path` 说客户端是哪种格式 +fn turn(id: i64, session: &str, model: &str, path: &str) -> tw_store::db::RequestRow { + tw_store::db::RequestRow { + session_log_bytes: None, + key_masked: None, + peer: None, + id, + at_ms: 1_790_000_000_000 + id, + client: "claude-code".into(), + client_hint: None, + session: Some(session.into()), + provider: "官方".into(), + model: model.into(), + sent_model: model.into(), + answered_model: None, + path: path.into(), + status: Some(200), + ttfb_ms: Some(100), + ttft_ms: None, + duration_ms: Some(200), + tokens_per_sec: None, + sent_bytes: Some(10), + received_bytes: Some(10), + egress: None, + input_tokens: Some(1234), + output_tokens: Some(20), + cache_read_tokens: Some(1000), + cache_write_tokens: None, + input_estimate: None, + cost_micros: Some(1_000), + cost_estimated: false, + error: None, + local: false, + cancelled: false, + routing: None, + billing: tw_api::Billing::PerToken, + cache_saved_micros: None, + price_source: None, + translated: None, + } +} + +/// 一个带着请求库和正文的控制面。`bodies`：哪一条的请求体（没存下的不在里面） +fn app( + rows: &[tw_store::db::RequestRow], + bodies: &[(i64, Value)], +) -> (tempfile::TempDir, axum::Router) { + let d = tempfile::tempdir().unwrap(); + let p = d.path().join("config.yaml"); + std::fs::write(&p, BASE).unwrap(); + let db = tw_store::Db::open(&d.path().join("data.db")).unwrap(); + for r in rows { + db.insert(r).unwrap(); + } + let blobs = tw_store::Blobs::new(d.path().join("blobs")); + for (id, request) in bodies { + let at = rows.iter().find(|r| r.id == *id).unwrap().at_ms; + assert!(blobs.put(at, *id, Which::Request, request.to_string().as_bytes())); + } + let rec = tw_store::Recorder::new( + db, + blobs, + tw_pricing::shared(tw_pricing::PriceBook::builtin().unwrap()), + ); + let cfg: tw_config::Config = serde_yaml_ng::from_str(BASE).unwrap(); + let gw = tw_gateway::AppState::new(cfg).unwrap(); + let bus = gw.bus.clone(); + let state = ControlState { + shutdown: Default::default(), + remote: Default::default(), + cfg: Arc::new(ConfigManager::new(p, gw.clone(), bus)), + gateway: gw, + store: Some(Arc::new(tokio::sync::Mutex::new(rec))), + started: std::time::Instant::now(), + price_updater: Default::default(), + chatgpt: Default::default(), + zai: Default::default(), + }; + (d, tw_control::router(state)) +} + +async fn get(app: &axum::Router, path: &str) -> (StatusCode, Value) { + let r = app + .clone() + .oneshot(Request::builder().uri(path).body(Body::empty()).unwrap()) + .await + .unwrap(); + let st = r.status(); + let b = axum::body::to_bytes(r.into_body(), 1 << 20).await.unwrap(); + (st, serde_json::from_slice(&b).unwrap()) +} + +/// 四块都是正数、加起来是 `total` +fn well_formed(parts: &Value) { + let n = |k: &str| { + parts[k] + .as_u64() + .unwrap_or_else(|| panic!("{k} in {parts}")) + }; + assert!(n("system") > 0, "{parts}"); + assert!(n("tools") > 0, "{parts}"); + assert!(n("history") > 0, "{parts}"); + assert!(n("last_user") > 0, "{parts}"); + assert_eq!( + n("total"), + n("system") + n("tools") + n("history") + n("last_user") + ); +} + +/// Anthropic：`system`、`tools`、消息；最后一条用户消息带着工具结果。实测的两个数来自 +/// 记录，窗口来自价目表，和会话详情里那一轮的 `context_window` 是同一个数 +#[tokio::test] +async fn an_anthropic_body_is_split_into_four_parts() { + let body = json!({"model": "claude-sonnet-4-5", "max_tokens": 100, + "system": "You are a careful assistant.", + "tools": [{"name": "Read", "description": "read a file", "input_schema": {"type": "object"}}], + "messages": [ + {"role": "user", "content": "first question about the code"}, + {"role": "assistant", "content": [{"type": "text", "text": "let me look"}, + {"type": "tool_use", "id": "t1", "name": "Read", "input": {"p": "a.rs"}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "t1", + "content": "fn main() {}"}]} + ]}); + let (_d, app) = app( + &[turn(1, "s1", "claude-sonnet-4-5", "/v1/messages")], + &[(1, body)], + ); + + let (st, v) = get(&app, "/sessions/s1/turns/1/context").await; + assert_eq!(st, StatusCode::OK, "{v}"); + assert_eq!(v["kept"], true); + assert_eq!(v["input_tokens"], 1234); + assert_eq!(v["cache_read_tokens"], 1000); + let window = v["window"] + .as_u64() + .expect("the price table knows the model"); + assert!(window >= 200_000, "{v}"); + well_formed(&v["parts"]); + // 系统提示 28 个 ASCII 字符 + assert_eq!(v["parts"]["system"], 7); + + let (st, detail) = get(&app, "/sessions/s1").await; + assert_eq!(st, StatusCode::OK); + assert_eq!(detail["turns"][0]["context_window"], window); +} + +/// Chat：system、developer 消息算系统提示；工具结果的消息是用户角色，所以最后一条 +/// 用户消息就是那条工具结果 +#[tokio::test] +async fn a_chat_body_is_split_into_four_parts() { + let body = json!({"model": "gpt-5", "stream": false, + "tools": [{"type": "function", "function": {"name": "read", "description": "read", + "parameters": {"type": "object"}}}], + "messages": [ + {"role": "system", "content": "You are a careful assistant."}, + {"role": "developer", "content": "Answer briefly."}, + {"role": "user", "content": "first question about the code"}, + {"role": "assistant", "content": null, "tool_calls": [{"id": "c1", "type": "function", + "function": {"name": "read", "arguments": "{\"p\":\"a.rs\"}"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "fn main() {}"} + ]}); + let (_d, app) = app( + &[turn(1, "s1", "gpt-5", "/v1/chat/completions")], + &[(1, body)], + ); + + let (st, v) = get(&app, "/sessions/s1/turns/1/context").await; + assert_eq!(st, StatusCode::OK, "{v}"); + assert_eq!(v["kept"], true); + well_formed(&v["parts"]); + // 28 + 15 个 ASCII 字符，两段各自取整 + assert_eq!(v["parts"]["system"], 7 + 4); + // 最后一条用户消息只有那条工具结果："fn main() {}" 12 个字符加上消息开销 3 + assert_eq!(v["parts"]["last_user"], 3 + 3); +} + +/// Responses：`instructions` 和 developer 项算系统提示，`function_call_output` 是用户角色 +#[tokio::test] +async fn a_responses_body_is_split_into_four_parts() { + let body = json!({"model": "gpt-5-codex", "instructions": "You are a careful assistant.", + "tools": [{"type": "function", "name": "read", "description": "read", + "parameters": {"type": "object"}}], + "input": [ + {"role": "developer", "content": "Answer briefly."}, + {"role": "user", "content": "first question about the code"}, + {"type": "function_call", "call_id": "c1", "name": "read", "arguments": "{\"p\":\"a.rs\"}"}, + {"type": "function_call_output", "call_id": "c1", "output": "fn main() {}"} + ]}); + let (_d, app) = app( + &[turn(1, "s1", "gpt-5-codex", "/v1/responses")], + &[(1, body)], + ); + + let (st, v) = get(&app, "/sessions/s1/turns/1/context").await; + assert_eq!(st, StatusCode::OK, "{v}"); + assert_eq!(v["kept"], true); + well_formed(&v["parts"]); + assert_eq!(v["parts"]["system"], 7 + 4); + assert_eq!(v["parts"]["last_user"], 3 + 3); +} + +/// Gemini：`systemInstruction`、`tools[].functionDeclarations`、`contents`；`functionResponse` +/// 在用户角色里 +#[tokio::test] +async fn a_gemini_body_is_split_into_four_parts() { + let body = json!({ + "systemInstruction": {"parts": [{"text": "You are a careful assistant."}]}, + "tools": [{"functionDeclarations": [{"name": "read", "description": "read", + "parameters": {"type": "object"}}]}], + "contents": [ + {"role": "user", "parts": [{"text": "first question about the code"}]}, + {"role": "model", "parts": [{"functionCall": {"name": "read", "args": {"p": "a.rs"}}}]}, + {"role": "user", "parts": [{"functionResponse": {"name": "read", + "response": {"output": "fn main() {}"}}}]} + ]}); + let (_d, app) = app( + &[turn( + 1, + "s1", + "gemini-2.5-pro", + "/v1beta/models/gemini-2.5-pro:generateContent", + )], + &[(1, body)], + ); + + let (st, v) = get(&app, "/sessions/s1/turns/1/context").await; + assert_eq!(st, StatusCode::OK, "{v}"); + assert_eq!(v["kept"], true); + well_formed(&v["parts"]); + assert_eq!(v["parts"]["system"], 7); +} + +/// 请求体没存下来（或者清掉了）：`kept` 是 false、`parts` 是 null，别的照给 +#[tokio::test] +async fn a_turn_whose_body_is_gone_says_so() { + let (_d, app) = app(&[turn(1, "s1", "claude-sonnet-4-5", "/v1/messages")], &[]); + + let (st, v) = get(&app, "/sessions/s1/turns/1/context").await; + assert_eq!(st, StatusCode::OK, "{v}"); + assert_eq!(v["kept"], false); + assert!(v["parts"].is_null(), "{v}"); + assert_eq!(v["input_tokens"], 1234); + assert_eq!(v["cache_read_tokens"], 1000); + assert!(v["window"].as_u64().is_some(), "{v}"); +} + +/// 手写的规格优先于价目表；价目表也不认识的模型，窗口是 null +#[tokio::test] +async fn the_window_is_the_spec_known_now() { + let (_d, app) = app( + &[ + turn(1, "s1", "handwritten", "/v1/messages"), + turn(2, "s1", "nobody-knows-this-model", "/v1/messages"), + ], + &[], + ); + + let (_, v) = get(&app, "/sessions/s1/turns/1/context").await; + assert_eq!(v["window"], 4096); + let (_, v) = get(&app, "/sessions/s1/turns/2/context").await; + assert!(v["window"].is_null(), "{v}"); + + let (_, detail) = get(&app, "/sessions/s1").await; + assert_eq!(detail["turns"][0]["context_window"], 4096); + assert!(detail["turns"][1]["context_window"].is_null(), "{detail}"); +} + +/// 会话不存在和那一轮不在这次会话里是两个 404，各用会话端点已有的码 +#[tokio::test] +async fn unknown_sessions_and_turns_are_two_different_404s() { + let (_d, app) = app( + &[ + turn(1, "s1", "claude-sonnet-4-5", "/v1/messages"), + turn(2, "s2", "claude-sonnet-4-5", "/v1/messages"), + ], + &[], + ); + + let (st, v) = get(&app, "/sessions/nope/turns/1/context").await; + assert_eq!(st, StatusCode::NOT_FOUND, "{v}"); + assert_eq!(v["code"], "control.session_not_found"); + + // 第 2 轮是别的会话的 + let (st, v) = get(&app, "/sessions/s1/turns/2/context").await; + assert_eq!(st, StatusCode::NOT_FOUND, "{v}"); + assert_eq!(v["code"], "control.request_not_found"); + + let (st, v) = get(&app, "/sessions/s1/turns/99/context").await; + assert_eq!(st, StatusCode::NOT_FOUND, "{v}"); + assert_eq!(v["code"], "control.request_not_found"); +} diff --git a/crates/tw-engine/src/facts.rs b/crates/tw-engine/src/facts.rs index 6d6a1c85..3458d231 100644 --- a/crates/tw-engine/src/facts.rs +++ b/crates/tw-engine/src/facts.rs @@ -6,7 +6,7 @@ //! 下面这些维度，没有一个在网络代理里有对应物。 use serde::{Deserialize, Serialize}; -use tw_dialect::ir::{Part, Request, ToolInput, ToolKind}; +use tw_dialect::ir::{Part, Request, Role, ToolInput, ToolKind}; #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct RequestFacts { @@ -118,6 +118,76 @@ pub fn estimate_tokens(r: &Request) -> u64 { text.tokens() + fixed } +/// 一个请求的输入 token 分成几块各是多少（估算，同 [`estimate_tokens`] 的算法）。 +/// +/// 会话里看一轮的上下文由什么占着：系统提示、工具定义、前面的历史、最后一条用户消息。 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct TokenParts { + /// 系统提示：开头的那些（[`Request::system`]）加上对话中途的系统消息 + pub system: u64, + /// 工具定义 + pub tools: u64, + /// 最后一条用户消息之前的每一条消息（上游回过的话、更早的用户消息和工具结果） + pub history: u64, + /// 最后一条用户消息，连同它带着的工具结果。一条用户消息都没有是 0 + pub last_user: u64, +} + +impl TokenParts { + /// 四块加起来 + pub fn total(&self) -> u64 { + self.system + self.tools + self.history + self.last_user + } +} + +/// 估一个请求的输入 token 数，按块分开。 +/// +/// 算法和 [`estimate_tokens`] 一样，一遍走完；每一块各自取整，所以四块之和和整体估的数 +/// 最多差几个 token。 +/// +/// 「最后一条用户消息」是角色为用户的最后一条：Chat、Responses 里只装着工具结果的消息 +/// 解码后也是用户角色，所以紧跟在工具调用之后的那一轮，这一块就是那些工具结果。排在它 +/// 后面的（助手的预填）算进历史。 +pub fn estimate_parts(r: &Request) -> TokenParts { + let mut system = Text::default(); + let mut tools = Text::default(); + let mut history = Text::default(); + let mut last_user = Text::default(); + let mut fixed = TokenParts::default(); + for s in &r.system { + system.add(s); + } + for t in &r.tools { + tools.add(&t.name); + if let Some(d) = &t.description { + tools.add(d); + } + if let ToolKind::Function { schema, .. } = &t.kind { + tools.add_json(schema); + } + } + let last = r.messages.iter().rposition(|m| m.role == Role::User); + for (i, m) in r.messages.iter().enumerate() { + let (text, count) = if m.role == Role::System { + (&mut system, &mut fixed.system) + } else if Some(i) == last { + (&mut last_user, &mut fixed.last_user) + } else { + (&mut history, &mut fixed.history) + }; + *count += PER_MESSAGE; + for p in &m.parts { + *count += part(text, p); + } + } + TokenParts { + system: system.tokens() + fixed.system, + tools: tools.tokens() + fixed.tools, + history: history.tokens() + fixed.history, + last_user: last_user.tokens() + fixed.last_user, + } +} + /// 一段 JSON 里所有字符串值，按 [`estimate_tokens`] 的算法估成 token。 /// /// 数 token 的请求体解不开时用：**数 token 不该因为我们的解析器不认识某个字段就失败**， @@ -444,6 +514,35 @@ mod tests { } } + /// 按块分：系统提示、工具、历史、最后一条用户消息各归各的，最后一条用户消息带着的 + /// 工具结果算它的；四块之和和整体估的数只差取整 + #[test] + fn parts_split_the_body_at_the_last_user_message() { + let v: serde_json::Value = serde_json::from_str( + r#"{"model":"m","max_tokens":1,"system":"abcdefgh", + "tools":[{"name":"Read","description":"abcd","input_schema":{"type":"object"}}], + "messages":[ + {"role":"user","content":"first question"}, + {"role":"assistant","content":[{"type":"text","text":"ok"}, + {"type":"tool_use","id":"t1","name":"Read","input":{"p":"a.rs"}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"t1","content":"the file 的内容"}]} + ]}"#, + ) + .unwrap(); + let r = tw_dialect::anthropic::decode_request(&v, &mut Dropped::new(Dialect::Anthropic)) + .unwrap(); + let p = estimate_parts(&r); + assert_eq!(p.system, 2); + // Read + abcd + {"type":"object"} = 4 + 4 + 17 + assert_eq!(p.tools, 25_u64.div_ceil(4)); + // 两条消息各 3 的开销；"first question" 14 + "ok" 2 + Read 4 + {"p":"a.rs"} 12 + assert_eq!(p.history, 6 + 32_u64.div_ceil(4)); + // "the file " 9 + 3 个汉字 = 3 + 3 + 3 + assert_eq!(p.last_user, 3 + 3 + 3); + let whole = estimate_tokens(&r); + assert!(p.total().abs_diff(whole) <= 3, "{p:?} vs {whole}"); + } + /// 解不开的请求体：按里面所有字符串估，不是 0 #[test] fn strings_anywhere_in_a_body_still_give_a_number() { diff --git a/crates/tw-engine/src/lib.rs b/crates/tw-engine/src/lib.rs index 38ab032a..8b0f65ee 100644 --- a/crates/tw-engine/src/lib.rs +++ b/crates/tw-engine/src/lib.rs @@ -12,5 +12,5 @@ pub use engine::{ RuleNotes, SetAction, Target, balance_factors, has_catch_all, is_builtin_group, notes, order_by, scalar_name, }; -pub use facts::{RequestFacts, estimate_strings, estimate_tokens}; +pub use facts::{RequestFacts, TokenParts, estimate_parts, estimate_strings, estimate_tokens}; pub use weighted::Member; diff --git a/crates/tw-store/src/db.rs b/crates/tw-store/src/db.rs index 4c0ba380..3709b62b 100644 --- a/crates/tw-store/src/db.rs +++ b/crates/tw-store/src/db.rs @@ -665,6 +665,8 @@ pub struct TurnRow { pub id: i64, pub at_ms: i64, pub model: String, + /// 发给上游的模型名（见 [`RequestRow::sent_model`]）。规格按它查：上游收到的是它 + pub sent_model: String, pub provider: String, pub input_tokens: Option, pub output_tokens: Option, @@ -784,7 +786,7 @@ impl Db { let mut st = self.conn.prepare( "SELECT id, at_ms, model, provider, input_tokens, output_tokens, cache_read_tokens, cost_micros, duration_ms, error, cancelled, - cost_estimated, billing, error_code, error_args, status + cost_estimated, billing, error_code, error_args, status, sent_model FROM requests WHERE session = ?1 AND local = 0 ORDER BY at_ms, id", )?; let rows = st.query_map([session], |r| { @@ -792,6 +794,7 @@ impl Db { id: r.get(0)?, at_ms: r.get(1)?, model: r.get(2)?, + sent_model: r.get(16)?, provider: r.get(3)?, input_tokens: r.get(4)?, output_tokens: r.get(5)?,