From 994c36d4a0cd364be365831d9ad642f14e456614 Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Fri, 24 Jul 2026 19:38:52 -0300 Subject: [PATCH 1/3] (MOT-4184) fix(session-manager): persist assistant usage updates --- .../src/functions/update_message.rs | 5 +- session-manager/src/service.rs | 9 +++ session-manager/src/types.rs | 39 +++++++++++ .../tests/features/update_message.feature | 20 +++++- .../schemas/session.update-message.json | 64 +++++++++++++++++++ 5 files changed, 135 insertions(+), 2 deletions(-) diff --git a/session-manager/src/functions/update_message.rs b/session-manager/src/functions/update_message.rs index 984b6a7a3..3d5b60a73 100644 --- a/session-manager/src/functions/update_message.rs +++ b/session-manager/src/functions/update_message.rs @@ -7,7 +7,7 @@ use serde_json::Value; use super::Deps; use crate::error::SessionError; -use crate::types::{ContentBlock, JsonMap}; +use crate::types::{ContentBlock, JsonMap, Usage}; #[derive(Debug, Clone, Deserialize, JsonSchema)] pub struct UpdateMessageRequest { @@ -17,6 +17,9 @@ pub struct UpdateMessageRequest { pub content: Vec, /// New `details` — only for `function_result` / `custom` messages. pub details: Option, + /// Final token/cost accounting — only for `assistant` messages. Omit for + /// streaming content updates that do not yet have terminal usage. + pub usage: Option, /// Optimistic concurrency: when supplied and it does not match the /// entry's current revision, nothing is written and /// `{ updated: false, revision }` returns the current revision. diff --git a/session-manager/src/service.rs b/session-manager/src/service.rs index d1c94b83c..a099dae01 100644 --- a/session-manager/src/service.rs +++ b/session-manager/src/service.rs @@ -689,6 +689,15 @@ impl SessionService { } message.set_content(req.content); + if let Some(new_usage) = req.usage { + if !message.set_usage(new_usage) { + return Err(SessionError::InvalidEntryKind(format!( + "entry {} has role {:?}; `usage` applies only to assistant messages", + req.entry_id, + message.role() + ))); + } + } if let Some(new_details) = req.details { match &mut message { AgentMessage::FunctionResult { details, .. } => *details = new_details, diff --git a/session-manager/src/types.rs b/session-manager/src/types.rs index 6decb1e61..e1d38728c 100644 --- a/session-manager/src/types.rs +++ b/session-manager/src/types.rs @@ -191,6 +191,16 @@ impl AgentMessage { | AgentMessage::Custom { content, .. } => *content = new_content, } } + + /// Set terminal usage on an assistant message. Returns false for roles + /// that cannot carry model usage. + pub fn set_usage(&mut self, new_usage: Usage) -> bool { + let AgentMessage::Assistant { usage, .. } = self else { + return false; + }; + *usage = Some(new_usage); + true + } } /// Bookkeeping payload of a `kind: "custom"` session entry. @@ -379,6 +389,35 @@ mod tests { assert_eq!(round[2]["details"], json!({ "x": 1 })); } + #[test] + fn usage_updates_only_assistant_messages() { + let mut assistant: AgentMessage = serde_json::from_value(json!({ + "role": "assistant", + "content": [], + "stop_reason": "end", + "model": "m1", + "provider": "p1", + "timestamp": 2 + })) + .unwrap(); + assert!(assistant.set_usage(Usage { + input: Some(10), + output: Some(4), + ..Usage::default() + })); + let value = serde_json::to_value(assistant).unwrap(); + assert_eq!(value["usage"]["input"], 10); + assert_eq!(value["usage"]["output"], 4); + + let mut user: AgentMessage = serde_json::from_value(json!({ + "role": "user", + "content": [], + "timestamp": 1 + })) + .unwrap(); + assert!(!user.set_usage(Usage::default())); + } + #[test] fn session_entry_kind_discrimination() { let entry: SessionEntry = serde_json::from_value(json!({ diff --git a/session-manager/tests/features/update_message.feature b/session-manager/tests/features/update_message.feature index 3c6110857..76b587435 100644 --- a/session-manager/tests/features/update_message.feature +++ b/session-manager/tests/features/update_message.feature @@ -2,7 +2,7 @@ Feature: session::update-message — streaming deltas and edited output Contract (session-manager.md § session::update-message): replaces the - content (and optionally details) of an existing message entry. Each + content (and optionally assistant usage or result details) of an existing message entry. Each successful update increments the entry's revision (echoed on the event, monotonic per entry — consumers keep the highest). With expected_revision set, a mismatch writes nothing and returns @@ -46,6 +46,24 @@ Feature: session::update-message — streaming deltas and edited output Then the response field "entry.revision" is 2 And the response field "entry.message.content.0.text" is "Hello world" + # Prevents: terminal provider usage being discarded when the harness + # replaces the final content of its streamed assistant placeholder. + Scenario: terminal usage is persisted on assistant messages + When I call "session::update-message" with: + """ + { "session_id": "s_001", "entry_id": "e_001", + "content": [{ "type": "text", "text": "done" }], + "usage": { "input": 120, "output": 18, "reasoning": 7 } } + """ + Then the response field "updated" is true + When I call "session::get-message" with: + """ + { "session_id": "s_001", "entry_id": "e_001" } + """ + Then the response field "entry.message.usage.input" is 120 + And the response field "entry.message.usage.output" is 18 + And the response field "entry.message.usage.reasoning" is 7 + # Prevents: updates rewriting history's position — the entry keeps its # creation timestamp (ordering anchor) while the event carries the # update time. diff --git a/session-manager/tests/golden/schemas/session.update-message.json b/session-manager/tests/golden/schemas/session.update-message.json index 8b3873ad3..f89bf24da 100644 --- a/session-manager/tests/golden/schemas/session.update-message.json +++ b/session-manager/tests/golden/schemas/session.update-message.json @@ -132,6 +132,59 @@ "type": "object" } ] + }, + "Usage": { + "description": "Token / cost accounting reported by providers.", + "properties": { + "cache_read": { + "format": "uint64", + "minimum": 0.0, + "type": [ + "integer", + "null" + ] + }, + "cache_write": { + "format": "uint64", + "minimum": 0.0, + "type": [ + "integer", + "null" + ] + }, + "cost_usd": { + "format": "double", + "type": [ + "number", + "null" + ] + }, + "input": { + "format": "uint64", + "minimum": 0.0, + "type": [ + "integer", + "null" + ] + }, + "output": { + "format": "uint64", + "minimum": 0.0, + "type": [ + "integer", + "null" + ] + }, + "reasoning": { + "format": "uint64", + "minimum": 0.0, + "type": [ + "integer", + "null" + ] + } + }, + "type": "object" } }, "properties": { @@ -167,6 +220,17 @@ }, "session_id": { "type": "string" + }, + "usage": { + "anyOf": [ + { + "$ref": "#/definitions/Usage" + }, + { + "type": "null" + } + ], + "description": "Final token/cost accounting — only for `assistant` messages. Omit for streaming content updates that do not yet have terminal usage." } }, "required": [ From 8e1f70516f6a7e4dea5dad7d6698519dcc0c1b80 Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Fri, 24 Jul 2026 20:00:21 -0300 Subject: [PATCH 2/3] (MOT-4184) test(harness): include assistant usage in E2E fixtures --- harness/tests/e2e/src/scenarios/dsl.rs | 15 +++++++++++--- .../src/scenarios/exactly_once_function.rs | 9 ++++++++- .../e2e/src/scenarios/multi_turn_traces.rs | 20 ++++++++++++++++--- 3 files changed, 37 insertions(+), 7 deletions(-) diff --git a/harness/tests/e2e/src/scenarios/dsl.rs b/harness/tests/e2e/src/scenarios/dsl.rs index b49c132dc..c5017686e 100644 --- a/harness/tests/e2e/src/scenarios/dsl.rs +++ b/harness/tests/e2e/src/scenarios/dsl.rs @@ -743,6 +743,8 @@ impl Message { function: &ControlledFunction, arguments: Value, model: &ModelFixtureV1, + input_tokens: u64, + output_tokens: u64, ) -> Value { json!({ "role": "assistant", @@ -754,17 +756,24 @@ impl Message { }], "stop_reason": "end", "model": model.id, - "provider": model.provider + "provider": model.provider, + "usage": usage(input_tokens, output_tokens) }) } - pub(super) fn assistant_text(text: &str, model: &ModelFixtureV1) -> Value { + pub(super) fn assistant_text( + text: &str, + model: &ModelFixtureV1, + input_tokens: u64, + output_tokens: u64, + ) -> Value { json!({ "role": "assistant", "content": [{ "type": "text", "text": text }], "stop_reason": "end", "model": model.id, - "provider": model.provider + "provider": model.provider, + "usage": usage(input_tokens, output_tokens) }) } diff --git a/harness/tests/e2e/src/scenarios/exactly_once_function.rs b/harness/tests/e2e/src/scenarios/exactly_once_function.rs index 64ad62e8f..f413c2326 100644 --- a/harness/tests/e2e/src/scenarios/exactly_once_function.rs +++ b/harness/tests/e2e/src/scenarios/exactly_once_function.rs @@ -65,7 +65,14 @@ pub(super) fn scenario() -> ScenarioFixture { .system_prompt_sha256("{{system_prompt_sha256}}") .messages_exact([ Message::user(MESSAGE), - Message::function_call(CALL_ID, &record, arguments.clone(), &model), + Message::function_call( + CALL_ID, + &record, + arguments.clone(), + &model, + 8, + 4, + ), Message::function_result(CALL_ID, &record, "recorded"), ]) .tools_exact([record.tool()]), diff --git a/harness/tests/e2e/src/scenarios/multi_turn_traces.rs b/harness/tests/e2e/src/scenarios/multi_turn_traces.rs index 0f84fe8e4..59dc382f8 100644 --- a/harness/tests/e2e/src/scenarios/multi_turn_traces.rs +++ b/harness/tests/e2e/src/scenarios/multi_turn_traces.rs @@ -70,7 +70,14 @@ pub(super) fn scenario() -> ScenarioFixture { .system_prompt_sha256("{{system_prompt_sha256}}") .messages_exact([ Message::user(FIRST_MESSAGE), - Message::function_call(CALL_ID, &record, arguments.clone(), &model), + Message::function_call( + CALL_ID, + &record, + arguments.clone(), + &model, + 8, + 4, + ), Message::function_result(CALL_ID, &record, "recorded"), ]) .tools_exact([record.tool()]), @@ -85,9 +92,16 @@ pub(super) fn scenario() -> ScenarioFixture { .system_prompt_regex("agent_trigger") .messages_exact([ Message::user(FIRST_MESSAGE), - Message::function_call(CALL_ID, &record, arguments.clone(), &model), + Message::function_call( + CALL_ID, + &record, + arguments.clone(), + &model, + 8, + 4, + ), Message::function_result(CALL_ID, &record, "recorded"), - Message::assistant_text(FIRST_TEXT, &model), + Message::assistant_text(FIRST_TEXT, &model, 18, 2), Message::user(SECOND_MESSAGE), ]) .tools_subset([Tool::named("agent_trigger")]), From 6d38c4ca5676a37ee3a1ea1d3e8b0da9e0e487b0 Mon Sep 17 00:00:00 2001 From: Ytallo Layon Date: Fri, 24 Jul 2026 20:09:33 -0300 Subject: [PATCH 3/3] (MOT-4184) style(harness): format usage fixture calls --- .../e2e/src/scenarios/exactly_once_function.rs | 9 +-------- .../e2e/src/scenarios/multi_turn_traces.rs | 18 ++---------------- 2 files changed, 3 insertions(+), 24 deletions(-) diff --git a/harness/tests/e2e/src/scenarios/exactly_once_function.rs b/harness/tests/e2e/src/scenarios/exactly_once_function.rs index f413c2326..7a3274923 100644 --- a/harness/tests/e2e/src/scenarios/exactly_once_function.rs +++ b/harness/tests/e2e/src/scenarios/exactly_once_function.rs @@ -65,14 +65,7 @@ pub(super) fn scenario() -> ScenarioFixture { .system_prompt_sha256("{{system_prompt_sha256}}") .messages_exact([ Message::user(MESSAGE), - Message::function_call( - CALL_ID, - &record, - arguments.clone(), - &model, - 8, - 4, - ), + Message::function_call(CALL_ID, &record, arguments.clone(), &model, 8, 4), Message::function_result(CALL_ID, &record, "recorded"), ]) .tools_exact([record.tool()]), diff --git a/harness/tests/e2e/src/scenarios/multi_turn_traces.rs b/harness/tests/e2e/src/scenarios/multi_turn_traces.rs index 59dc382f8..19be1d10f 100644 --- a/harness/tests/e2e/src/scenarios/multi_turn_traces.rs +++ b/harness/tests/e2e/src/scenarios/multi_turn_traces.rs @@ -70,14 +70,7 @@ pub(super) fn scenario() -> ScenarioFixture { .system_prompt_sha256("{{system_prompt_sha256}}") .messages_exact([ Message::user(FIRST_MESSAGE), - Message::function_call( - CALL_ID, - &record, - arguments.clone(), - &model, - 8, - 4, - ), + Message::function_call(CALL_ID, &record, arguments.clone(), &model, 8, 4), Message::function_result(CALL_ID, &record, "recorded"), ]) .tools_exact([record.tool()]), @@ -92,14 +85,7 @@ pub(super) fn scenario() -> ScenarioFixture { .system_prompt_regex("agent_trigger") .messages_exact([ Message::user(FIRST_MESSAGE), - Message::function_call( - CALL_ID, - &record, - arguments.clone(), - &model, - 8, - 4, - ), + Message::function_call(CALL_ID, &record, arguments.clone(), &model, 8, 4), Message::function_result(CALL_ID, &record, "recorded"), Message::assistant_text(FIRST_TEXT, &model, 18, 2), Message::user(SECOND_MESSAGE),