Fix agent wake claim token races#4616
Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #4616 +/- ##
===========================================
+ Coverage 34.72% 59.36% +24.63%
===========================================
Files 233 385 +152
Lines 19843 42845 +23002
Branches 7097 12302 +5205
===========================================
+ Hits 6890 25433 +18543
- Misses 12919 17336 +4417
- Partials 34 76 +42
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Electric Agents Mobile BuildLocal mobile checks ran for commit The EAS Android preview build was skipped because the |
samwillis
left a comment
There was a problem hiding this comment.
Review
Interactive review with GPT-5.5.
Blocking Finding
The token store now correctly allows a single consumer to own multiple stream-scoped write tokens, but the surrounding callback routing still stores only one primaryStream per (tenantId, consumerId) in consumer_callbacks. If the same consumer claims stream A, then stream B before A sends done, A's later done will look up B's primaryStream, check/clear B's in-memory token, and potentially mark B idle while leaving A's token/state stale. That can recreate the invalid-token/stale-state race this PR is trying to remove.
I think the fix is to make the release path use the released claim identity rather than the mutable callback row's current primaryStream: either key callback state by (consumerId, epoch) / stream, or use materializeReleasedClaim()'s returned claim.stream_path and claim.entity_url for token ownership, token clearing, and status updates.
Verification
I attempted the focused PR tests, but the worktree shell is on Node v18.20.5 and pnpm refused to run because the repo requires >=20.19.0 || >=22.12.0.
| export class ClaimWriteTokenStore { | ||
| private readonly claimsByStream = new Map<string, ActiveClaimWriteToken>() | ||
| private readonly streamByConsumer = new Map<string, string>() | ||
| private readonly streamKeysByConsumer = new Map<string, Set<string>>() |
There was a problem hiding this comment.
This fixes the in-memory side of multi-stream ownership, but the callback release path still appears to be keyed through a single mutable consumer_callbacks.primaryStream per consumer. If the same consumer claims A then B, A's later done can resolve to B's stream and clear/update the wrong active token/status. The release path should use the released (consumerId, epoch) claim's stream_path/entity_url, or callback state needs to be scoped by claim/stream too.
samwillis
left a comment
There was a problem hiding this comment.
Follow-up interactive review with GPT-5.5.
The latest update addresses my previous blocking concern: the wake done path now uses the released claim's stream identity instead of the mutable callback row's current primary stream, and the new regression coverage exercises the multi-stream single-consumer case.
Verification run locally in the review worktree:
@electric-ax/agents-runtime:test/pull-wake-runner.test.tsandtest/process-wake.test.tspassed, 115 tests@electric-ax/agents-server:test/claim-write-token-store.test.tsandtest/subscription-webhooks-routing.test.tspassed, 19 tests
Approved.
Fixes agent wake handling so concurrent desktop sessions no longer invalidate each other's claim write tokens, and same-stream duplicate wakes are retried after the active wake drains instead of being dropped.
This should stop the recent local desktop stall pattern where an unrelated wake or SIGINT caused another active session to start receiving
Invalid write token401s. It also narrows failed-parent child cleanup so a parent handler failure detaches the parent wake registration instead of sendingSIGKILLto a child that may already be running.Root Cause
The claim write token store treated a consumer as owning at most one stream. When the same desktop runner/consumer claimed a wake for a second agent stream, it deleted the first stream's token, so the still-active wake failed later writes with
UNAUTHORIZED: Invalid write token.A second race existed for duplicate same-stream wake events. The runner could claim another wake for a stream while the previous wake was still active, minting a replacement token for that same stream. The first fix avoided concurrent same-stream claims, but review caught that a simple local skip could strand pending work if the skipped event was the only retry signal.
Separately, parent handler failure cleanup used
deleteEntity()for children spawned before the failure. That client method actually sends semanticSIGKILL, which is too broad for rolling back a failed parent manifest entry.Approach
clearConsumer()still clears all streams owned by that consumer.SIGINTof an active or just-starting run as a shutdown request, so the wake closes instead of returning to idle and immediately processing fresh input in the aborted invocation.Key Invariants
SIGINTaborts the current handler and closes that wake; later work runs in a later wake.Non-goals
SIGKILLAPIs; it stops using them for failed-parent child cleanup.Verification
Files Changed
packages/agents-server/src/claim-write-token-store.ts: allow one consumer to hold multiple stream-scoped claim tokens.packages/agents-runtime/src/create-handler.ts: expose stream-path active wake tracking to the pull runner.packages/agents-runtime/src/pull-wake-runner.ts: prevent concurrent same-stream claims and defer/retry skipped active-stream wake events.packages/agents-runtime/src/process-wake.ts: close the wake after active-runSIGINTaborts, and detach child wake registrations on parent handler failure without killing children.packages/agents-runtime/src/runtime-server-client.tsandpackages/agents-server/src/routing/internal-router.ts: add a narrow internal wake unregister call..changeset/fix-agent-wake-claim-token-races.md: patch changeset for the affected agents packages.