fix: keep token_usage DB-resident in hybrid mode instead of offloading to object storage - #4732
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe PR keeps ChangesToken usage retention and offload behavior
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Comment |
|
Warning This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
This stack of pull requests is managed by Graphite. Learn more about stacking. |
Confidence Score: 4/5Safe to merge; the fix correctly targets a well-scoped storage routing bug with matching tests for both new and legacy rows. The core logic — removing token_usage from payloadFields/ClearPayload and adding a denormalized-column fallback — is correct and well-tested. Three minor issues prevent a perfect score: the ExtractPayload map capacity hint is off-by-one after the change; token_usage is always written into the snapshot even when empty (unlike the metadata guard pattern); and MergePayloadFromJSON restores token_usage from the snapshot on FindByID while metadata is intentionally left DB-authoritative, creating a design asymmetry that could produce stale data if token counts were ever amended after upload. framework/logstore/payload.go — map capacity hint, unconditional empty token_usage in snapshot, and asymmetric snapshot-restore behaviour in MergePayloadFromJSON. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant App
participant Hybrid
participant DB
participant ObjStore
App->>Hybrid: CreateIfNotExists(log)
Hybrid->>Hybrid: SerializeFields() sets TokenUsage JSON
Hybrid->>Hybrid: ExtractPayload() snapshot includes token_usage
Hybrid->>Hybrid: ClearPayload() clears large fields (NOT token_usage)
Hybrid->>DB: INSERT (token_usage JSON stays in row)
Hybrid->>ObjStore: Upload snapshot (token_usage also in snapshot)
App->>Hybrid: SearchLogs()
Hybrid->>DB: SELECT rows (token_usage JSON present)
DB-->>Hybrid: rows with TokenUsage
Hybrid->>Hybrid: DeserializeFields() parses TokenUsageParsed
Note over Hybrid: Fallback: if TokenUsage empty rebuild from denormalized cols
Hybrid-->>App: Logs with TotalTokens populated
App->>Hybrid: FindByID(id)
Hybrid->>DB: SELECT row
Hybrid->>ObjStore: Download snapshot
Hybrid->>Hybrid: MergePayloadFromJSON() restores large fields + token_usage
Hybrid-->>App: Full log with all fields
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant App
participant Hybrid
participant DB
participant ObjStore
App->>Hybrid: CreateIfNotExists(log)
Hybrid->>Hybrid: SerializeFields() sets TokenUsage JSON
Hybrid->>Hybrid: ExtractPayload() snapshot includes token_usage
Hybrid->>Hybrid: ClearPayload() clears large fields (NOT token_usage)
Hybrid->>DB: INSERT (token_usage JSON stays in row)
Hybrid->>ObjStore: Upload snapshot (token_usage also in snapshot)
App->>Hybrid: SearchLogs()
Hybrid->>DB: SELECT rows (token_usage JSON present)
DB-->>Hybrid: rows with TokenUsage
Hybrid->>Hybrid: DeserializeFields() parses TokenUsageParsed
Note over Hybrid: Fallback: if TokenUsage empty rebuild from denormalized cols
Hybrid-->>App: Logs with TotalTokens populated
App->>Hybrid: FindByID(id)
Hybrid->>DB: SELECT row
Hybrid->>ObjStore: Download snapshot
Hybrid->>Hybrid: MergePayloadFromJSON() restores large fields + token_usage
Hybrid-->>App: Full log with all fields
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
framework/logstore/payload.go (1)
59-59: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueBump the map capacity hint to match actual entries.
token_usageis no longer inpayloadFieldsbut is always written, andmetadatamay also be added — so the map can hold up tolen(payloadFields)+2entries while the hint is+1. This is a harmless minor under-allocation (the test already expects+2), but aligning the hint avoids an extra growth.♻️ Capacity hint tweak
- m := make(map[string]string, len(payloadFields)+1) + m := make(map[string]string, len(payloadFields)+2)🤖 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 `@framework/logstore/payload.go` at line 59, The map capacity hint in payload construction is too low for the actual number of entries being inserted. Update the map allocation in the payload builder around the map initialization so the hint matches the maximum written keys, including token_usage and the optional metadata field; use the payloadFields-based constructor logic to reflect up to len(payloadFields)+2 entries instead of len(payloadFields)+1.framework/logstore/tables.go (1)
730-739: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueFallback drops
CachedReadTokens.
SerializeFieldsalso denormalizesCachedReadTokens(fromPromptTokensDetails.CachedReadTokens), but this reconstruction only restores prompt/completion/total. For older offloaded rows that had cached-read tokens, the rebuiltTokenUsageParsedwill report them as zero. If cache analytics from list/search responses matter for those legacy rows, consider repopulatingPromptTokensDetailsfroml.CachedReadTokenswhen non-zero. Otherwise this is acceptable as a best-effort hydration.🤖 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 `@framework/logstore/tables.go` around lines 730 - 739, The fallback hydration in SerializeFields rebuilds TokenUsageParsed from PromptTokens, CompletionTokens, and TotalTokens only, so it loses CachedReadTokens for legacy offloaded rows. Update the reconstruction logic in tables.go to also repopulate PromptTokensDetails with CachedReadTokens when l.CachedReadTokens is non-zero, alongside the existing BifrostLLMUsage fields, so list/search responses preserve the denormalized cache-read data.
🤖 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 `@framework/logstore/payload.go`:
- Line 59: The map capacity hint in payload construction is too low for the
actual number of entries being inserted. Update the map allocation in the
payload builder around the map initialization so the hint matches the maximum
written keys, including token_usage and the optional metadata field; use the
payloadFields-based constructor logic to reflect up to len(payloadFields)+2
entries instead of len(payloadFields)+1.
In `@framework/logstore/tables.go`:
- Around line 730-739: The fallback hydration in SerializeFields rebuilds
TokenUsageParsed from PromptTokens, CompletionTokens, and TotalTokens only, so
it loses CachedReadTokens for legacy offloaded rows. Update the reconstruction
logic in tables.go to also repopulate PromptTokensDetails with CachedReadTokens
when l.CachedReadTokens is non-zero, alongside the existing BifrostLLMUsage
fields, so list/search responses preserve the denormalized cache-read data.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 77e6f32b-6d9b-40d8-bb28-fe8a4751bd2d
📒 Files selected for processing (5)
framework/logstore/hybrid_test.goframework/logstore/payload.goframework/logstore/payload_test.goframework/logstore/tables.gotransports/config.schema.json
1 participant
Summary
In hybrid mode,
token_usagewas being offloaded to object storage and cleared from the database row, causing token counts to be unavailable in log list views and breaking token-based sorting without fetching every object from storage. This PR keepstoken_usageDB-resident (likemetadata) while still including it in the object-store snapshot for downstream consumers.Changes
token_usagefrompayloadFieldsso it is no longer cleared from the DB row during hybrid offload.token_usageandTokenUsageParsedfromClearPayloadso the JSON column and parsed struct are preserved after upload.ExtractPayloadcontinues to writetoken_usageinto the object-store snapshot so object consumers still receive the full payload.DeserializeFieldsto reconstructTokenUsageParsedfrom the denormalizedprompt_tokens,completion_tokens, andtotal_tokenscolumns for older rows where the JSON column was previously cleared.object_storage_exclude_fieldsschema description to note thattoken_usageis always kept in the database and does not need to be listed explicitly.Type of change
Affected areas
How to test
go test ./framework/logstore/...The new
TestHybrid_TokenUsageStaysInDBtest verifies that after a log is written and uploaded in hybrid mode, thetoken_usageJSON column andTotalTokensremain in the DB row and are returned correctly bySearchLogs.TestDeserializeFields_TokenUsageFromDenormalizedColumnsverifies the fallback reconstruction path for older rows.Breaking changes
Security considerations
No security implications. This change only affects which columns are retained in the database versus offloaded to object storage.
Checklist
docs/contributing/README.mdand followed the guidelines