Skip to content

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

Merged
ytallo merged 3 commits into
mainfrom
fix/session-manager-assistant-usage
Jul 24, 2026
Merged

(MOT-4184) fix(session-manager): persist assistant usage updates#602
ytallo merged 3 commits into
mainfrom
fix/session-manager-assistant-usage

Conversation

@ytallo

@ytallo ytallo commented Jul 24, 2026

Copy link
Copy Markdown
Contributor

Summary

  • persist terminal token and cost usage when session::update-message finalizes an assistant message
  • reject usage updates for message roles that cannot carry provider usage
  • update the function schema and add unit and BDD coverage

Root cause

The Harness included provider usage in its final session::update-message call, but Session Manager did not model that field in the request. Deserialization ignored it, so the call succeeded while the durable assistant message lost its token accounting. As a result, harness::metrics had no persisted usage to aggregate.

Validation

  • cargo test — 62 unit tests, 127 BDD scenarios, and 4 schema tests passed
  • replaced Session Manager in the local validation stack and ran harness::send with openai-codex / codex/gpt-5.6-luna
  • confirmed session::messages and harness::metrics both reported input, output, and reasoning tokens

Refs MOT-4184

Summary by CodeRabbit

  • New Features

    • Added support for recording token usage and cost details on completed assistant messages.
    • Message updates can now include input, output, reasoning, cache, and cost metrics.
    • Stored usage information is returned when retrieving updated messages.
  • Bug Fixes

    • Prevented usage details from being applied to non-assistant messages.
    • Improved validation of terminal message updates.

@vercel

vercel Bot commented Jul 24, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
workers Ready Ready Preview, Comment Jul 24, 2026 11:10pm
workers-tech-spec Building Building Preview, Comment Jul 24, 2026 11:10pm

Request Review

@coderabbitai

coderabbitai Bot commented Jul 24, 2026

Copy link
Copy Markdown

Review Change Stack

Warning

Review limit reached

@ytallo, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 50 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 21f45b1b-ead0-4c36-9480-a8379a6b4749

📥 Commits

Reviewing files that changed from the base of the PR and between 8e1f705 and 6d38c4c.

📒 Files selected for processing (2)
  • harness/tests/e2e/src/scenarios/exactly_once_function.rs
  • harness/tests/e2e/src/scenarios/multi_turn_traces.rs
📝 Walkthrough

Walkthrough

session::update-message now accepts optional terminal usage, persists it only on assistant messages, exposes the schema, and verifies the behavior through unit, feature, and end-to-end trace tests.

Changes

Terminal usage propagation

Layer / File(s) Summary
Usage request contract
session-manager/src/functions/update_message.rs, session-manager/tests/golden/schemas/session.update-message.json
The update request and schema define optional terminal usage fields for token, cache, reasoning, and cost data.
Assistant usage persistence
session-manager/src/types.rs, session-manager/src/service.rs
Usage is applied only to assistant messages; unsupported roles return InvalidEntryKind, with unit coverage for both cases.
Usage-aware trace expectations
harness/tests/e2e/src/scenarios/*, session-manager/tests/features/update_message.feature
DSL helpers, exact traces, and feature scenarios include and verify assistant usage values.
Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant UpdateMessageRequest
  participant SessionService
  participant AgentMessage
  UpdateMessageRequest->>SessionService: provide optional usage
  SessionService->>AgentMessage: set_usage(new_usage)
  AgentMessage-->>SessionService: return role support
  SessionService-->>UpdateMessageRequest: persist or return InvalidEntryKind
Loading

Suggested reviewers: sergiofilhowz

Poem

A rabbit counts tokens, one, two, three,
Assistant usage joins the decree.
Calls carry numbers, traces hop,
Stored messages show the final crop.
Flopsy approves with a joyful stop!

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly matches the main change: persisting assistant usage updates in session-manager.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/session-manager-assistant-usage

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@github-actions

Copy link
Copy Markdown
Contributor

skill-check — worker

0 verified, 49 skipped (no docs/).

Layer Result
structure
vale
ai
render

Four for four. Nicely done.

@ytallo
ytallo marked this pull request as ready for review July 24, 2026 22:41

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
session-manager/src/service.rs (1)

692-700: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add an integration test for usage rejection.

Line 692 adds the public non-assistant validation path, but current tests only call set_usage directly. Add a session::update-message case that supplies usage to a user message and asserts the error plus unchanged persisted content.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@session-manager/src/service.rs` around lines 692 - 700, Add an integration
test for the session::update-message flow that submits usage for a user message,
verifies the request returns the expected invalid-entry-kind error, and confirms
the message’s persisted content remains unchanged. Exercise the public update
path rather than calling set_usage directly.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@session-manager/src/service.rs`:
- Around line 692-700: Add an integration test for the session::update-message
flow that submits usage for a user message, verifies the request returns the
expected invalid-entry-kind error, and confirms the message’s persisted content
remains unchanged. Exercise the public update path rather than calling set_usage
directly.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: cccd6274-63a4-49e2-abf7-1cb4018d0864

📥 Commits

Reviewing files that changed from the base of the PR and between 83bc78b and 8e1f705.

📒 Files selected for processing (8)
  • harness/tests/e2e/src/scenarios/dsl.rs
  • harness/tests/e2e/src/scenarios/exactly_once_function.rs
  • harness/tests/e2e/src/scenarios/multi_turn_traces.rs
  • session-manager/src/functions/update_message.rs
  • session-manager/src/service.rs
  • session-manager/src/types.rs
  • session-manager/tests/features/update_message.feature
  • session-manager/tests/golden/schemas/session.update-message.json

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant