Skip to content

Harvest quota headers on relay transports - #130

Closed
iceteaSA wants to merge 2 commits into
cortexkit:mainfrom
iceteaSA:feat/relay-header-harvest
Closed

Harvest quota headers on relay transports#130
iceteaSA wants to merge 2 commits into
cortexkit:mainfrom
iceteaSA:feat/relay-header-harvest

Conversation

@iceteaSA

@iceteaSA iceteaSA commented Jul 17, 2026

Copy link
Copy Markdown
Contributor

Stacked on #129 (first 14 commits are that PR; merge it first — this diff then reduces to 5 commits).

#129 gated quota-header harvesting to direct fetches. This extends it to both relay transports, closing the gap where enabling the relay silently stops passive quota freshening:

  • sendViaRelay gains an optional onResponseHeaders callback that delivers genuine upstream response headers to the caller:
    • WebSocket: fired from the response_start control message (the Worker already forwards real upstream headers there). The optimistically-resolved Response keeps its synthetic headers; the callback is the only channel that carries the real ones. Fires at most once per request, survives retry-before-response, never fires on error paths, and a throwing callback cannot break the stream.
    • HTTP relay: fired with the relay response's headers (the Worker copies upstream headers verbatim).
    • Relay-to-direct fallback: the callback does not fire; the direct path's existing harvest covers it (no double push).
  • The plugin wires the callback at the send site with the same served-account binding the direct harvest uses; all merge/rotation/attribution semantics are unchanged and transport-agnostic.
  • e2e now covers the relay transport both ways: quota headers land as source: headers bound to the served account, and header-less relay responses leave the store untouched. The previous direct-only e2e assertion was timing-dependent (it read state before async persistence); it now uses a condition-based wait.

Tests: core relay callback lifecycle (once-per-request, retry carry-over, error paths, optimistic mode), opencode wiring (single push on fallback, synthetic headers never harvested), Miniflare Worker round-trip asserting response_start carries the unified quota headers, e2e both transports. Full suites green (core 49, opencode 876, pi 48, e2e 6).


View with Codesmith Autofix with Codesmith
Need help on this PR? Tag /codesmith with what you need. Autofix is disabled.


Summary by cubic

Extends passive quota-header harvesting to relay transports and direct fetches, and normalizes Anthropic unified headers into the shared quota store. Surfaces tier/plan, credits, binding-window, and fallback advice across the UI with token-fenced, fire-and-forget profile persistence so commands never block on profile writes.

  • New Features

    • Relay/header harvest: sendViaRelay exposes genuine upstream headers (HTTP + WebSocket response_start) via onResponseHeaders; fires once on commit; ignored on errors, relay→direct fallback, and optimistic frames; retry swaps staged headers so last-accepted wins.
    • Quotas: quota-headers normalizes anthropic-ratelimit-unified-*, derives binding-window and fallback advice; QuotaManager merges header pushes with poll (headers replace five_hour/seven_day; poll remains source for scoped windows, idle fallbacks, and boot).
    • Account metadata: oauth-profile fetch with 7‑day, token‑fingerprint TTL; UI shows tier labels on account lists/dialogs/sidebar/TUI, extra‑usage credits with stable currency formatting and exhaustion, plus a binding marker and fallback advice in quota rows.
  • Bug Fixes

    • Quota correctness: preserve persisted poll backoff; don’t clobber scoped windows or cached data on partial header frames; header pushes idempotent across retries/rotation; header‑sourced exhaustion shares poll staleness to block unsafe API‑key fallback.
    • Relay lifecycle: genuine headers delivered across retries with last‑accepted wins; single fire per request; no double‑push on relay→direct fallback; synthetic optimistic headers never harvested.
    • Profile safety: delayed‑clear protection; boot hydration owns profile persistence; saves are detached from command rendering; per‑plugin fetch isolation; generic account saves no longer overwrite profile state.

Written for commit 9a1a72f. Summary will update on new commits.

Review in cubic

Greptile Summary

This PR extends quota and account metadata handling across relay and display paths. The main changes are:

  • Relay transports now expose genuine upstream response headers for quota harvesting.
  • OpenCode wires harvested quota headers to the served OAuth account.
  • Quota snapshots now merge header-sourced updates with poll-owned state.
  • Account profile metadata is fetched and shown in quota, account, sidebar, and TUI surfaces.
  • Tests and docs cover relay header delivery, quota surfaces, and profile display behavior.

Confidence Score: 5/5

This looks safe to merge.

  • No blocking issues found in the changed code.

Important Files Changed

Filename Overview
packages/core/src/relay.ts Adds relay response-header delivery for HTTP and WebSocket transports, including retry-state cleanup for websocket attempts.
packages/opencode/src/index.ts Connects relay and direct response headers to quota harvesting, profile hydration, sidebar updates, and command display paths.
packages/core/src/accounts.ts Adds profile persistence, header-aware quota merging, extra-usage fields, and account state lock handling.
packages/core/src/oauth-profile.ts Introduces OAuth profile fetching, token matching, freshness checks, and tier-label formatting.
packages/core/src/quota-headers.ts Normalizes Anthropic unified quota headers into quota snapshots for passive quota updates.

Reviews (11): Last reviewed commit: "feat: harvest quota headers on relay tra..." | Re-trigger Greptile

Context used (5)

Comment thread packages/core/src/relay.ts Outdated
Comment thread packages/opencode/src/index.ts

@cubic-dev-ai cubic-dev-ai 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.

All reported issues were addressed across 30 files

Confidence score: 5/5

  • Safe to merge after the addressed issues were fixed.
Architecture diagram
sequenceDiagram
    participant Plugin as Plugin (onResponseHeaders)
    participant Relay as sendViaRelay
    participant Worker as Relay Worker
    participant Anthropic as Anthropic API
    participant QuotaMgr as QuotaManager
    participant Storage as AccountStorage
    participant Sidebar as Sidebar State

    Note over Plugin,Sidebar: NEW: Harvest quota headers from relay transports

    Plugin->>Relay: request() with onResponseHeaders callback
    Note over Relay: transports: HTTP, WebSocket, or fallback

    alt HTTP relay
        Relay->>Worker: POST /v1/messages (patched body)
        Worker->>Anthropic: forward request
        Anthropic-->>Worker: response with upstream headers (e.g. anthropic-ratelimit-unified-*)
        Worker-->>Relay: relay response (status, headers, body)
        Note over Relay: CHANGED: onResponseHeaders fired with relay response headers
        Relay->>Relay: onResponseHeaders(response.headers)
        Relay-->>Plugin: Response (via sendViaRelay return)
    else WebSocket relay
        Relay->>Worker: WebSocket connect + send payload
        Worker->>Anthropic: forward request
        Anthropic-->>Worker: upstream response headers
        Worker-->>Relay: response_start control message (status, headers)
        Note over Relay: CHANGED: onResponseHeaders fired once from response_start headers
        Relay->>Relay: onResponseHeaders(new Headers(message.headers))
        Worker-->>Relay: stream body (SSE)
        Relay-->>Plugin: Response (optimistic for early resolution)
    else relay-to-direct fallback
        Note over Relay: No onResponseHeaders call (direct path harvests independently)
        Relay->>Plugin: fallback() Response (direct Anthropic request)
        Note over Plugin: Existing direct harvest path continues unchanged
    end

    Note over Plugin,Storage: Header harvest (fired inside onResponseHeaders callback)

    Plugin->>Plugin: isQuotaBearingHeaderFrame(headers) ?
    alt Has quota headers
        Plugin->>Plugin: normalizeQuotaHeaders(headers)
        Note over Plugin: Merge header snapshot with existing poll state<br/>Preserve scoped, extraUsage, poll bindingWindow
        Plugin->>QuotaMgr: pushMainFromHeaders(accessToken, incoming)
        QuotaMgr->>QuotaMgr: mergeHeaderQuotaSnapshot(existing, incoming)
        QuotaMgr->>Storage: persist quota (mainQuota)
        Plugin->>Storage: persist quota for fallback accounts (if applicable)
        Plugin->>Sidebar: refresh sidebar state
    else No quota headers
        Plugin->>Plugin: trace skip (no harvest)
    end

    Note over Plugin,Storage: Profile hydration (for tier labels)

    opt On boot or /claude-quota command
        Plugin->>Plugin: ensureProfilesForQuotaDisplay()
        Plugin->>Anthropic: GET /api/oauth/profile (with access token)
        Anthropic-->>Plugin: profile data (tier, orgType)
        Plugin->>Storage: persist mainProfile / account profiles
        Plugin->>Sidebar: update tierLabel
    end
Loading

Tip: cubic can generate docs of your entire codebase and keep them up to date. Try it here.

Re-trigger cubic

Comment thread packages/core/src/quota-manager.ts
Comment thread packages/core/src/quota-headers.ts Outdated
Comment thread packages/core/src/accounts.ts
Comment thread packages/opencode/src/index.ts
Comment thread packages/opencode/src/index.ts Outdated
Comment thread packages/core/src/quotas.ts Outdated
Comment thread packages/core/src/oauth-profile.ts Outdated
Comment thread docs/quota-surfaces.md Outdated
Comment thread docs/quota-surfaces.md Outdated
@iceteaSA
iceteaSA force-pushed the feat/relay-header-harvest branch from bf8749c to e3f231c Compare July 17, 2026 21:56
@ualtinok

Copy link
Copy Markdown
Contributor

The relay propagation design is directionally sound. I tested it stacked on #129 and the current sticky-routing work across HTTP relay, optimistic WebSocket relay, relay-to-direct fallback, typecheck, build, unit tests, and E2E.

This should remain blocked on #129 because it depends on that quota-harvesting API and inherits the unresolved persistence/test-isolation issues there.

Please also make onUpstreamHeaders idempotent per logical request. The WebSocket recovery path can safely reconnect and send a new full_sync after response_start but before downstream bytes; the callback is currently invoked for each accepted response_start. The public contract says it fires exactly once, so either enforce that with per-request state and a recovery regression test, or explicitly define and test last-response-wins semantics.

Once #129 is corrected and rebased, I will re-check this stack together.

@iceteaSA
iceteaSA force-pushed the feat/relay-header-harvest branch 2 times, most recently from 25f0b82 to a58b9c7 Compare July 18, 2026 20:10
@iceteaSA

Copy link
Copy Markdown
Contributor Author

Done in a58b9c7, rebased onto the corrected #129 (9f1b85a).

onUpstreamHeaders is now idempotent per logical request with explicit last-accepted-response-wins semantics: headers are staged at response_start and fire exactly once when the response commits (first downstream bytes or completion). A pre-byte recovery that reconnects and resends full_sync replaces the staged headers with the retry's frame — the retry's response is the one that actually streams to the consumer, and since the unified headers are cumulative account state it is also the fresher quota evidence. Duplicate response_start frames within one attempt are ignored. The doc comment on the callback states these semantics.

Regression covers your exact sequence: response_start (headers A) → drop before downstream bytes → reconnect + full_sync resend → response_start (headers B) → body streams → callback fired exactly once with B (red on the previous code, which fired A then B), plus a duplicate-frame guard test.

One declared edge: if the connection drops after response_start and the request cannot retry (attempt budget exhausted), the staged headers are discarded rather than delivered for a failed request — no response was delivered, and the next successful request re-harvests within seconds.

@iceteaSA
iceteaSA force-pushed the feat/relay-header-harvest branch 2 times, most recently from e856931 to 2f07769 Compare July 19, 2026 14:37
@ualtinok

Copy link
Copy Markdown
Contributor

The updated callback semantics resolve my previous finding. Headers are staged until the response commits, a pre-byte recovery replaces them with the retry response's headers, duplicate response_start frames are ignored, and the callback fires once with the final accepted response.

I validated the #129/#130 stack with two consecutive full unit-suite runs (1,032 passing each), typecheck, and an isolated E2E run (23 passing). I found no additional #130-specific blocker.

Please keep this stacked on #129 and rebase it after #129 replaces the broad fallback-account profile save with the token-checked profile-only mutation described there. I will revalidate both together after that update.

@iceteaSA
iceteaSA force-pushed the feat/relay-header-harvest branch from 2f07769 to 62118a4 Compare July 19, 2026 17:55
@iceteaSA

Copy link
Copy Markdown
Contributor Author

Restacked on #129's token-fenced profile fix (dadd240) as requested — new head 62118a4, still the two-commit stack shape (squashed #129 base + the relay harvest commit). Gates on the restack: opencode 923, e2e 23, typecheck/lint clean.

@ualtinok

Copy link
Copy Markdown
Contributor

No new relay-specific blocker was found in the latest review. PR #128 has now merged, and #129 has one remaining conditional profile-clear race documented in the new review comment there.

Please restack this branch after #129 fixes that race and rebases onto current main. I will rerun the combined #129/#130 gate on the resulting heads.

@iceteaSA
iceteaSA force-pushed the feat/relay-header-harvest branch from 62118a4 to e4e2a55 Compare July 19, 2026 22:56
@iceteaSA

Copy link
Copy Markdown
Contributor Author

Restacked on the corrected #129 head (which is now rebased onto current main). No relay-specific changes needed — the onUpstreamHeaders idempotency fix from the prior round is intact.

Gates green on the restacked head: typecheck · opencode 955 · core 58 · pi 53 · e2e 23 · lint clean. Ready for the combined #129/#130 gate.

Comment thread packages/opencode/src/index.ts Outdated
Harvest anthropic-ratelimit-unified-* headers in the fetch/relay path and
normalize them into the shared OAuthQuotaSnapshot alongside poll data; surface
account tier/credits/failover metadata in the expanded sidebar. Token-fenced
profile-only persistence with delayed-clear protection. Profile persistence is
detached from the /claude-quota and /claude-account display path (render
immediately, persist fire-and-forget) so an optional tier-label save can never
block a command on the account write lock.
@iceteaSA
iceteaSA force-pushed the feat/relay-header-harvest branch from e4e2a55 to fa083f7 Compare July 19, 2026 23:33
@iceteaSA

Copy link
Copy Markdown
Contributor Author

Resolved the profile-persistence-blocks-display finding on the #129 base (this branch is restacked on it). ensureProfilesForQuotaDisplay now updates the in-memory display profile immediately and runs persistProfileStateBestEffort fire-and-forget (void … .catch()) at all three sites — the optional tier-label save can no longer block /claude-quota or /claude-account status on the account write lock. The persist stays fully token-fenced internally; only the await was removed. Regression holds the config-write lock and proves the command renders before the persist completes, then confirms the persist still lands after lock release.

Restacked head gated green: typecheck · opencode 956 · e2e 23 · lint clean.

Relay responses now feed the passive quota-header harvester: HTTP relay
responses expose upstream headers directly, and the WebSocket relay
delivers genuine upstream headers from response_start control messages
via an onResponseHeaders callback in sendViaRelay. Synthetic optimistic
response headers are never harvested; relay-to-direct fallback keeps
header harvesting on the direct path.

The WebSocket callback is idempotent per logical request with
last-accepted-response wins: headers are staged at response_start and
fire exactly once when the response commits (first downstream bytes or
completion). A pre-byte recovery that resends full_sync replaces the
staged headers with the retry's frame — the response that actually
streams is also the fresher quota evidence — and duplicate
response_start frames within one attempt are ignored.
@iceteaSA
iceteaSA force-pushed the feat/relay-header-harvest branch from fa083f7 to 9a1a72f Compare July 20, 2026 05:15
@ualtinok

Copy link
Copy Markdown
Contributor

Integrated the relay-only commit onto main as 1c141720 after #129 was squash-merged. The original stacked branch became structurally conflicting because its #129 parent commit had a different SHA, so applying only the relay commit kept history clean while preserving the original author.

Post-integration validation passed: typecheck, lint, format, 956 OpenCode/core tests, 53 Pi tests, 23 end-to-end tests, and the packed TUI smoke test. Closing this PR as integrated.

@ualtinok ualtinok closed this Jul 21, 2026
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.

2 participants