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)?,