Skip to content

Fix agent wake claim token races#4616

Merged
KyleAMathews merged 6 commits into
mainfrom
fix/agents-wake-claim-token-race
Jun 18, 2026
Merged

Fix agent wake claim token races#4616
KyleAMathews merged 6 commits into
mainfrom
fix/agents-wake-claim-token-race

Conversation

@KyleAMathews

@KyleAMathews KyleAMathews commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

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 token 401s. It also narrows failed-parent child cleanup so a parent handler failure detaches the parent wake registration instead of sending SIGKILL to 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 semantic SIGKILL, which is too broad for rolling back a failed parent manifest entry.

Approach

  • Scope claim write tokens by stream, while keeping a reverse index so clearConsumer() still clears all streams owned by that consumer.
  • Track active wake stream paths in the runtime router, and let the pull wake runner ask whether a stream is already in flight.
  • Prevent concurrent claims for the same stream while a claim or wake is active.
  • Defer locally skipped same-stream wake events and retry them after the active stream wake drains.
  • Treat SIGINT of 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.
  • Add a narrow internal wake unregister endpoint and use it to detach failed-parent child wake registrations without signaling the child entity.

Key Invariants

  • A consumer may hold claim write tokens for multiple streams at the same time.
  • A stream has at most one active claim token at a time.
  • A duplicate wake-stream event must not mint a replacement token while that stream is active.
  • A locally skipped wake must be retried after the active wake drains so pending work is not stranded.
  • SIGINT aborts the current handler and closes that wake; later work runs in a later wake.
  • Failed parent cleanup may detach wake registrations created by that failed run, but must not terminate the child entity.

Non-goals

  • This does not change Durable Streams lease/heartbeat semantics.
  • This does not rewrite the runtime as per-agent XState actors; it adds guardrails to the existing promise-based wake dispatch model.
  • This does not remove the explicit public/legacy SIGKILL APIs; it stops using them for failed-parent child cleanup.

Verification

pnpm --filter @electric-ax/agents-runtime exec vitest run test/pull-wake-runner.test.ts test/process-wake.test.ts --pool-options.threads.maxThreads=2
pnpm --filter @electric-ax/agents-server exec vitest run test/claim-write-token-store.test.ts test/subscription-webhooks-routing.test.ts --maxWorkers=2
pnpm --filter @electric-ax/agents-runtime typecheck
pnpm --filter @electric-ax/agents-server typecheck
GITHUB_BASE_REF=main node scripts/check-changeset.mjs

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-run SIGINT aborts, and detach child wake registrations on parent handler failure without killing children.
  • packages/agents-runtime/src/runtime-server-client.ts and packages/agents-server/src/routing/internal-router.ts: add a narrow internal wake unregister call.
  • Runtime and server tests: cover multi-stream consumers, same-stream duplicate deferral, retry after active wake drains, SIGINT wake closure, and failed-parent child wake detach.
  • .changeset/fix-agent-wake-claim-token-races.md: patch changeset for the affected agents packages.

@github-actions

github-actions Bot commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

Electric Agents Desktop Builds

Build artifacts for commit 135ba33.

Platform Status Artifact
macOS Apple Silicon Passed DMG
macOS Intel Passed DMG
Windows x64 Passed Installer
Linux x64 Passed AppImage / deb

Workflow run

@KyleAMathews KyleAMathews requested a review from samwillis June 17, 2026 16:49
@codecov

codecov Bot commented Jun 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 83.59375% with 21 lines in your changes missing coverage. Please review.
✅ Project coverage is 59.36%. Comparing base (aa6ca85) to head (135ba33).
⚠️ Report is 8 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines Patch % Lines
packages/agents-runtime/src/process-wake.ts 40.00% 9 Missing ⚠️
packages/agents-runtime/src/pull-wake-runner.ts 91.22% 5 Missing ⚠️
...ckages/agents-runtime/src/runtime-server-client.ts 71.42% 4 Missing ⚠️
packages/agents-runtime/src/create-handler.ts 80.00% 1 Missing ⚠️
packages/agents-runtime/src/setup-context.ts 0.00% 1 Missing ⚠️
...kages/agents-server/src/routing/internal-router.ts 93.33% 1 Missing ⚠️
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     
Flag Coverage Δ
packages/agents 72.96% <ø> (?)
packages/agents-mcp 77.70% <ø> (?)
packages/agents-mobile 80.67% <ø> (ø)
packages/agents-runtime 83.45% <78.26%> (?)
packages/agents-server 75.46% <97.22%> (+0.17%) ⬆️
packages/agents-server-ui 7.51% <ø> (-0.03%) ⬇️
packages/electric-ax 47.62% <ø> (?)
packages/experimental 87.73% <ø> (?)
packages/react-hooks 86.48% <ø> (?)
packages/start 82.83% <ø> (?)
packages/typescript-client 91.71% <ø> (?)
packages/y-electric 56.05% <ø> (?)
typescript 59.36% <83.59%> (+24.63%) ⬆️
unit-tests 59.36% <83.59%> (+24.63%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@github-actions

github-actions Bot commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

Electric Agents Mobile Build

Local mobile checks ran for commit 135ba33.

The EAS Android preview build was skipped because the mobile-eas-build label is not present.
Add the mobile-eas-build label to this PR to produce an installable preview build.

Workflow run

@samwillis samwillis left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>>()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 samwillis left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.ts and test/process-wake.test.ts passed, 115 tests
  • @electric-ax/agents-server: test/claim-write-token-store.test.ts and test/subscription-webhooks-routing.test.ts passed, 19 tests

Approved.

@KyleAMathews KyleAMathews merged commit 37b9922 into main Jun 18, 2026
71 checks passed
@KyleAMathews KyleAMathews deleted the fix/agents-wake-claim-token-race branch June 18, 2026 16:11
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