Skip to content

Fix token budget lock contention for streaming queries and Supabase setup#3

Merged
Aryan1718 merged 2 commits into
developfrom
fix-token-budget-timeout
Jun 6, 2026
Merged

Fix token budget lock contention for streaming queries and Supabase setup#3
Aryan1718 merged 2 commits into
developfrom
fix-token-budget-timeout

Conversation

@Aryan1718

Copy link
Copy Markdown
Owner

Summary

This PR fixes token budget lock contention that could cause streaming query failures with PostgreSQL statement
timeouts.

The issue showed up around workspace_daily_usage updates during query execution. The previous flow could keep the
token budget row lock open too long while a request was still doing retrieval, LLM work, or SSE streaming. Under
concurrent requests for the same workspace, that created blocking and timeout failures.

What changed

  • updated the token budget flow to use short isolated DB transactions for:
    • token reservation
    • token commit
    • token release
  • avoided holding the workspace_daily_usage row lock across long-running query and streaming work
  • added a regression test around budget operations while the caller session already has an active transaction
  • included helper scripts used during schema/bootstrap work:
    • server/scripts/bootstrap_schema.py
    • server/scripts/inspect_schema.py
  • included follow-up task documentation captured during debugging:
    • docs/task/ingestion-follow-up-hardening-and-validation.md

Root cause

The budget reservation path used the request DB session. In API flows, that session could already be inside an open
transaction before reserve_tokens() was called. That meant the row lock on workspace_daily_usage could remain open
much longer than intended, especially in /query/stream, where the request stays alive while tokens are streamed
back.

Concurrent requests for the same workspace/day could then block on the same row and eventually fail with a statement
timeout.

Result

This change keeps budget mutations atomic, but limits the lock lifetime to the budget update itself instead of the
full request lifecycle.

Validation

  • verified runtime import and server startup after the change
  • verified the fix branch was created on top of the merged ingestion workflow branch
  • existing token budget test execution hit a separate pre-existing SQLite/UUID test setup problem in this repo, not a
    regression from this change

@Aryan1718 Aryan1718 merged commit 8111647 into develop Jun 6, 2026
1 of 2 checks passed
@Aryan1718 Aryan1718 deleted the fix-token-budget-timeout branch June 6, 2026 23:42
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