(MOT-4184) fix(session-manager): persist assistant usage updates

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)
         })
     }
 

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()]),

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")]),

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<ContentBlock>,
     /// New `details` — only for `function_result` / `custom` messages.
     pub details: Option<Value>,
+    /// Final token/cost accounting — only for `assistant` messages. Omit for
+    /// streaming content updates that do not yet have terminal usage.
+    pub usage: Option<Usage>,
     /// 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.

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,

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!({

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.

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": [