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..7a3274923 100644 --- a/harness/tests/e2e/src/scenarios/exactly_once_function.rs +++ b/harness/tests/e2e/src/scenarios/exactly_once_function.rs @@ -65,7 +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), + 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..19be1d10f 100644 --- a/harness/tests/e2e/src/scenarios/multi_turn_traces.rs +++ b/harness/tests/e2e/src/scenarios/multi_turn_traces.rs @@ -70,7 +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), + Message::function_call(CALL_ID, &record, arguments.clone(), &model, 8, 4), Message::function_result(CALL_ID, &record, "recorded"), ]) .tools_exact([record.tool()]), @@ -85,9 +85,9 @@ 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")]), 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": [