修复: token 汇总正确性(会话级双计/is_estimated跨端/流式usage传递) + is_estimated SELECT 缺失修复
This commit is contained in:
@@ -46,6 +46,7 @@ fn ai_message_from_row(row: &Row<'_>) -> std::result::Result<AiMessageRecord, ru
|
||||
prompt_cache_hit_tokens: row.get("prompt_cache_hit_tokens")?,
|
||||
prompt_cache_miss_tokens: row.get("prompt_cache_miss_tokens")?,
|
||||
reasoning_tokens: row.get("reasoning_tokens")?,
|
||||
is_estimated: row.get("is_estimated")?,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -86,8 +87,8 @@ impl AiMessageRepo {
|
||||
(id, conversation_id, seq, role, content, parts, tool_call_id,
|
||||
tool_calls, model, status, reasoning_content, timestamp, created_at,
|
||||
prompt_tokens, completion_tokens,
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens, is_estimated)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19)",
|
||||
)
|
||||
.map_err(storage_err)?;
|
||||
for rec in &records {
|
||||
@@ -96,7 +97,7 @@ impl AiMessageRepo {
|
||||
rec.parts, rec.tool_call_id, rec.tool_calls, rec.model, rec.status,
|
||||
rec.reasoning_content, rec.timestamp, rec.created_at,
|
||||
rec.prompt_tokens, rec.completion_tokens,
|
||||
rec.prompt_cache_hit_tokens, rec.prompt_cache_miss_tokens, rec.reasoning_tokens
|
||||
rec.prompt_cache_hit_tokens, rec.prompt_cache_miss_tokens, rec.reasoning_tokens, rec.is_estimated
|
||||
])
|
||||
.map_err(storage_err)?;
|
||||
}
|
||||
@@ -122,7 +123,8 @@ impl AiMessageRepo {
|
||||
"SELECT id, conversation_id, seq, role, content, parts, tool_call_id,
|
||||
tool_calls, model, status, reasoning_content, timestamp, created_at,
|
||||
prompt_tokens, completion_tokens,
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens,
|
||||
is_estimated
|
||||
FROM ai_messages WHERE conversation_id = ?1 ORDER BY seq ASC",
|
||||
)
|
||||
.map_err(storage_err)?;
|
||||
@@ -164,13 +166,13 @@ impl AiMessageRepo {
|
||||
"SELECT id, conversation_id, seq, role, content, parts, tool_call_id,
|
||||
tool_calls, model, status, reasoning_content, timestamp, created_at,
|
||||
prompt_tokens, completion_tokens,
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens, is_estimated
|
||||
FROM ai_messages WHERE conversation_id = ?1 AND seq < ?2 ORDER BY seq DESC LIMIT ?3"
|
||||
} else {
|
||||
"SELECT id, conversation_id, seq, role, content, parts, tool_call_id,
|
||||
tool_calls, model, status, reasoning_content, timestamp, created_at,
|
||||
prompt_tokens, completion_tokens,
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens, is_estimated
|
||||
FROM ai_messages WHERE conversation_id = ?1 ORDER BY seq DESC LIMIT ?2"
|
||||
};
|
||||
let mut stmt = guard.prepare(sql).map_err(storage_err)?;
|
||||
@@ -289,8 +291,8 @@ impl AiMessageRepo {
|
||||
(id, conversation_id, seq, role, content, parts, tool_call_id,
|
||||
tool_calls, model, status, reasoning_content, timestamp, created_at,
|
||||
prompt_tokens, completion_tokens,
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
|
||||
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens, is_estimated)
|
||||
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19)",
|
||||
)
|
||||
.map_err(storage_err)?;
|
||||
for rec in &records {
|
||||
@@ -299,7 +301,7 @@ impl AiMessageRepo {
|
||||
rec.parts, rec.tool_call_id, rec.tool_calls, rec.model, rec.status,
|
||||
rec.reasoning_content, rec.timestamp, rec.created_at,
|
||||
rec.prompt_tokens, rec.completion_tokens,
|
||||
rec.prompt_cache_hit_tokens, rec.prompt_cache_miss_tokens, rec.reasoning_tokens
|
||||
rec.prompt_cache_hit_tokens, rec.prompt_cache_miss_tokens, rec.reasoning_tokens, rec.is_estimated
|
||||
])
|
||||
.map_err(storage_err)?;
|
||||
}
|
||||
@@ -373,6 +375,7 @@ mod tests {
|
||||
prompt_cache_hit_tokens: None,
|
||||
prompt_cache_miss_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
is_estimated: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -535,6 +538,7 @@ mod tests {
|
||||
prompt_cache_hit_tokens: None,
|
||||
prompt_cache_miss_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
is_estimated: None,
|
||||
},
|
||||
AiMessageRecord {
|
||||
id: "new_1".into(),
|
||||
@@ -555,6 +559,7 @@ mod tests {
|
||||
prompt_cache_hit_tokens: None,
|
||||
prompt_cache_miss_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
is_estimated: None,
|
||||
},
|
||||
];
|
||||
repo.replace_conversation("conv", records).await.expect("replace");
|
||||
@@ -627,6 +632,7 @@ mod tests {
|
||||
prompt_cache_hit_tokens: None,
|
||||
prompt_cache_miss_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
is_estimated: None,
|
||||
}],
|
||||
)
|
||||
.await
|
||||
@@ -668,6 +674,7 @@ mod tests {
|
||||
prompt_cache_hit_tokens: None,
|
||||
prompt_cache_miss_tokens: None,
|
||||
reasoning_tokens: None,
|
||||
is_estimated: None,
|
||||
};
|
||||
repo.replace_conversation("c", vec![rec()]).await.expect("1st");
|
||||
repo.replace_conversation("c", vec![rec()]).await.expect("2nd");
|
||||
|
||||
@@ -46,7 +46,8 @@ pub fn run(conn: &Connection) -> Result<()> {
|
||||
// V33 = 审批重启恢复:ai_conversations 加 pending_approvals TEXT 列,持久化挂起审批快照,
|
||||
// 重启后从 DB 恢复 pending_approvals 内存态,使待审批不丢。
|
||||
// V41 = 任务关联工程模块:tasks.module_id 列(工程系统打底,项目多工程下任务落到具体 module)。
|
||||
let steps: [(i32, fn(&Connection) -> Result<()>); 41] = [
|
||||
// V42 = ai_messages 加 is_estimated 列(消息级估算标记,reload 逐条回显「估算」标注)。
|
||||
let steps: [(i32, fn(&Connection) -> Result<()>); 42] = [
|
||||
(1, migrate_v1),
|
||||
(2, migrate_v2),
|
||||
(3, migrate_v3),
|
||||
@@ -88,6 +89,7 @@ pub fn run(conn: &Connection) -> Result<()> {
|
||||
(39, migrate_v39),
|
||||
(40, migrate_v40),
|
||||
(41, migrate_v41),
|
||||
(42, migrate_v42),
|
||||
];
|
||||
|
||||
for (version, migrate_fn) in steps {
|
||||
@@ -1257,6 +1259,22 @@ fn migrate_v41(conn: &Connection) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// V42: ai_messages 加 is_estimated 列(消息级估算标记)
|
||||
///
|
||||
/// 消息级 token 持久化(V38/V39)已落 prompt/completion/cache/reasoning,但「该轮 prompt 是否
|
||||
/// estimated 兜底」未存——reload 逐条回显时无法对齐 live 态 AiCompleted.is_estimated 标注。
|
||||
/// 本迁移补 INTEGER 列(NULL=老消息未标记,向前兼容;0=false 真实,1=true 估算)。
|
||||
/// 用 PRAGMA 探测列存在性,缺失才 ALTER(同 v38/v39 模式),对新库/老库均安全。
|
||||
fn migrate_v42(conn: &Connection) -> Result<()> {
|
||||
if !column_exists(conn, "ai_messages", "is_estimated") {
|
||||
conn.execute("ALTER TABLE ai_messages ADD COLUMN is_estimated INTEGER", [])?;
|
||||
tracing::info!("v42: ai_messages 加 is_estimated 列(消息级估算标记)");
|
||||
}
|
||||
conn.execute("INSERT INTO schema_version (version) VALUES (?)", [42])?;
|
||||
tracing::info!("迁移 v42 完成: ai_messages 加 is_estimated 列");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// V21 建表 SQL — 消息拆分存储 ai_messages 表
|
||||
///
|
||||
/// 与 V9_SQL 中的 ai_messages 镜像(V9 给新库,此 const 给老库 V21 迁移用 IF NOT EXISTS)。
|
||||
@@ -1281,6 +1299,7 @@ CREATE TABLE IF NOT EXISTS ai_messages (
|
||||
prompt_cache_hit_tokens INTEGER,
|
||||
prompt_cache_miss_tokens INTEGER,
|
||||
reasoning_tokens INTEGER,
|
||||
is_estimated INTEGER,
|
||||
UNIQUE(conversation_id, seq)
|
||||
);
|
||||
|
||||
@@ -1532,6 +1551,7 @@ CREATE TABLE IF NOT EXISTS ai_messages (
|
||||
prompt_cache_hit_tokens INTEGER,
|
||||
prompt_cache_miss_tokens INTEGER,
|
||||
reasoning_tokens INTEGER,
|
||||
is_estimated INTEGER,
|
||||
UNIQUE(conversation_id, seq)
|
||||
);
|
||||
|
||||
@@ -2023,7 +2043,7 @@ mod tests {
|
||||
"tasks.module_id 列缺失(V41 加)"
|
||||
);
|
||||
|
||||
// 3. schema_version 应推进到 41(全量迁移成功落版本号)
|
||||
// 3. schema_version 应推进到 42(全量迁移成功落版本号)
|
||||
let max_version: i64 = conn
|
||||
.query_row(
|
||||
"SELECT COALESCE(MAX(version), 0) FROM schema_version",
|
||||
@@ -2032,8 +2052,8 @@ mod tests {
|
||||
)
|
||||
.expect("查 schema_version 应成功");
|
||||
assert_eq!(
|
||||
max_version, 41,
|
||||
"全量迁移后 schema_version 应为 41(实际 {}),说明某条 migrate_vN 链路断在中间",
|
||||
max_version, 42,
|
||||
"全量迁移后 schema_version 应为 42(实际 {}),说明某条 migrate_vN 链路断在中间",
|
||||
max_version
|
||||
);
|
||||
|
||||
@@ -2046,6 +2066,11 @@ mod tests {
|
||||
column_exists(&conn, "project_modules", "status"),
|
||||
"project_modules.status 列缺失(V40 加)"
|
||||
);
|
||||
// 5. V42 抽查:ai_messages.is_estimated 列存在(消息级估算标记)
|
||||
assert!(
|
||||
column_exists(&conn, "ai_messages", "is_estimated"),
|
||||
"ai_messages.is_estimated 列缺失(V42 加)"
|
||||
);
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
@@ -2105,7 +2130,7 @@ mod tests {
|
||||
|r| r.get(0),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(max_v, 41, "首轮应推进到 41");
|
||||
assert_eq!(max_v, 42, "首轮应推进到 42");
|
||||
|
||||
// 清空版本表强制全链第二遍(每步 execute 第二次)
|
||||
conn.execute("DELETE FROM schema_version", []).unwrap();
|
||||
@@ -2117,7 +2142,7 @@ mod tests {
|
||||
|r| r.get(0),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(max_v2, 41, "重跑后应重新推进到 41");
|
||||
assert_eq!(max_v2, 42, "重跑后应重新推进到 42");
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
|
||||
@@ -440,6 +440,9 @@ pub struct AiMessageRecord {
|
||||
pub prompt_cache_miss_tokens: Option<u32>,
|
||||
/// 思考 token(deepseek-reasoner/o1 reasoning_tokens,隐藏输出)。
|
||||
pub reasoning_tokens: Option<u32>,
|
||||
/// 本轮 prompt 是否估算值(round_usage.prompt_tokens==0 → estimated_prompt 兜底打标)。
|
||||
/// reload 逐条回显「估算」标注;老消息 NULL → None(向前兼容)。
|
||||
pub is_estimated: Option<bool>,
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
|
||||
Reference in New Issue
Block a user