fix(server): fail in-flight bindings on same-token reconnect (#365) (#386)

* fix(server): fail in-flight bindings on same-token reconnect (#365)

When a daemon reconnects with the same CLIENT_TOKEN, `registerAgent`
closes the old WS and immediately swaps the `agents` map to the new
connection. The old socket's `close` handler eventually calls
`unregisterAgent(agentId, oldWs)`, but by then the map points at the
new conn, so the ws-identity guard (`existing.ws !== ws`) makes it a
no-op — and the displaced connection's in-flight task bindings are
never failed. Those tasks' HTTP streams hang forever, since the bound
sink's `finish()` (which closes `AsyncEventQueue.iterate()`) is never
called, and the new daemon process never knew the old taskId so it
can't complete them either.

Extract the binding-fail loop from `unregisterAgent` into a reusable
`failBindingsForAgent` helper and invoke it from the reconnect
replacement branch *before* the map swap (while the bindings still
belong to the old conn). Superseded tasks now receive a terminal
`failed` status with a `superseded` error code; the normal disconnect
path keeps its existing `disconnected` code and message verbatim.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* test(server): strengthen #365 reconnect test (multi-binding + messageId)

Per code review: assert the per-path messageId suffix
(`-superseded`) that failBindingsForAgent wires through, and cover
the multiple-in-flight-binding case so a future early-return in the
fail loop regresses loudly.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: longfin

packages/server/src/registry.test.ts

@@ -401,6 +401,111 @@ test('registerAgent emits client_collision and closes the prior ws with the desc
   assert.equal(parsed.previousConnectedAt, 1000);
 });
 
+test('same-token reconnect fails the displaced connection in-flight bindings (issue #365)', () => {
+  // A second daemon authenticating with the same CLIENT_TOKEN replaces the
+  // incumbent. The old connection's in-flight task must receive a terminal
+  // `failed` status and have its sink finished + binding dropped — otherwise
+  // the task's HTTP stream hangs forever (the new daemon is a separate process
+  // that never knew the old taskId, so it can't complete it either).
+  const registry = new Registry();
+  const oldWs = makeWs();
+  const base = {
+    agentId: 'a1',
+    clientId: 'c1',
+    ownerPrincipal: 'eth:0x0',
+    agentCard: makeCard(false),
+    allowedCallers: [],
+    connectedAt: 0,
+  };
+  registry.registerAgent({ ...base, ws: oldWs });
+
+  // Two in-flight tasks on the displaced connection — both must be failed, so
+  // a future reintroduction of an early-return in the loop regresses loudly.
+  const statuses: Array<{
+    status: { state?: string; message?: { messageId?: string; metadata?: unknown } };
+  }> = [];
+  const finished = new Set<string>();
+  for (const taskId of ['t-live', 't-live-2']) {
+    registry.bindTask({
+      agentId: 'a1',
+      taskId,
+      contextId: `ctx-${taskId}`,
+      requestedExtensions: [OPENAI_COMPAT_EXTENSION_URI],
+      sink: {
+        pushStatus: (event) => statuses.push(event),
+        pushArtifact: () => undefined,
+        finish: () => {
+          finished.add(taskId);
+        },
+      },
+    });
+  }
+
+  // Same token (clientId) reconnects on a fresh socket → replacement branch.
+  registry.registerAgent({ ...base, ws: makeWs(), connectedAt: 1 });
+
+  assert.deepEqual([...finished].sort(), ['t-live', 't-live-2'], 'every displaced sink must be finished');
+  assert.equal(registry.getBinding('t-live'), undefined, 'binding must be dropped');
+  assert.equal(registry.getBinding('t-live-2'), undefined, 'binding must be dropped');
+  assert.equal(statuses[0]?.status.state, 'failed');
+  // messageId carries the per-path suffix wired through failBindingsForAgent —
+  // the disconnect path uses `-disc`, this reconnect path uses `-superseded`.
+  assert.equal(statuses[0]?.status.message?.messageId, 't-live-superseded');
+  assert.deepEqual(statuses[0]?.status.message?.metadata, {
+    [OPENAI_COMPAT_EXTENSION_URI]: {
+      terminal_error: {
+        code: 'superseded',
+        message: 'superseded by a reconnect from the same client token',
+      },
+    },
+    error: {
+      code: 'superseded',
+      message: 'superseded by a reconnect from the same client token',
+    },
+  });
+});
+
+test('a late close from the superseded socket does not double-fail the new connection bindings (issue #365)', () => {
+  // After the reconnect path has already terminated the old connection's
+  // bindings, the old socket's close handler eventually fires
+  // unregisterAgent(agentId, oldWs). The ws-identity guard must make that a
+  // no-op so it can't touch a task the *new* connection has since bound.
+  const registry = new Registry();
+  const oldWs = makeWs();
+  const base = {
+    agentId: 'a1',
+    clientId: 'c1',
+    ownerPrincipal: 'eth:0x0',
+    agentCard: makeCard(false),
+    allowedCallers: [],
+    connectedAt: 0,
+  };
+  registry.registerAgent({ ...base, ws: oldWs });
+  const newWs = makeWs();
+  registry.registerAgent({ ...base, ws: newWs, connectedAt: 1 });
+
+  // New connection binds a fresh task after the swap.
+  let finished = false;
+  registry.bindTask({
+    agentId: 'a1',
+    taskId: 't-new',
+    contextId: 'ctx-new',
+    sink: {
+      pushStatus: () => undefined,
+      pushArtifact: () => undefined,
+      finish: () => {
+        finished = true;
+      },
+    },
+  });
+
+  // Late close from the displaced socket — must not disturb t-new.
+  registry.unregisterAgent('a1', oldWs);
+
+  assert.equal(finished, false, 'new connection binding must survive a stale unregister');
+  assert.ok(registry.getBinding('t-new'), 'new binding must still be present');
+});
+
 test('disconnectClient returns 0 when no agents are bound to the client', () => {
   const registry = new Registry();
   registry.registerAgent({

packages/server/src/registry.ts

@@ -88,6 +88,18 @@ export class Registry {
         // the foreground log can recognize the duplicate-token scenario
         // without cross-referencing close codes.
         existing.ws.close(4009, 'another client with the same token connected');
+        // Fail the displaced connection's in-flight tasks *before* swapping the
+        // map. The old socket's close fires asynchronously and, by the time its
+        // unregisterAgent runs, the map already points at `conn` — so the
+        // ws-identity guard makes it a no-op and it can't clean these up. Doing
+        // it here (while the bindings still belong to the old conn) is the only
+        // point that sees them. The new conn hasn't bound any tasks yet, so
+        // this only terminates the superseded ones.
+        this.failBindingsForAgent(conn.agentId, {
+          code: 'superseded',
+          message: 'superseded by a reconnect from the same client token',
+          messageIdSuffix: 'superseded',
+        });
         this.agents.set(conn.agentId, conn);
         this.notifyAgentChange(conn.agentId);
         return { ok: true };
@@ -139,6 +151,27 @@ export class Registry {
     if (!existing || existing.ws !== ws) return;
     this.agents.delete(agentId);
     this.notifyAgentChange(agentId);
+    this.failBindingsForAgent(agentId, {
+      code: 'disconnected',
+      message: 'client disconnected mid-task',
+      messageIdSuffix: 'disc',
+    });
+  }
+
+  // Push a terminal `failed` status to every in-flight task bound to this
+  // agent, finish its sink, and drop the binding. Both the normal disconnect
+  // (unregisterAgent) and the same-token reconnect replacement (registerAgent)
+  // funnel through here so an in-flight task is never left without a terminal
+  // event — the HTTP stream for that task would otherwise hang forever waiting
+  // on AsyncEventQueue.iterate(), since the bound sink's finish() is what
+  // closes the iterator. The reconnect path in particular MUST call this
+  // explicitly before swapping the agents map: once the map points at the new
+  // conn, the old socket's late close handler hits the ws-identity guard in
+  // unregisterAgent and early-returns, so it can no longer fail these bindings.
+  private failBindingsForAgent(
+    agentId: string,
+    error: { code: string; message: string; messageIdSuffix: string },
+  ): void {
     for (const binding of [...this.bindings.values()]) {
       if (binding.agentId !== agentId) continue;
       binding.sink.pushStatus({
@@ -154,13 +187,13 @@ export class Registry {
           state: 'failed' as never,
           timestamp: new Date().toISOString(),
           message: {
-            messageId: `${binding.taskId}-disc`,
+            messageId: `${binding.taskId}-${error.messageIdSuffix}`,
             role: 'agent',
-            parts: [{ text: 'client disconnected mid-task' }],
+            parts: [{ text: error.message }],
             ...terminalErrorMessageFields(
               {
-                code: 'disconnected',
-                message: 'client disconnected mid-task',
+                code: error.code,
+                message: error.message,
               },
               binding.requestedExtensions,
             ),