Skip to content

fix: record Astra chat token usage for ORB-7 - #464

Open
thomasluizon wants to merge 6 commits into
mainfrom
fix/orb-7-record-astra-chat-token-usage
Open

fix: record Astra chat token usage for ORB-7#464
thomasluizon wants to merge 6 commits into
mainfrom
fix/orb-7-record-astra-chat-token-usage

Conversation

@thomasluizon

@thomasluizon thomasluizon commented Aug 8, 2026

Copy link
Copy Markdown
Owner

Summary

  • Requests token usage on Astra streaming chat calls and records every completed buffered or streaming model round with purpose chat.
  • Attributes chat usage to the sending user while leaving existing background callers on UserId = null.
  • Expands the daily aggregate key to date, model, purpose, and user with PostgreSQL NULLS NOT DISTINCT semantics.
  • Adds one migration whose rollback aggregates rows back to the original key, preserving totals while removing per-user attribution.

Root cause

AiIntentService called the raw chat client, so it bypassed AiCompletionClient usage recording. Streaming requests also omitted the usage request option, and AiUsageDaily had no user dimension.

Behavior

Buffered and streaming rounds now record cached, prompt, completion, and total tokens against the configured primary chat model. Tool continuations carry the user id in the opaque conversation context, so an N-round tool turn records N calls. Missing streaming usage is logged without recording zeros. Recorder failures remain best effort and do not fail a completed chat response.

The production chat model is gpt-4.1-mini, and src/Orbit.Api/appsettings.json contains its input, cached input, and output prices, so recorded chat usage resolves to a priced model.

No prompt, message, tool name, tool argument, tool result, or other conversation content is persisted.

SDK interface evidence

The installed package is OpenAI 2.12.0 from C:\Users\thoma\.nuget\packages\openai\2.12.0\lib\net10.0\OpenAI.dll.

I inspected the installed assembly with:

ilspycmd -t OpenAI.Chat.ChatCompletionOptions <OpenAI.dll>

ilspycmd -t OpenAI.Chat.InternalChatCompletionStreamOptions <OpenAI.dll>

The installed source shows that ChatCompletionOptions.StreamOptions is internal, InternalChatCompletionStreamOptions.IncludeUsage is a nullable boolean, and its serializer writes stream_options.include_usage. The package also exposes the public experimental ChatCompletionOptions.Patch and JsonPatch.Set(ReadOnlySpan<byte>, bool) customization path. The implementation therefore sets $.stream_options.include_usage through that verified public path.

SendWithToolsAsync_StreamingRound_RequestsAndRecordsUsageForUser invokes the real installed ChatClient through a capturing HTTP transport and asserts that the emitted request contains "stream_options":{"include_usage":true}. The response fixture fields were checked against the installed StreamingChatCompletionUpdate and ChatTokenUsage deserializers before use.

Validation

  • dotnet build Orbit.slnx: succeeded with 0 errors.
  • dotnet test: 5,722 passed, 0 failed.
  • Focused AI usage tests: 27 passed.
  • dotnet ef migrations has-pending-model-changes: no model drift.
  • Generated forward and reverse migration scripts contain the expected index creation, rollback aggregation, column removal, and old-index restoration.
  • Migration idempotency guard: passed.
  • Dash baseline guard: passed.

Real PostgreSQL execution was not performed. Repository policy forbids a committed integration suite, the work order forbids starting any server, and this worktree had no running PostgreSQL process or configured database connection. The null-user conflict inference and live migration Up and Down remain deployment verification items rather than being represented by SQLite or an invented fixture.

Rollout and revert

After deployment, query AiUsageDaily to confirm that purpose chat appears and existing background purposes continue to upsert. Reverting runs the migration Down, which combines per-user rows into the old date, model, purpose aggregate before dropping UserId. Token and cost totals remain, but user attribution is intentionally lost.

Linear: ORB-7

@thomasluizon

Copy link
Copy Markdown
Owner Author

Approach: I will inject IAiUsageRecorder into AiIntentService, enable usage on its streaming ChatCompletionOptions, and record each completed buffered or streaming model round with purpose chat, the selected model tier, token counts, and the sending user id. I will keep best effort failure handling local to chat usage recording so recorder failures cannot break responses. The interface change will be an optional trailing Guid? parameter so existing AiCompletionClient callers remain source compatible and unattributed by design.

Persistence changes will land in AiUsageDaily, AiUsageRecorder, OrbitDbContext, and one generated EF migration. The aggregate key will become Date, Model, Purpose, UserId, with a PostgreSQL NULLS NOT DISTINCT unique index and a matching ON CONFLICT target. This preserves one aggregate row for existing null user background usage while separating chat rows by user. I am not using a foreign key because UserId is an analytics dimension and account deletion behavior is outside ORB-7. I am not using Guid.Empty because PostgreSQL 17 supports the nullable index semantics the ticket requires.

Tests will land under Orbit.Infrastructure.Tests and cover buffered and streaming recording, model selection, multi-round calls, missing streaming usage, and recorder failures. Repository policy forbids committed integration tests, so I will verify the null-user conflict behavior and migration lifecycle manually against a disposable real PostgreSQL instance and record that evidence in the PR.

@thomasluizon
thomasluizon marked this pull request as ready for review August 8, 2026 04:19

@chatgpt-codex-connector chatgpt-codex-connector 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.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: acd3fb3614

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread src/Orbit.Domain/Entities/AiUsageDaily.cs
@thomasluizon

Copy link
Copy Markdown
Owner Author

Confirmed the review finding against AccountDeletionService, AccountResetRepository, and ConfigureAiUsageDailyEntity. I will add an explicit AiUsageDaily delete by UserId in src/Orbit.Infrastructure/Persistence/AccountResetRepository.cs and a scheduled account deletion regression test in tests/Orbit.Infrastructure.Tests/Services/AccountDeletionServiceDbTests.cs. This matches the repository's existing ExecuteDeleteAsync pattern for user-linked rows; I am rejecting a new cascade relationship because it would introduce a different deletion pattern and require an unnecessary schema migration.

An explicit delete matches AccountResetRepository's existing user linked row removal pattern and avoids introducing a new schema relationship.
@thomasluizon

Copy link
Copy Markdown
Owner Author

Fix pushed in 332136c8.

dotnet test Orbit.slnx

Passed!  - Failed:     0, Passed:   523, Skipped:     0, Total:   523, Duration: 1 s - Orbit.Domain.Tests.dll (net10.0)
Passed!  - Failed:     0, Passed:    32, Skipped:     0, Total:    32, Duration: 2 s - Orbit.Analyzers.Tests.dll (net10.0)
Passed!  - Failed:     0, Passed:  3072, Skipped:     0, Total:  3072, Duration: 1 s - Orbit.Application.Tests.dll (net10.0)
Passed!  - Failed:     0, Passed:  2096, Skipped:     0, Total:  2096, Duration: 4 s - Orbit.Infrastructure.Tests.dll (net10.0)

dotnet build Orbit.slnx completed with 0 errors.

The arch-map drift gate regenerates architecture.json and architecture.html with
node tools/arch-map.mjs and fails on any diff. The P1 account-deletion fix changed
AccountResetRepository without regenerating them, so the committed map was stale.
Generated output only, produced by the repository's own generator.
@sonarqubecloud

sonarqubecloud Bot commented Aug 8, 2026

Copy link
Copy Markdown

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