feat: add OpenAI batch pricing support - #4884
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR adds end-to-end batch-job accounting for provider batch results, including persistent lifecycle state, batch-specific pricing, usage extraction, aggregate logging, polling, governance reporting, provider adapters, and server startup wiring. ChangesBatch Accounting Feature
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 1 | ❌ 4❌ Failed checks (4 warnings)
✅ Passed checks (1 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
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 |
This stack of pull requests is managed by Graphite. Learn more about stacking. |
3d28874 to
d27f680
Compare
e4bb51b to
3ad12a9
Compare
![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.
Actionable comments posted: 5
🧹 Nitpick comments (8)
framework/batchaccounting/sweeper.go (1)
226-249: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueKV lease TTL can override early backoff.
acquireProviderPollLeasesets the lease forKVLeaseTTL(default 5m), butnextCheckAtreturns delays as low asInterval(default 1m) for early attempts. With a KV store present, a job rescheduled ~1m out stays due but its lease is still held, so it is repeatedly picked up byFindDueBatchJobsand silently skipped (Line 118) until the lease expires. This is functionally safe but wastes sweep/DB cycles; consider aligning the lease TTL with the reschedule delay (or releasing the lease on non-terminal completion).🤖 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/batchaccounting/sweeper.go` around lines 226 - 249, The provider poll lease in Sweeper.acquireProviderPollLease uses a fixed KVLeaseTTL that can outlast the early reschedule delay computed by Sweeper.nextCheckAt, causing the same due job to be repeatedly rediscovered and skipped. Update the lease timing so it cannot exceed the next scheduled poll window for a job, either by deriving the TTL from nextCheckAt/deterministicJitter or by releasing the lease when the job is re-queued for a later attempt. Keep the behavior aligned in acquireProviderPollLease, nextCheckAt, and any poll-scheduling path that uses FindDueBatchJobs.framework/logstore/tables.go (1)
273-317: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueSolid schema design; consider a defense-in-depth unique constraint.
idx_batch_jobs_identity(Provider, BatchID) is non-unique. Uniqueness is currently only guaranteed becauseUpsertBatchJobalways derivesIDviaBatchJobID(provider, batchID)whenIDis empty — if any future caller ever setsProvider/BatchIDwith a different explicitID, duplicate logical rows for the same batch could be created. Marking this indexuniquewould enforce the invariant at the DB layer rather than relying solely on caller discipline.🛡️ Optional: enforce identity uniqueness at the DB layer
- Provider string `gorm:"type:varchar(255);index:idx_batch_jobs_identity,priority:1;index:idx_batch_jobs_sweeper,priority:1;not null" json:"provider"` - BatchID string `gorm:"type:varchar(255);index:idx_batch_jobs_identity,priority:2;not null" json:"batch_id"` + Provider string `gorm:"type:varchar(255);uniqueIndex:idx_batch_jobs_identity,priority:1;index:idx_batch_jobs_sweeper,priority:1;not null" json:"provider"` + BatchID string `gorm:"type:varchar(255);uniqueIndex:idx_batch_jobs_identity,priority:2;not null" json:"batch_id"`🤖 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 273 - 317, The BatchJob identity index is not enforcing uniqueness at the database layer, so duplicate logical rows can still be created if a caller sets a different explicit ID. Update the BatchJob schema in tables.go to make the Provider + BatchID identity index unique, and keep the existing composite identity on BatchJob/TableName-backed persistence so the invariant is enforced by the DB rather than only by UpsertBatchJob/BatchJobID caller behavior.framework/logstore/batch_jobs_test.go (2)
1-1: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueFilename violates the repo's Go naming convention.
batch_jobs_test.gohas an extra underscore beyond the permitted_test.gosuffix. Should bebatchjobs_test.go.As per path instructions: "Go filenames: No underscores. The only permitted underscore is the
_test.gosuffix... Concatenate words (lowercase, no separators) for multi-word filenames." This is reinforced by a retrieved learning noting the underscore exception for non-test files "applies only to non-test .go files; the general rule may still apply to test files."🤖 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/batch_jobs_test.go` at line 1, The test file name violates the Go repository naming convention because it contains an extra underscore before the _test.go suffix. Rename the file to use concatenated lowercase words without separators, matching the repo rule for multi-word Go filenames, and keep the existing package and test contents unchanged; the fix is only in the filename for batchjobs_test.go.Sources: Path instructions, Learnings
15-167: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLGTM! Test coverage correctly validates claim-token guarding, phase-marker token checks, and
NextCheckAtpreserve/clear semantics across terminal vs. non-terminal provider statuses.One optional gap:
MarkBatchJobUnpriceable/FailBatchJobAccounting(and the sharedfinishBatchJobAccountingclaim-token guard) aren't directly exercised against the real store here — only indirectly via thebatchaccountingpackage's fake store tests. Direct coverage would strengthen confidence in the DB-level guard for these paths too.🤖 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/batch_jobs_test.go` around lines 15 - 167, Add direct sqlite-store coverage for the claim-token guard on the accounting finish paths, since `TestBatchJobAccountingClaimTokenGuardsCompletion` and `TestBatchJobAccountingPhaseMarkersUseClaimToken` only cover completion and phase markers. Create/extend tests around `finishBatchJobAccounting`, `MarkBatchJobUnpriceable`, and `FailBatchJobAccounting` to verify wrong tokens return `ErrNotFound` and the correct token succeeds against the real `newSqliteLogStore`, not just the fake store tests.framework/logstore/rdb.go (1)
417-474: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueUpsert always does two round trips.
Every call performs a
Create(...OnConflict DoNothing)followed by an unconditionalUpdates(...), even on first insert where the just-inserted values already match. This is correct but costs an extra DB round trip per call; given batch-job upserts aren't a per-request hot path (they occur per batch lifecycle transition), this is likely an acceptable tradeoff.🤖 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/rdb.go` around lines 417 - 474, UpsertBatchJob currently always does a Create with OnConflict DoNothing followed by a separate Updates call, causing an unnecessary second DB round trip on insert. Update RDBLogStore.UpsertBatchJob to perform a single upsert path using the existing BatchJobID/ID conflict key and clause.OnConflict DoUpdates so the insert and field refresh happen in one write, while preserving the current nil/empty-field handling for BatchJob fields like Model, ProviderStatus, and NextCheckAt.framework/batchaccounting/accounting_test.go (1)
108-111: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
bifrost.Ptr()instead of the address operator for consistency.These spots take
&of local variables (&claimedBy,&expires,&token,&reason,&inputRate,&outputRate) to build pointer fields. The siblingcost_test.goin this same PR already follows the project convention of usingbifrost.Ptr(...)for this. Based on learnings, preferbifrost.Ptr()over&valueconsistently, including in test utilities.Also applies to: 168-168, 205-220
🤖 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/batchaccounting/accounting_test.go` around lines 108 - 111, Replace the local-variable address usage in the batch accounting tests with the project’s pointer helper for consistency. In the test setup around `entry.ClaimedBy`, `entry.ClaimExpiresAt`, `entry.ClaimToken`, and the other affected pointer fields in this file, use `bifrost.Ptr(...)` instead of taking `&` of locals. Keep the existing test behavior the same and align these helpers with the pattern already used in `cost_test.go`.Source: Learnings
plugins/logging/main.go (2)
659-665: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winSweeper is hardcoded to only sweep
schemas.OpenAIbatch jobs.
SweeperConfig.Provideris fixed toschemas.OpenAI.FindDueBatchJobs(and thus the sweeper) filters by this provider, so anyBatchJobcreated for a different provider (viarecordBatchJobLifecycle, which is provider-agnostic — it usesentry.Provider) would never be picked up: itsNextCheckAtwould stay set forever, and it would never transition out ofBatchJobAccountingStatusPending. Given thebatchaccountingpackage,BatchJobschema, andBatchResultFetcherinterface are all provider-generic (not OpenAI-specific), this looks like scope-limiting for the current PR ("add OpenAI batch pricing support") rather than an inherent constraint — but it's easy to forget to revisit when another provider's Batch API support ships, leaving those jobs silently stuck.Consider passing an empty
Provider(sweep all providers) or making this configurable, unless OpenAI-only scoping is a deliberate, documented decision for now.🤖 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 `@plugins/logging/main.go` around lines 659 - 665, The sweeper setup in main.go is incorrectly locked to schemas.OpenAI, which means FindDueBatchJobs will ignore batch jobs from other providers and leave them stuck in pending state. Update the SweeperConfig.Provider passed to batchaccounting.NewSweeper so the sweeper can cover all providers, or make the provider selection configurable if OpenAI-only behavior is intentional. Keep the fix aligned with the provider-agnostic flow used by recordBatchJobLifecycle, BatchJob, and BatchResultFetcher.
197-233: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winBatch results accounting runs synchronously inside
PostLLMHook, blocking the client response.
accountBatchResults→batchaccounting.AccountBatchResultsperforms multiple sequential store round-trips (upsert, claim, find, summarize every result item, create-if-not-exists, mark-written, report usage, mark-reported, complete) directly in thePostLLMHookcall path forBatchResultsRequest. OpenAI batches can contain up to tens of thousands of result lines, so this can add meaningful latency to the client'sBatchResultsRequestresponse — unlike the sweeper path, which does the same work off the request path.The plugin already has a precedent for deferring this kind of work off the hot path:
p.scheduleDeferredUsageUpdate+p.deferredUsageSembound concurrent deferred DB work elsewhere in this same file. Consider routingaccountBatchResultsthrough the same deferred/bounded-goroutine mechanism instead of calling it inline.Also applies to: 1378-1383
🤖 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 `@plugins/logging/main.go` around lines 197 - 233, Batch results accounting is still running inline from LoggerPlugin.PostLLMHook via accountBatchResults, which blocks the response path for BatchResultsRequest. Move the batchaccounting.AccountBatchResults call out of the synchronous path by scheduling it through the same deferred/bounded goroutine pattern used by p.scheduleDeferredUsageUpdate and p.deferredUsageSem, so the request can return immediately while accounting continues in the background. Keep accountBatchResults as the worker function, but have the hook enqueue it instead of invoking it 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.
Inline comments:
In `@framework/batchaccounting/accounting_test.go`:
- Around line 186-191: The fakeBatchPricing helper contains a broken leftover
method: CalculateCostForUsage references an undefined cost identifier and is not
needed by the PricingManager interface used in these tests. Remove the unused
CalculateCostForUsage method from fakeBatchPricing so the test package compiles
cleanly and only the required CalculateBatchCostForUsage behavior remains.
In `@framework/batchaccounting/accounting.go`:
- Around line 408-461: Remove the stale duplicate extractUsage switch
implementation and keep only the map-based extractor path using usageExtractors.
Fix the compile break by ensuring extractUsage returns extractedUsage, error in
all paths, and replace the invalid default return with the supported no-op
behavior used elsewhere. Also add the missing extractor functions referenced by
usageExtractors, namely extractResponseBodyUsage and extractAnthropicUsage, so
they encapsulate the response parsing logic currently embedded in the old
extractUsage block.
In `@framework/batchaccounting/sweeper.go`:
- Around line 121-131: `Run` in the sweeper can nil-deref if
`fetcher.RetrieveBatch` returns a nil batch with no error; add an explicit
`retrieved == nil` check immediately after the retrieve call and handle it the
same way as other empty responses by warning, rescheduling with `reschedule`,
and returning before calling `batchJobFromRetrieve`. Keep the fix localized
around `RetrieveBatch`, `batchJobFromRetrieve`, and the existing `results ==
nil` pattern below so the empty-response behavior stays consistent.
In `@plugins/governance/main.go`:
- Around line 1738-1752: ReportBatchUsage stops on the first BumpBudgetUsage or
BumpRateLimitUsageBy error, which can skip remaining IDs and cause undercounting
for the batch. Update ReportBatchUsage to continue iterating through all
BudgetIDs and RateLimitIDs, collect any errors from p.store.BumpBudgetUsage and
p.store.BumpRateLimitUsageBy, and return the combined error after both loops
finish so one bad ID does not block unrelated updates.
- Around line 1747-1751: The batch usage path in the rate-limit update logic is
always charging a single request via requestDelta in the usage bump loop, which
undercounts batch jobs. Update the BatchUsageReport flow and the call site that
uses p.store.BumpRateLimitUsageBy so it threads through
schemas.BatchRequestCounts from the batch pipeline and uses that total instead
of a hardcoded 1 when request-based limits should reflect item volume.
---
Nitpick comments:
In `@framework/batchaccounting/accounting_test.go`:
- Around line 108-111: Replace the local-variable address usage in the batch
accounting tests with the project’s pointer helper for consistency. In the test
setup around `entry.ClaimedBy`, `entry.ClaimExpiresAt`, `entry.ClaimToken`, and
the other affected pointer fields in this file, use `bifrost.Ptr(...)` instead
of taking `&` of locals. Keep the existing test behavior the same and align
these helpers with the pattern already used in `cost_test.go`.
In `@framework/batchaccounting/sweeper.go`:
- Around line 226-249: The provider poll lease in
Sweeper.acquireProviderPollLease uses a fixed KVLeaseTTL that can outlast the
early reschedule delay computed by Sweeper.nextCheckAt, causing the same due job
to be repeatedly rediscovered and skipped. Update the lease timing so it cannot
exceed the next scheduled poll window for a job, either by deriving the TTL from
nextCheckAt/deterministicJitter or by releasing the lease when the job is
re-queued for a later attempt. Keep the behavior aligned in
acquireProviderPollLease, nextCheckAt, and any poll-scheduling path that uses
FindDueBatchJobs.
In `@framework/logstore/batch_jobs_test.go`:
- Line 1: The test file name violates the Go repository naming convention
because it contains an extra underscore before the _test.go suffix. Rename the
file to use concatenated lowercase words without separators, matching the repo
rule for multi-word Go filenames, and keep the existing package and test
contents unchanged; the fix is only in the filename for batchjobs_test.go.
- Around line 15-167: Add direct sqlite-store coverage for the claim-token guard
on the accounting finish paths, since
`TestBatchJobAccountingClaimTokenGuardsCompletion` and
`TestBatchJobAccountingPhaseMarkersUseClaimToken` only cover completion and
phase markers. Create/extend tests around `finishBatchJobAccounting`,
`MarkBatchJobUnpriceable`, and `FailBatchJobAccounting` to verify wrong tokens
return `ErrNotFound` and the correct token succeeds against the real
`newSqliteLogStore`, not just the fake store tests.
In `@framework/logstore/rdb.go`:
- Around line 417-474: UpsertBatchJob currently always does a Create with
OnConflict DoNothing followed by a separate Updates call, causing an unnecessary
second DB round trip on insert. Update RDBLogStore.UpsertBatchJob to perform a
single upsert path using the existing BatchJobID/ID conflict key and
clause.OnConflict DoUpdates so the insert and field refresh happen in one write,
while preserving the current nil/empty-field handling for BatchJob fields like
Model, ProviderStatus, and NextCheckAt.
In `@framework/logstore/tables.go`:
- Around line 273-317: The BatchJob identity index is not enforcing uniqueness
at the database layer, so duplicate logical rows can still be created if a
caller sets a different explicit ID. Update the BatchJob schema in tables.go to
make the Provider + BatchID identity index unique, and keep the existing
composite identity on BatchJob/TableName-backed persistence so the invariant is
enforced by the DB rather than only by UpsertBatchJob/BatchJobID caller
behavior.
In `@plugins/logging/main.go`:
- Around line 659-665: The sweeper setup in main.go is incorrectly locked to
schemas.OpenAI, which means FindDueBatchJobs will ignore batch jobs from other
providers and leave them stuck in pending state. Update the
SweeperConfig.Provider passed to batchaccounting.NewSweeper so the sweeper can
cover all providers, or make the provider selection configurable if OpenAI-only
behavior is intentional. Keep the fix aligned with the provider-agnostic flow
used by recordBatchJobLifecycle, BatchJob, and BatchResultFetcher.
- Around line 197-233: Batch results accounting is still running inline from
LoggerPlugin.PostLLMHook via accountBatchResults, which blocks the response path
for BatchResultsRequest. Move the batchaccounting.AccountBatchResults call out
of the synchronous path by scheduling it through the same deferred/bounded
goroutine pattern used by p.scheduleDeferredUsageUpdate and p.deferredUsageSem,
so the request can return immediately while accounting continues in the
background. Keep accountBatchResults as the worker function, but have the hook
enqueue it instead of invoking it directly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: b86efa85-37a4-4e2c-83d0-99be2b303ddb
📒 Files selected for processing (20)
framework/batchaccounting/accounting.goframework/batchaccounting/accounting_test.goframework/batchaccounting/sweeper.goframework/logstore/batch_jobs_test.goframework/logstore/hybrid.goframework/logstore/migrations.goframework/logstore/rdb.goframework/logstore/store.goframework/logstore/tables.goframework/modelcatalog/datasheet/cost.goframework/modelcatalog/datasheet/cost_test.goframework/modelcatalog/datasheet/types.goframework/modelcatalog/pricing.goplugins/governance/main.goplugins/governance/store.goplugins/logging/main.goplugins/logging/operations_test.goplugins/logging/utils.gotransports/bifrost-http/server/batch_accounting.gotransports/bifrost-http/server/server.go
3ad12a9 to
e0f2d32
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
framework/batchaccounting/accounting.go (1)
408-417: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winDuplicate
extractUsagedeclaration breaks the build.
extractUsageis declared twice at package scope — here (lines 408-417) and again at lines 433-439 — a redeclaration error. This stale switch-based version also hasreturn "", falseat line 415, which cannot satisfy the(extractedUsage, error)return type (the exact failure in the Snyk/CI build logs). It further omits theGeminicase that the map-based version at lines 421-426 handles. Delete this block and keep the map-drivenextractUsage.🐛 Remove the stale switch-based implementation
-func extractUsage(provider schemas.ModelProvider, fallbackModel string, item schemas.BatchResultItem) (extractedUsage, error) { - switch provider { - case schemas.OpenAI, schemas.Bedrock: - return extractResponseBodyUsage(fallbackModel, item) - case schemas.Anthropic: - return extractAnthropicUsage(fallbackModel, item) - default: - return "", false - } -} - type usageExtractor func(fallbackModel string, item schemas.BatchResultItem) (extractedUsage, error)🤖 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/batchaccounting/accounting.go` around lines 408 - 417, Remove the stale switch-based extractUsage implementation so only the map-driven extractUsage remains at package scope. The duplicate declaration conflicts with the later extractUsage function and the fallback return is invalid for the extractedUsage,error signature, so delete this block and keep the version that handles all providers including Gemini.Source: Pipeline failures
🧹 Nitpick comments (1)
framework/batchaccounting/sweeper.go (1)
195-197: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse a named reason constant for consistency.
markTerminalAsUnpriceableis otherwise called with constants (e.g.UnpriceableReasonMaxPollAttempts); this hardcoded"terminal_without_results"string diverges from that convention and is easy to typo/mismatch against reason filters.🤖 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/batchaccounting/sweeper.go` around lines 195 - 197, The markTerminalWithoutResults helper is passing a hardcoded reason string to markTerminalAsUnpriceable, which is inconsistent with the existing named reason constants pattern. Add a dedicated constant for the terminal-without-results reason and use it in markTerminalWithoutResults, matching the style used by symbols like UnpriceableReasonMaxPollAttempts so the reason stays consistent and typo-safe.
🤖 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.
Inline comments:
In `@framework/batchaccounting/sweeper.go`:
- Around line 106-167: The provider poll lease in Sweeper.sweepJob is being left
in KV for the full KVLeaseTTL, which blocks the intended retry cadence after
rescheduling. Update the sweepJob flow (and any helper like reschedule or
acquireProviderPollLease) so that once NextCheckAt is persisted, the lease key
for provider:batchID is cleared via schemas.KVStore.Delete, or reduce the lease
TTL to match the actual polling interval. Keep the fix centered on sweepJob and
reschedule so the lock does not survive past the next scheduled check.
---
Duplicate comments:
In `@framework/batchaccounting/accounting.go`:
- Around line 408-417: Remove the stale switch-based extractUsage implementation
so only the map-driven extractUsage remains at package scope. The duplicate
declaration conflicts with the later extractUsage function and the fallback
return is invalid for the extractedUsage,error signature, so delete this block
and keep the version that handles all providers including Gemini.
---
Nitpick comments:
In `@framework/batchaccounting/sweeper.go`:
- Around line 195-197: The markTerminalWithoutResults helper is passing a
hardcoded reason string to markTerminalAsUnpriceable, which is inconsistent with
the existing named reason constants pattern. Add a dedicated constant for the
terminal-without-results reason and use it in markTerminalWithoutResults,
matching the style used by symbols like UnpriceableReasonMaxPollAttempts so the
reason stays consistent and typo-safe.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 3b26d7e2-b479-4c6b-a4ac-255877887918
📒 Files selected for processing (20)
framework/batchaccounting/accounting.goframework/batchaccounting/accounting_test.goframework/batchaccounting/sweeper.goframework/logstore/batch_jobs_test.goframework/logstore/hybrid.goframework/logstore/migrations.goframework/logstore/rdb.goframework/logstore/store.goframework/logstore/tables.goframework/modelcatalog/datasheet/cost.goframework/modelcatalog/datasheet/cost_test.goframework/modelcatalog/datasheet/types.goframework/modelcatalog/pricing.goplugins/governance/main.goplugins/governance/store.goplugins/logging/main.goplugins/logging/operations_test.goplugins/logging/utils.gotransports/bifrost-http/server/batch_accounting.gotransports/bifrost-http/server/server.go
✅ Files skipped from review due to trivial changes (1)
- plugins/logging/utils.go
🚧 Files skipped from review as they are similar to previous changes (15)
- framework/modelcatalog/pricing.go
- framework/logstore/tables.go
- framework/logstore/batch_jobs_test.go
- framework/logstore/migrations.go
- framework/logstore/store.go
- framework/logstore/hybrid.go
- transports/bifrost-http/server/server.go
- framework/modelcatalog/datasheet/types.go
- framework/modelcatalog/datasheet/cost_test.go
- transports/bifrost-http/server/batch_accounting.go
- plugins/governance/store.go
- plugins/logging/operations_test.go
- framework/modelcatalog/datasheet/cost.go
- plugins/logging/main.go
- framework/logstore/rdb.go
292a9f2 to
0d22ac5
Compare
0d22ac5 to
2a1789c
Compare
c3108b9 to
1764325
Compare
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 5
♻️ Duplicate comments (1)
plugins/governance/main.go (1)
1753-1759: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winCharge the batch’s item count, not one request.
This still sets
requestDeltato1, while the batch pipeline tracks aggregate request counts. Request-based limits are therefore undercounted for multi-item batches; the test currently cements that behavior. Thread the aggregate total throughBatchUsageReport, use it here, and assert a multi-item delta.
plugins/governance/main.go#L1753-L1759: use the report’s aggregate request count instead of a fixed one.plugins/governance/accounting_test.go#L27-L41: supply a multi-item count and assert that exact request usage.As per coding guidelines, “fallback and retry attempts do not double-count or undercount usage unexpectedly.”
🤖 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 `@plugins/governance/main.go` around lines 1753 - 1759, The batch billing path in plugins/governance/main.go lines 1753-1759 must use BatchUsageReport’s aggregate request count instead of the fixed requestDelta value of 1; thread that total through the report and pass it to BumpRateLimitUsageBy while preserving retry deduplication. Update plugins/governance/accounting_test.go lines 27-41 to provide a multi-item count and assert the exact aggregate request usage.Source: Coding guidelines
🧹 Nitpick comments (1)
framework/batchaccounting/sweeper.go (1)
75-104: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd per-job panic recovery to the sweep loop.
Runticks forever callingSweepOnce, which processes all due jobs synchronously with norecover(). A panic while handling any single job (fetcher/store/pricing code) propagates out of the loop and permanently kills the background sweeper for every provider it serves, with no automatic restart visible in this file.♻️ Proposed fix
func (s *Sweeper) SweepOnce(ctx context.Context) { if s == nil || s.store == nil || s.pricing == nil || s.fetcher == nil { return } now := time.Now().UTC() jobs, err := s.store.FindDueBatchJobs(ctx, string(s.config.Provider), now, s.config.Limit) if err != nil { s.warn("batch accounting sweeper failed to find due jobs: %v", err) return } for _, job := range jobs { - s.sweepJob(ctx, job, now) + s.sweepJobSafe(ctx, job, now) } } + +func (s *Sweeper) sweepJobSafe(ctx context.Context, job *logstore.BatchJob, now time.Time) { + defer func() { + if r := recover(); r != nil { + s.warn("batch accounting sweeper recovered from panic job_id=%s: %v", job.ID, r) + } + }() + s.sweepJob(ctx, job, now) +}🤖 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/batchaccounting/sweeper.go` around lines 75 - 104, Add per-job panic recovery to the loop in SweepOnce, isolating each s.sweepJob(ctx, job, now) invocation so a panic from one job is recovered and does not terminate processing of later jobs or the Run loop; preserve normal error handling and continue iterating through all due jobs.
🤖 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.
Inline comments:
In `@core/providers/bedrock/bedrock.go`:
- Line 726: Update the SetExtraHeadersHTTP call to use the key-only Mantle
project resolver, replacing the ctx-based resolveMantleProjectID invocation with
resolveMantleProjectID(key). Keep the existing header configuration and
surrounding Bedrock request flow unchanged.
In `@core/providers/openai/openai.go`:
- Around line 2617-2632: Update the large-payload response handling around the
OpenAI transcription unmarshal logic to detect the plain-text response format
before JSON decoding. When request.Params.ResponseFormat is "text", assign the
raw lpResult.ResponseBody to response.Text and skip sonic.Unmarshal; preserve
the existing diarized and structured JSON paths for other formats.
In `@framework/batchaccounting/sweeper.go`:
- Around line 175-203: Update markTerminalAsUnpriceable to defer
s.deletePollLease(job) immediately on entry so claim failures, unclaimed jobs,
marking errors, and successful completion all release the poll lease. Remove the
now-redundant explicit s.deletePollLease(job) call from the maxPollAttempts
branch in reschedule.
- Around line 209-221: Update batchJobFromRetrieve to avoid the shallow job :=
*existing copy. Construct the returned BatchJob explicitly, preserving required
existing values while cloning pointer fields such as NextCheckAt, ClaimedBy,
ClaimToken, and ResultsURL so the returned job does not share mutable references
with existing.
In `@plugins/governance/tracker.go`:
- Around line 63-79: Replace the process-local, TTL-based batchBilled tracking
in plugins/governance/tracker.go:63-79 and its handling at
plugins/governance/tracker.go:299-335 with durable per-target settlement state
that survives restarts and cache eviction. Update the batch reporting flow in
plugins/governance/main.go:1742-1760 to claim each target through this durable
mechanism before incrementing usage, preserving retry safety when
GovernanceReportedAt is written after partial success. Add a fresh-tracker retry
test in plugins/governance/accounting_test.go:24-42 covering partial reporting
followed by marker failure, verifying already-settled targets are neither
double-counted nor skipped.
---
Duplicate comments:
In `@plugins/governance/main.go`:
- Around line 1753-1759: The batch billing path in plugins/governance/main.go
lines 1753-1759 must use BatchUsageReport’s aggregate request count instead of
the fixed requestDelta value of 1; thread that total through the report and pass
it to BumpRateLimitUsageBy while preserving retry deduplication. Update
plugins/governance/accounting_test.go lines 27-41 to provide a multi-item count
and assert the exact aggregate request usage.
---
Nitpick comments:
In `@framework/batchaccounting/sweeper.go`:
- Around line 75-104: Add per-job panic recovery to the loop in SweepOnce,
isolating each s.sweepJob(ctx, job, now) invocation so a panic from one job is
recovered and does not terminate processing of later jobs or the Run loop;
preserve normal error handling and continue iterating through all due jobs.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 9ae88a5e-0778-4045-bbaa-78de63c1a4fc
📒 Files selected for processing (32)
core/mcp/agent.gocore/providers/anthropic/anthropic.gocore/providers/bedrock/bedrock.gocore/providers/openai/openai.gocore/schemas/batch.gocore/schemas/batch_test.gocore/schemas/chatcompletions.gocore/schemas/usage_test.goframework/batchaccounting/accounting.goframework/batchaccounting/accounting_test.goframework/batchaccounting/sweeper.goframework/logstore/batch_jobs_test.goframework/logstore/clickhousemigrate.goframework/logstore/clickhousestore.goframework/logstore/hybrid.goframework/logstore/migrations.goframework/logstore/rdb.goframework/logstore/store.goframework/logstore/tables.goframework/modelcatalog/datasheet/cost.goframework/modelcatalog/datasheet/cost_test.goframework/modelcatalog/datasheet/types.goframework/modelcatalog/pricing.goplugins/governance/accounting_test.goplugins/governance/main.goplugins/governance/store.goplugins/governance/tracker.goplugins/logging/main.goplugins/logging/operations_test.goplugins/logging/utils.gotransports/bifrost-http/server/batch_accounting.gotransports/bifrost-http/server/server.go
🚧 Files skipped from review as they are similar to previous changes (20)
- core/schemas/usage_test.go
- transports/bifrost-http/server/server.go
- framework/modelcatalog/pricing.go
- plugins/logging/utils.go
- framework/logstore/hybrid.go
- core/schemas/batch_test.go
- framework/logstore/batch_jobs_test.go
- transports/bifrost-http/server/batch_accounting.go
- framework/modelcatalog/datasheet/types.go
- framework/logstore/tables.go
- core/schemas/chatcompletions.go
- framework/logstore/store.go
- plugins/governance/store.go
- core/mcp/agent.go
- framework/batchaccounting/accounting_test.go
- framework/logstore/migrations.go
- framework/logstore/rdb.go
- plugins/logging/operations_test.go
- plugins/logging/main.go
- framework/batchaccounting/accounting.go
🛑 Comments failed to post (5)
core/providers/bedrock/bedrock.go (1)
726-726: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Use the key-only Mantle project resolver.
resolveMantleProjectID(ctx, key)can dereference Vertex-only alias state for a Bedrock key. CallresolveMantleProjectID(key)instead.Based on learnings:
resolveMantleProjectIDmust be implemented at the key level only; the ctx-based variant can panic for non-Vertex keys.🤖 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 `@core/providers/bedrock/bedrock.go` at line 726, Update the SetExtraHeadersHTTP call to use the key-only Mantle project resolver, replacing the ctx-based resolveMantleProjectID invocation with resolveMantleProjectID(key). Keep the existing header configuration and surrounding Bedrock request flow unchanged.Source: Learnings
core/providers/openai/openai.go (1)
2617-2632: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Handle plain-text transcription responses in the large-payload path.
For
response_format: "text", OpenAI returns raw text, sosonic.Unmarshalfails here. Mirror the regular response path before attempting JSON decoding.Proposed fix
- if request.Params != nil && schemas.IsDiarizedTranscriptionFormat(request.Params.ResponseFormat) { + if request.Params != nil && schemas.IsPlainTextTranscriptionFormat(request.Params.ResponseFormat) { + response.Text = string(lpResult.ResponseBody) + } else if request.Params != nil && schemas.IsDiarizedTranscriptionFormat(request.Params.ResponseFormat) { var diarized openAIDiarizedTranscriptionResponse if err := sonic.Unmarshal(lpResult.ResponseBody, &diarized); err != nil {📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.// Unmarshal the upstream response body to preserve transcription text and fields if len(lpResult.ResponseBody) > 0 { response := &schemas.BifrostTranscriptionResponse{} if request.Params != nil && schemas.IsPlainTextTranscriptionFormat(request.Params.ResponseFormat) { response.Text = string(lpResult.ResponseBody) } else if request.Params != nil && schemas.IsDiarizedTranscriptionFormat(request.Params.ResponseFormat) { var diarized openAIDiarizedTranscriptionResponse if err := sonic.Unmarshal(lpResult.ResponseBody, &diarized); err != nil { return nil, providerUtils.NewBifrostOperationError(schemas.ErrProviderResponseUnmarshal, err) } response.Duration = diarized.Duration response.Task = diarized.Task response.Text = diarized.Text response.DiarizedSegments = diarized.Segments response.Usage = diarized.Usage } else if err := sonic.Unmarshal(lpResult.ResponseBody, response); err != nil { return nil, providerUtils.NewBifrostOperationError(schemas.ErrProviderResponseUnmarshal, err) }🤖 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 `@core/providers/openai/openai.go` around lines 2617 - 2632, Update the large-payload response handling around the OpenAI transcription unmarshal logic to detect the plain-text response format before JSON decoding. When request.Params.ResponseFormat is "text", assign the raw lpResult.ResponseBody to response.Text and skip sonic.Unmarshal; preserve the existing diarized and structured JSON paths for other formats.framework/batchaccounting/sweeper.go (2)
175-203: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Poll lease is leaked when the accounting claim fails or errors.
markTerminalAsUnpriceableonly callss.deletePollLease(job)after a successfulMarkBatchJobUnpriceable. IfClaimBatchJobAccountingerrors or returns!claimed(line 193-198), the function returns without releasing the lease acquired earlier insweepJob. That job'sbatch-accounting:poll:<provider>:<batchID>key then blocks re-polling for the fullKVLeaseTTL(5m default) instead of the intended cadence — the same class of bug flagged and fixed forreschedulepreviously. The explicit extras.deletePollLease(job)call right aftermarkTerminalAsUnpriceable(...)at line 179 (in the max-poll-attempts branch) is a workaround for exactly this gap, butmarkTerminalWithoutResults(called fromsweepJobline 135) has no equivalent fallback, so that call path still leaks the lease.Move the release into a
deferinsidemarkTerminalAsUnpriceableso every exit path releases the lease, and drop the now-redundant explicit call inreschedule.🔒 Proposed fix
func (s *Sweeper) markTerminalAsUnpriceable(ctx context.Context, job *logstore.BatchJob, reason string) { + defer s.deletePollLease(job) token, claimed, err := s.store.ClaimBatchJobAccounting(ctx, job.ID, s.config.ClaimedBy, defaultClaimTTL) if err != nil || !claimed { if err != nil { s.warn("batch accounting sweeper failed to claim for unpriceable provider=%s batch_id=%s job_id=%s reason=%s: %v", job.Provider, job.BatchID, job.ID, reason, err) } return } if err := s.store.MarkBatchJobUnpriceable(ctx, job.ID, token, reason, nil); err != nil { s.warn("batch accounting sweeper failed to mark unpriceable batch provider=%s batch_id=%s job_id=%s reason=%s: %v", job.Provider, job.BatchID, job.ID, reason, err) } - s.deletePollLease(job) }if job.PollAttempts >= maxPollAttempts { s.markTerminalAsUnpriceable(ctx, job, UnpriceableReasonMaxPollAttempts) - s.deletePollLease(job) return }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.func (s *Sweeper) reschedule(ctx context.Context, job *logstore.BatchJob, now time.Time) { job.PollAttempts++ if job.PollAttempts >= maxPollAttempts { s.markTerminalAsUnpriceable(ctx, job, UnpriceableReasonMaxPollAttempts) return } next := s.nextCheckAt(job, now) job.NextCheckAt = &next if err := s.store.UpsertBatchJob(ctx, job); err != nil { s.warn("batch accounting sweeper failed to reschedule provider=%s batch_id=%s job_id=%s: %v", job.Provider, job.BatchID, job.ID, err) return } s.deletePollLease(job) } func (s *Sweeper) markTerminalAsUnpriceable(ctx context.Context, job *logstore.BatchJob, reason string) { defer s.deletePollLease(job) token, claimed, err := s.store.ClaimBatchJobAccounting(ctx, job.ID, s.config.ClaimedBy, defaultClaimTTL) if err != nil || !claimed { if err != nil { s.warn("batch accounting sweeper failed to claim for unpriceable provider=%s batch_id=%s job_id=%s reason=%s: %v", job.Provider, job.BatchID, job.ID, reason, err) } return } if err := s.store.MarkBatchJobUnpriceable(ctx, job.ID, token, reason, nil); err != nil { s.warn("batch accounting sweeper failed to mark unpriceable batch provider=%s batch_id=%s job_id=%s reason=%s: %v", job.Provider, job.BatchID, job.ID, reason, err) } }🤖 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/batchaccounting/sweeper.go` around lines 175 - 203, Update markTerminalAsUnpriceable to defer s.deletePollLease(job) immediately on entry so claim failures, unclaimed jobs, marking errors, and successful completion all release the poll lease. Remove the now-redundant explicit s.deletePollLease(job) call from the maxPollAttempts branch in reschedule.
209-221: 🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash ast-grep run --pattern 'type BatchJob struct { $$$ }' --lang go framework/logstoreRepository: maximhq/bifrost
Length of output: 3465
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== batchJobFromRetrieve context ==" sed -n '180,250p' framework/batchaccounting/sweeper.go echo echo "== BatchJob struct and nearby helpers ==" sed -n '288,360p' framework/logstore/tables.go echo echo "== references to BatchJob field mutations in framework/batchaccounting and framework/logstore ==" rg -n "BatchJob|ClaimedBy|ClaimToken|BudgetIDs|RateLimitIDs|LastError|NextCheckAt|PollAttempts|AccountingStatus|SelectedKeyID|VirtualKeyID|UnpriceableReason|AggregateLogWrittenAt|GovernanceReportedAt" framework/batchaccounting framework/logstoreRepository: maximhq/bifrost
Length of output: 50371
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== sweepJob flow around batchJobFromRetrieve ==" sed -n '100,170p' framework/batchaccounting/sweeper.go echo echo "== store upsert behavior for BatchJob pointer fields ==" sed -n '421,485p' framework/logstore/rdb.go echo sed -n '156,225p' framework/logstore/clickhousestore.go echo echo "== exact helper definitions that merge/copy BatchJob pointer fields ==" sed -n '240,270p' framework/batchaccounting/accounting.goRepository: maximhq/bifrost
Length of output: 7620
Avoid the shallow struct copy here
BatchJobstill has pointer fields (NextCheckAt,ClaimedBy,ClaimToken,ResultsURL, etc.), sojob := *existingshares those references withexisting. Build the copy explicitly or clone the pointer fields you keep.🤖 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/batchaccounting/sweeper.go` around lines 209 - 221, Update batchJobFromRetrieve to avoid the shallow job := *existing copy. Construct the returned BatchJob explicitly, preserving required existing values while cloning pointer fields such as NextCheckAt, ClaimedBy, ClaimToken, and ResultsURL so the returned job does not share mutable references with existing.Source: Path instructions
plugins/governance/tracker.go (1)
63-79: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Persist per-target batch billing idempotency.
batchBilledis lost on restart and evicted after seven days, butGovernanceReportedAtis written only after reporting succeeds. A retry after either gap re-applies targets that succeeded before the marker write, double-counting budget/rate-limit usage.
plugins/governance/tracker.go#L63-L79: replace the process-local/TTL-only guarantee with durable target-level settlement state.plugins/governance/tracker.go#L299-L335: do not make retry safety depend on cache retention.plugins/governance/main.go#L1742-L1760: claim/apply each target through the durable settlement mechanism before incrementing usage.plugins/governance/accounting_test.go#L24-L42: add a retry case with a fresh tracker after a partial report/marker failure.As per coding guidelines, “fallback and retry attempts do not double-count or undercount usage unexpectedly.”
📍 Affects 3 files
plugins/governance/tracker.go#L63-L79(this comment)plugins/governance/tracker.go#L299-L335plugins/governance/main.go#L1742-L1760plugins/governance/accounting_test.go#L24-L42🤖 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 `@plugins/governance/tracker.go` around lines 63 - 79, Replace the process-local, TTL-based batchBilled tracking in plugins/governance/tracker.go:63-79 and its handling at plugins/governance/tracker.go:299-335 with durable per-target settlement state that survives restarts and cache eviction. Update the batch reporting flow in plugins/governance/main.go:1742-1760 to claim each target through this durable mechanism before incrementing usage, preserving retry safety when GovernanceReportedAt is written after partial success. Add a fresh-tracker retry test in plugins/governance/accounting_test.go:24-42 covering partial reporting followed by marker failure, verifying already-settled targets are neither double-counted nor skipped.Source: Coding guidelines
There was a problem hiding this comment.
🧹 Nitpick comments (1)
core/providers/openai/files_test.go (1)
13-24: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover a non-zero batch item index.
This test only passes a one-item slice, so it cannot catch an indexing regression in the new error context. Add a valid item before the invalid item and assert
batch request item 1: custom_id is required.Suggested test
+func TestConvertRequestsToJSONLReportsNonZeroCustomIDIndex(t *testing.T) { + _, err := ConvertRequestsToJSONL([]schemas.BatchRequestItem{ + {CustomID: "request-0"}, + {CustomID: " "}, + }) + require.Error(t, err) + assert.Contains(t, err.Error(), "batch request item 1: custom_id is required") +}🤖 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 `@core/providers/openai/files_test.go` around lines 13 - 24, Update TestConvertRequestsToJSONLRequiresCustomID to include a valid batch request item before the invalid CustomID item, then assert the error references batch request item 1 and retains the custom_id-required message for both empty and whitespace-only IDs.
🤖 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 `@core/providers/openai/files_test.go`:
- Around line 13-24: Update TestConvertRequestsToJSONLRequiresCustomID to
include a valid batch request item before the invalid CustomID item, then assert
the error references batch request item 1 and retains the custom_id-required
message for both empty and whitespace-only IDs.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 165d4439-9527-4aa1-9bfb-6246270293b3
📒 Files selected for processing (26)
core/mcp/agent.gocore/providers/anthropic/anthropic.gocore/providers/bedrock/bedrock.gocore/providers/openai/files.gocore/providers/openai/files_test.gocore/providers/openai/openai.gocore/schemas/batch.gocore/schemas/batch_test.gocore/schemas/chatcompletions.gocore/schemas/usage_test.goframework/batchaccounting/accounting.goframework/batchaccounting/accounting_test.goframework/batchaccounting/sweeper.goframework/logstore/batch_jobs_test.goframework/logstore/clickhousemigrate.goframework/logstore/clickhousestore.goframework/logstore/migrations.goframework/logstore/rdb.goframework/logstore/tables.goframework/modelcatalog/datasheet/cost.goframework/modelcatalog/datasheet/cost_test.goplugins/governance/accounting_test.goplugins/governance/main.goplugins/governance/tracker.goplugins/logging/main.gotransports/bifrost-http/server/batch_accounting.go
🚧 Files skipped from review as they are similar to previous changes (24)
- core/schemas/usage_test.go
- core/providers/openai/openai.go
- plugins/governance/accounting_test.go
- core/schemas/batch_test.go
- framework/logstore/clickhousemigrate.go
- framework/logstore/tables.go
- framework/logstore/batch_jobs_test.go
- core/providers/anthropic/anthropic.go
- transports/bifrost-http/server/batch_accounting.go
- framework/logstore/migrations.go
- core/mcp/agent.go
- core/schemas/batch.go
- core/providers/bedrock/bedrock.go
- framework/logstore/clickhousestore.go
- core/schemas/chatcompletions.go
- framework/batchaccounting/accounting.go
- framework/batchaccounting/sweeper.go
- framework/modelcatalog/datasheet/cost.go
- plugins/governance/main.go
- plugins/governance/tracker.go
- framework/modelcatalog/datasheet/cost_test.go
- plugins/logging/main.go
- framework/batchaccounting/accounting_test.go
- framework/logstore/rdb.go
The merge-base changed after approval.
## Summary Briefly explain the purpose of this PR and the problem it solves. ## Changes - What was changed and why - Any notable design decisions or trade-offs ## Type of change - [ ] Bug fix - [ ] Feature - [ ] Refactor - [x] Documentation - [ ] Chore/CI ## Affected areas - [ ] Core (Go) - [ ] Transports (HTTP) - [ ] Providers/Integrations - [ ] Plugins - [ ] UI (React) - [x] Docs ## How to test Describe the steps to validate this change. Include commands and expected outcomes. ```sh # Core/Transports go version go test ./... # UI cd ui pnpm i || npm i pnpm test || npm test pnpm build || npm run build ``` If adding new configs or environment variables, document them here. ## Screenshots/Recordings If UI changes, add before/after screenshots or short clips. ## Breaking changes - [ ] Yes - [ ] No If yes, describe impact and migration instructions. ## Related issues Link related issues and discussions. Example: Closes #123 ## Security considerations Note any security implications (auth, secrets, PII, sandboxing, etc.). ## Checklist - [ ] I read `docs/contributing/README.md` and followed the guidelines - [ ] I added/updated tests where appropriate - [ ] I updated documentation where needed - [ ] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
1764325 to
1ecd99e
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@framework/batchaccounting/sweeper.go`:
- Around line 106-207: Defer poll-lease release immediately after `sweepJob`
successfully acquires the lease, ensuring every subsequent return path cleans it
up. Remove the branch-specific `deletePollLease` calls from `reschedule` and
`markTerminalAsUnpriceable`, since lease ownership and release should be handled
by `sweepJob`, including upsert, accounting, claim-error, and unclaimed paths.
In `@framework/logstore/tables.go`:
- Around line 288-291: Enforce canonical identity in BatchJob persistence:
update UpsertBatchJob to validate that a non-empty ID matches the canonical
value derived from Provider and BatchID, rejecting mismatches, and make
idx_batch_jobs_identity unique to prevent duplicate provider/batch pairs.
Preserve valid upserts and existing identity behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: bacc7c6a-165b-4b68-96fe-777c986b7c85
📒 Files selected for processing (32)
core/mcp/agent.gocore/providers/anthropic/anthropic.gocore/providers/bedrock/bedrock.gocore/providers/openai/files.gocore/providers/openai/files_test.gocore/providers/openai/openai.gocore/schemas/batch.gocore/schemas/batch_test.gocore/schemas/chatcompletions.gocore/schemas/usage_test.goframework/batchaccounting/accounting.goframework/batchaccounting/accounting_test.goframework/batchaccounting/sweeper.goframework/logstore/batch_jobs_test.goframework/logstore/clickhousemigrate.goframework/logstore/clickhousestore.goframework/logstore/hybrid.goframework/logstore/migrations.goframework/logstore/rdb.goframework/logstore/store.goframework/logstore/tables.goframework/modelcatalog/datasheet/cost.goframework/modelcatalog/datasheet/cost_test.goframework/modelcatalog/datasheet/types.goframework/modelcatalog/pricing.goplugins/governance/accounting_test.goplugins/governance/main.goplugins/governance/store.goplugins/governance/tracker.goplugins/logging/main.goplugins/logging/operations_test.goplugins/logging/utils.go
🚧 Files skipped from review as they are similar to previous changes (30)
- framework/modelcatalog/datasheet/types.go
- core/schemas/usage_test.go
- core/providers/openai/files.go
- core/schemas/batch_test.go
- core/schemas/chatcompletions.go
- framework/logstore/store.go
- core/schemas/batch.go
- core/providers/openai/files_test.go
- framework/logstore/hybrid.go
- framework/logstore/batch_jobs_test.go
- plugins/governance/accounting_test.go
- core/providers/anthropic/anthropic.go
- core/providers/openai/openai.go
- plugins/logging/utils.go
- plugins/governance/tracker.go
- core/mcp/agent.go
- plugins/governance/store.go
- framework/logstore/migrations.go
- core/providers/bedrock/bedrock.go
- framework/modelcatalog/datasheet/cost_test.go
- framework/modelcatalog/pricing.go
- framework/logstore/clickhousemigrate.go
- framework/modelcatalog/datasheet/cost.go
- framework/logstore/clickhousestore.go
- framework/batchaccounting/accounting.go
- plugins/governance/main.go
- framework/batchaccounting/accounting_test.go
- plugins/logging/operations_test.go
- plugins/logging/main.go
- framework/logstore/rdb.go
1ecd99e to
7c3c2c6
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
framework/logstore/migrations.go (1)
1617-1619: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winRollback should use
dropColumnIfExistsfor consistency.
RollbackcallsMigrator().DropColumndirectly without an existence check, unlike the siblingmigrationAddMetadataColumnrollback (lines 1643-1649) in this same file, which uses thedropColumnIfExistshelper. Using the guarded helper here avoids an error if this rollback is ever invoked when the column is already absent.♻️ Proposed fix
Rollback: func(tx *gorm.DB) error { - return tx.WithContext(ctx).Migrator().DropColumn(&BatchJob{}, "endpoint") + return dropColumnIfExists(tx.WithContext(ctx), logger, &BatchJob{}, "endpoint") },🤖 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/migrations.go` around lines 1617 - 1619, Update the Rollback function for the endpoint migration to call the existing dropColumnIfExists helper instead of invoking Migrator().DropColumn directly, matching the guarded behavior used by migrationAddMetadataColumn.framework/logstore/clickhousestore.go (1)
238-249: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
bifrost.Ptr()/new(expr)over&for pointer fields.
job.ClaimedBy = &claimedByandjob.ClaimToken = &tokentake the address of simple unmodified values, andexpires := now.Add(ttl); job.ClaimExpiresAt = &expiresintroduces an intermediate variable just to address a computed value. Based on learnings, this repo prefersbifrost.Ptr()for simple values and Go 1.26'snew(expr)form for computed values instead of&. The same pattern recurs at line 303 (job.UnpriceableReason = &reason) infinishBatchJobAccounting.Based on learnings: In the maximhq/bifrost repository, prefer using bifrost.Ptr() to create pointers instead of the address operator (&) even when & would be valid syntactically. and separately, for computed values, the convention is to use Go 1.26's
new(expr)form rather than an intermediate variable.♻️ Proposed fix
token := uuid.NewString() job.AccountingStatus = BatchJobAccountingStatusProcessing - job.ClaimedBy = &claimedBy - job.ClaimToken = &token - expires := now.Add(ttl) - job.ClaimExpiresAt = &expires + job.ClaimedBy = bifrost.Ptr(claimedBy) + job.ClaimToken = bifrost.Ptr(token) + job.ClaimExpiresAt = new(now.Add(ttl))Since this touches a cross-module
bifrost.Ptr()import fromframework/logstore, please confirm the import direction is acceptable (no cycle) before applying.🤖 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/clickhousestore.go` around lines 238 - 249, Update the batch-job pointer assignments in the surrounding claim flow to use bifrost.Ptr() for claimedBy and token, and use new(now.Add(ttl)) for ClaimExpiresAt instead of the expires temporary and address operator. Apply the same bifrost.Ptr() convention to UnpriceableReason in finishBatchJobAccounting, adding the import only after confirming it does not introduce an import cycle.
🤖 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.
Inline comments:
In `@framework/logstore/clickhousestore.go`:
- Around line 273-311: Update finishBatchJobAccounting so job.LastError is set
to nil when reportedErr is nil, and only store a non-nil error message when
reportedErr is present. Preserve the existing accounting status, claim cleanup,
and unpriceable-reason updates.
In `@plugins/governance/tracker.go`:
- Around line 63-67: Replace the process-local batchBilled map with durable
idempotency in the ReportBatchUsage flow: persist each target-level billing key
alongside the usage mutation, or atomically combine the billing bump with
MarkBatchJobGovernanceReported/CompleteBatchJobAccounting. Ensure retries and
restarts cannot bill the same RequestID and target more than once.
---
Nitpick comments:
In `@framework/logstore/clickhousestore.go`:
- Around line 238-249: Update the batch-job pointer assignments in the
surrounding claim flow to use bifrost.Ptr() for claimedBy and token, and use
new(now.Add(ttl)) for ClaimExpiresAt instead of the expires temporary and
address operator. Apply the same bifrost.Ptr() convention to UnpriceableReason
in finishBatchJobAccounting, adding the import only after confirming it does not
introduce an import cycle.
In `@framework/logstore/migrations.go`:
- Around line 1617-1619: Update the Rollback function for the endpoint migration
to call the existing dropColumnIfExists helper instead of invoking
Migrator().DropColumn directly, matching the guarded behavior used by
migrationAddMetadataColumn.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 14668390-2522-42e1-a296-f53e895cddbd
📒 Files selected for processing (26)
core/mcp/agent.gocore/providers/anthropic/anthropic.gocore/providers/bedrock/bedrock.gocore/providers/openai/files.gocore/providers/openai/files_test.gocore/providers/openai/openai.gocore/schemas/batch.gocore/schemas/batch_test.gocore/schemas/chatcompletions.gocore/schemas/usage_test.goframework/batchaccounting/accounting.goframework/batchaccounting/accounting_test.goframework/batchaccounting/sweeper.goframework/logstore/batch_jobs_test.goframework/logstore/clickhousemigrate.goframework/logstore/clickhousestore.goframework/logstore/migrations.goframework/logstore/rdb.goframework/logstore/tables.goframework/modelcatalog/datasheet/cost.goframework/modelcatalog/datasheet/cost_test.goplugins/governance/accounting_test.goplugins/governance/main.goplugins/governance/tracker.goplugins/logging/main.gotransports/bifrost-http/server/batch_accounting.go
🚧 Files skipped from review as they are similar to previous changes (22)
- core/providers/openai/files_test.go
- transports/bifrost-http/server/batch_accounting.go
- framework/logstore/tables.go
- core/providers/openai/files.go
- core/providers/openai/openai.go
- core/schemas/chatcompletions.go
- core/providers/bedrock/bedrock.go
- core/providers/anthropic/anthropic.go
- core/schemas/batch_test.go
- framework/logstore/clickhousemigrate.go
- plugins/governance/accounting_test.go
- core/schemas/usage_test.go
- framework/batchaccounting/accounting.go
- framework/modelcatalog/datasheet/cost_test.go
- plugins/governance/main.go
- plugins/logging/main.go
- framework/logstore/batch_jobs_test.go
- core/schemas/batch.go
- framework/modelcatalog/datasheet/cost.go
- framework/batchaccounting/accounting_test.go
- framework/batchaccounting/sweeper.go
- framework/logstore/rdb.go
7c3c2c6 to
9a3d94a
Compare
9a3d94a to
c071763
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
framework/batchaccounting/sweeper.go (1)
109-117: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftOnly release the poll lease when this worker still owns it.
The TTL can expire while provider fetching or accounting is running. A second worker may then acquire the same key, after which this deferred unconditional
Deleteremoves the second worker’s lease. Store a unique lease token and use compare-and-delete, or otherwise verify ownership atomically before release.As per coding guidelines, shared state must be race-safe.
Also applies to: 241-248
🤖 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/batchaccounting/sweeper.go` around lines 109 - 117, Update the lease lifecycle around acquireProviderPollLease and the deferred s.deletePollLease call to retain a unique ownership token from acquisition and release the lease only through an atomic compare-and-delete (or equivalent ownership check). Ensure a later worker’s renewed or newly acquired lease is never deleted, and keep shared lease state race-safe.Source: Coding guidelines
🤖 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.
Inline comments:
In `@framework/batchaccounting/accounting.go`:
- Around line 506-547: Validate all untrusted provider usage and cost fields
before pricing, including the token fields parsed by anthropicUsageFromValue and
the related parsing paths through the covered range. Reject any negative value
with an error before constructing usage or marking a row priced, while
preserving valid zero and positive values. Add malformed-usage coverage for
negative token and total_cost inputs.
---
Outside diff comments:
In `@framework/batchaccounting/sweeper.go`:
- Around line 109-117: Update the lease lifecycle around
acquireProviderPollLease and the deferred s.deletePollLease call to retain a
unique ownership token from acquisition and release the lease only through an
atomic compare-and-delete (or equivalent ownership check). Ensure a later
worker’s renewed or newly acquired lease is never deleted, and keep shared lease
state race-safe.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 2a8d91c9-be82-41b4-80aa-d5cefb26187f
📒 Files selected for processing (28)
core/mcp/agent.gocore/providers/anthropic/anthropic.gocore/providers/bedrock/bedrock.gocore/providers/openai/files.gocore/providers/openai/files_test.gocore/providers/openai/openai.gocore/schemas/batch.gocore/schemas/batch_test.gocore/schemas/chatcompletions.gocore/schemas/usage_test.goframework/batchaccounting/accounting.goframework/batchaccounting/accounting_test.goframework/batchaccounting/sweeper.goframework/logstore/batch_jobs_test.goframework/logstore/clickhousemigrate.goframework/logstore/clickhousestore.goframework/logstore/migrations.goframework/logstore/rdb.goframework/logstore/tables.goframework/modelcatalog/datasheet/cost.goframework/modelcatalog/datasheet/cost_test.goframework/modelcatalog/pricing.goplugins/governance/accounting_test.goplugins/governance/main.goplugins/governance/tracker.goplugins/logging/main.goplugins/logging/operations_test.gotransports/bifrost-http/server/batch_accounting.go
💤 Files with no reviewable changes (2)
- framework/modelcatalog/pricing.go
- framework/logstore/migrations.go
🚧 Files skipped from review as they are similar to previous changes (22)
- core/providers/openai/files.go
- core/providers/openai/files_test.go
- core/mcp/agent.go
- core/schemas/usage_test.go
- core/providers/openai/openai.go
- plugins/governance/accounting_test.go
- core/schemas/batch.go
- framework/modelcatalog/datasheet/cost_test.go
- framework/logstore/tables.go
- core/schemas/chatcompletions.go
- core/providers/bedrock/bedrock.go
- core/schemas/batch_test.go
- framework/logstore/batch_jobs_test.go
- framework/logstore/clickhousemigrate.go
- plugins/governance/main.go
- plugins/governance/tracker.go
- framework/modelcatalog/datasheet/cost.go
- plugins/logging/operations_test.go
- framework/logstore/rdb.go
- framework/logstore/clickhousestore.go
- plugins/logging/main.go
- transports/bifrost-http/server/batch_accounting.go
| func anthropicUsageFromValue(value interface{}) (*schemas.BifrostLLMUsage, error) { | ||
| bytes, err := sonic.Marshal(value) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| var usage struct { | ||
| InputTokens int `json:"input_tokens"` | ||
| CacheCreationInputTokens int `json:"cache_creation_input_tokens"` | ||
| CacheReadInputTokens int `json:"cache_read_input_tokens"` | ||
| CacheCreation struct { | ||
| Ephemeral5mInputTokens int `json:"ephemeral_5m_input_tokens"` | ||
| Ephemeral1hInputTokens int `json:"ephemeral_1h_input_tokens"` | ||
| } `json:"cache_creation"` | ||
| OutputTokens int `json:"output_tokens"` | ||
| } | ||
| if err := sonic.Unmarshal(bytes, &usage); err != nil { | ||
| return nil, err | ||
| } | ||
| promptTokens := usage.InputTokens + usage.CacheCreationInputTokens + usage.CacheReadInputTokens | ||
| totalTokens := promptTokens + usage.OutputTokens | ||
| if totalTokens == 0 { | ||
| return &schemas.BifrostLLMUsage{}, nil | ||
| } | ||
| out := &schemas.BifrostLLMUsage{ | ||
| PromptTokens: promptTokens, | ||
| CompletionTokens: usage.OutputTokens, | ||
| TotalTokens: totalTokens, | ||
| } | ||
| if usage.CacheCreationInputTokens > 0 || usage.CacheReadInputTokens > 0 { | ||
| out.PromptTokensDetails = &schemas.ChatPromptTokensDetails{ | ||
| CachedReadTokens: usage.CacheReadInputTokens, | ||
| CachedWriteTokens: usage.CacheCreationInputTokens, | ||
| } | ||
| if usage.CacheCreation.Ephemeral5mInputTokens > 0 || usage.CacheCreation.Ephemeral1hInputTokens > 0 { | ||
| out.PromptTokensDetails.CachedWriteTokenDetails = &schemas.ChatCachedWriteTokenDetails{ | ||
| CachedWriteTokens5m: usage.CacheCreation.Ephemeral5mInputTokens, | ||
| CachedWriteTokens1h: usage.CacheCreation.Ephemeral1hInputTokens, | ||
| } | ||
| } | ||
| } | ||
| return out, nil | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Reject negative provider usage and cost values before pricing.
These external payloads accept signed token counts and provider costs unchecked. A row with negative tokens or total_cost can be marked priced and reduce the aggregate bill. Return an error for any negative usage/cost field and add malformed-usage coverage.
As per coding guidelines, validate all untrusted input.
Also applies to: 549-627
🤖 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/batchaccounting/accounting.go` around lines 506 - 547, Validate all
untrusted provider usage and cost fields before pricing, including the token
fields parsed by anthropicUsageFromValue and the related parsing paths through
the covered range. Reject any negative value with an error before constructing
usage or marking a row priced, while preserving valid zero and positive values.
Add malformed-usage coverage for negative token and total_cost inputs.
Source: Coding guidelines
44564de to
493bff0
Compare
|
closed because this is the new stack: #5296 (comment) |
2 participants
Summary
Briefly explain the purpose of this PR and the problem it solves.
Changes
Type of change
Affected areas
How to test
Describe the steps to validate this change. Include commands and expected outcomes.
If adding new configs or environment variables, document them here.
Screenshots/Recordings
If UI changes, add before/after screenshots or short clips.
Breaking changes
If yes, describe impact and migration instructions.
Related issues
Link related issues and discussions. Example: Closes #123
Security considerations
Note any security implications (auth, secrets, PII, sandboxing, etc.).
Checklist
docs/contributing/README.mdand followed the guidelines