engine: export node-control provider-disconnect + fix dropped-provider aggregate leak (#244)

willwashburn · claude · web-flow · commit 8b3ec7101bf3 · 2026-07-08T17:11:45.000-04:00
* engine: export handleProviderDisconnect + markNodeOffline from node-control Out-of-process node-control socket owners (the relaycast-cloud NodeDO) own their sockets but not the DB. On socket close they must flip the right liveness — one provider offline while the node stays online, or the whole node offline when the last connection drops — and that logic already lives in `engine/node.js` (`handleProviderDisconnect`, `markNodeOffline`) but was unreachable from the `@relaycast/engine/node-control` entrypoint, forcing the host to hand-roll a SQL mirror. Re-export both so the host delegates instead. Both encapsulate the full effect (provider row + agents + node aggregate + invocation reschedule); a conformance test asserts the public exports produce the provider-scoped, last-connection, and node-wide offline states. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test: optional-chain stack teardown so a failed setup surfaces its own error Addresses review feedback on #244. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix: zero a dropped provider's active_agents so the node aggregate is correct `markProviderOffline` and `markNodeOffline` left `node_providers.active_agents` at its last value, but `recomputeNodeAggregate` sums active_agents across ALL providers (online and offline). So after a provider's socket dropped with others still connected, `nodes.active_agents` kept counting the dropped provider's now- offline agents until it reconnected. Reset it to 0 on offline; a reconnect / heartbeat repopulates it from the provider's own report. Regression test in nodeControlExports asserts the node aggregate drops the dropped provider's count and restores it on the next heartbeat, and that markNodeOffline zeros every provider's count. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * test: exercise the real reconnect path in the aggregate restore assertion Close the dropped provider's socket (driving the disconnect through the adapter's own close path) and register a fresh socket with a new instance for the restore, instead of re-heartbeating the stale handle — matches how a provider actually reconnects. Addresses review feedback on #244. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
diff --git a/packages/engine/src/__tests__/conformance/nodeControlExports.test.ts b/packages/engine/src/__tests__/conformance/nodeControlExports.test.ts
@@ -0,0 +1,176 @@
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { and, eq } from 'drizzle-orm';
+import { handleProviderDisconnect, markNodeOffline } from '../../node-control.js';
+import { makeNodeStack, createWorkspace, FakeSocket, type TestStack } from './harness.js';
+import { nodeProviders, nodes } from '../../db/schema.js';
+
+// The provider-disconnect lifecycle is exported from @relaycast/engine/node-control
+// so an out-of-process socket owner — the relaycast-cloud NodeDO — drives it on
+// socket close instead of hand-rolling the SQL (the mirror the node-providers spec
+// is removing). These assert the public exports produce the same DB liveness the
+// in-process adapter's own close path produces.
+describe('node-control provider-disconnect exports', () => {
+  let stack: TestStack;
+  beforeEach(() => { stack = makeNodeStack({ ttlMs: 60_000 }); });
+  // Optional-chain so a failed beforeEach (undefined stack) surfaces its real
+  // error instead of a TypeError from teardown.
+  afterEach(() => stack?.close());
+
+  const db = () => stack.runtime.handle.db;
+
+  async function enrollNode(ws: { workspaceKey: string }, nodeId: string, name: string) {
+    const res = await stack.app.request('/v1/nodes', {
+      method: 'POST',
+      headers: { 'content-type': 'application/json', authorization: `Bearer ${ws.workspaceKey}` },
+      body: JSON.stringify({ node_id: nodeId, name, role: 'broker', capabilities: [], max_agents: 4, tags: ['test'], version: 'v0' }),
+    });
+    expect(res.status).toBe(201);
+  }
+
+  async function attachProvider(
+    workspaceId: string,
+    nodeId: string,
+    nodeName: string,
+    providerName: string,
+    capability: string,
+    activeAgents = 0,
+    instanceId = `${providerName}-i1`,
+  ) {
+    const provider = { name: providerName, instance_id: instanceId };
+    const sock = new FakeSocket();
+    const handle = stack.runtime.realtime.attachNodeSocket(workspaceId, nodeId, sock);
+    await handle.handleMessage(JSON.stringify({
+      v: 1,
+      id: `reg-${instanceId}`,
+      type: 'node.register',
+      name: nodeName,
+      node_id: nodeId,
+      provider,
+      capabilities: [{ name: capability, kind: 'action' }],
+      max_agents: 4,
+      tags: ['test'],
+      version: 'v1',
+      resume_cursor: null,
+    }));
+    await heartbeat(handle, providerName, activeAgents, instanceId);
+    return { sock, handle };
+  }
+
+  function heartbeat(
+    handle: { handleMessage(raw: string): Promise<void> },
+    providerName: string,
+    activeAgents: number,
+    instanceId = `${providerName}-i1`,
+  ) {
+    return handle.handleMessage(JSON.stringify({
+      v: 1,
+      type: 'node.heartbeat',
+      provider: { name: providerName, instance_id: instanceId },
+      load: 0,
+      active_agents: activeAgents,
+      handlers_live: true,
+    }));
+  }
+
+  function nodeActiveAgents(workspaceId: string, nodeId: string) {
+    return db()
+      .select({ activeAgents: nodes.activeAgents })
+      .from(nodes)
+      .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId)))
+      .then((rows) => rows[0]?.activeAgents);
+  }
+
+  function providerActiveAgents(workspaceId: string, nodeId: string, name: string) {
+    return db()
+      .select({ activeAgents: nodeProviders.activeAgents })
+      .from(nodeProviders)
+      .where(and(eq(nodeProviders.workspaceId, workspaceId), eq(nodeProviders.nodeId, nodeId), eq(nodeProviders.name, name)))
+      .then((rows) => rows[0]?.activeAgents);
+  }
+
+  function providerStatuses(workspaceId: string, nodeId: string) {
+    return db()
+      .select({ name: nodeProviders.name, status: nodeProviders.status })
+      .from(nodeProviders)
+      .where(and(eq(nodeProviders.workspaceId, workspaceId), eq(nodeProviders.nodeId, nodeId)))
+      .then((rows) => Object.fromEntries(rows.map((r) => [r.name, r.status])));
+  }
+
+  function nodeStatus(workspaceId: string, nodeId: string) {
+    return db()
+      .select({ status: nodes.status })
+      .from(nodes)
+      .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId)))
+      .then((rows) => rows[0]?.status);
+  }
+
+  it('handleProviderDisconnect (remaining connections) flips only that provider offline; node stays online', async () => {
+    const ws = await createWorkspace(stack.app, 'exp-provider-remaining');
+    await enrollNode(ws, 'node_a', 'alpha');
+    await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'py', 'run-etl');
+    await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'rb', 'build');
+
+    await handleProviderDisconnect(db(), stack.runtime.realtime, ws.workspaceId, 'node_a', 'py', true);
+
+    expect(await providerStatuses(ws.workspaceId, 'node_a')).toEqual({ py: 'offline', rb: 'online' });
+    expect(await nodeStatus(ws.workspaceId, 'node_a')).toBe('online');
+  });
+
+  it('handleProviderDisconnect (no remaining connections) marks the whole node offline', async () => {
+    const ws = await createWorkspace(stack.app, 'exp-provider-last');
+    await enrollNode(ws, 'node_a', 'alpha');
+    await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'py', 'run-etl');
+
+    await handleProviderDisconnect(db(), stack.runtime.realtime, ws.workspaceId, 'node_a', 'py', false);
+
+    expect(await providerStatuses(ws.workspaceId, 'node_a')).toEqual({ py: 'offline' });
+    expect(await nodeStatus(ws.workspaceId, 'node_a')).toBe('offline');
+  });
+
+  it('markNodeOffline flips the node and every provider offline', async () => {
+    const ws = await createWorkspace(stack.app, 'exp-node-offline');
+    await enrollNode(ws, 'node_a', 'alpha');
+    await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'py', 'run-etl');
+    await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'rb', 'build');
+
+    await markNodeOffline(db(), stack.runtime.realtime, ws.workspaceId, 'node_a');
+
+    expect(await providerStatuses(ws.workspaceId, 'node_a')).toEqual({ py: 'offline', rb: 'offline' });
+    expect(await nodeStatus(ws.workspaceId, 'node_a')).toBe('offline');
+  });
+
+  it('drops a disconnected provider from the node aggregate, and restores it on reconnect', async () => {
+    const ws = await createWorkspace(stack.app, 'exp-aggregate');
+    await enrollNode(ws, 'node_a', 'alpha');
+    const py = await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'py', 'run-etl', 2);
+    await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'rb', 'build', 3);
+    // The node aggregate sums both providers' active agents.
+    expect(await nodeActiveAgents(ws.workspaceId, 'node_a')).toBe(5);
+
+    // py's socket drops (others remain): its agents are gone, so the node
+    // aggregate must no longer count them — recomputeNodeAggregate would resurrect
+    // them if the provider row kept its stale activeAgents.
+    await py.handle.handleClose();
+    expect(await providerActiveAgents(ws.workspaceId, 'node_a', 'py')).toBe(0);
+    expect(await nodeActiveAgents(ws.workspaceId, 'node_a')).toBe(3);
+
+    // Symmetric restore through a real reconnect: a fresh socket registers py
+    // (new instance) and heartbeats, repopulating its count.
+    await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'py', 'run-etl', 2, 'py-i2');
+    expect(await providerActiveAgents(ws.workspaceId, 'node_a', 'py')).toBe(2);
+    expect(await nodeActiveAgents(ws.workspaceId, 'node_a')).toBe(5);
+  });
+
+  it('markNodeOffline zeros every provider active-agent count', async () => {
+    const ws = await createWorkspace(stack.app, 'exp-node-offline-aggregate');
+    await enrollNode(ws, 'node_a', 'alpha');
+    await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'py', 'run-etl', 2);
+    await attachProvider(ws.workspaceId, 'node_a', 'alpha', 'rb', 'build', 3);
+
+    await markNodeOffline(db(), stack.runtime.realtime, ws.workspaceId, 'node_a');
+
+    expect(await providerActiveAgents(ws.workspaceId, 'node_a', 'py')).toBe(0);
+    expect(await providerActiveAgents(ws.workspaceId, 'node_a', 'rb')).toBe(0);
+    expect(await nodeActiveAgents(ws.workspaceId, 'node_a')).toBe(0);
+  });
+});
diff --git a/packages/engine/src/engine/node.ts b/packages/engine/src/engine/node.ts
@@ -476,10 +476,11 @@ export async function markNodeOffline(
     .where(and(eq(nodes.workspaceId, workspaceId), eq(nodes.id, nodeId)));
 
   // Provider capability sets persist (offline nodes still show their manifest);
-  // only their liveness flips.
+  // only their liveness flips. Zero activeAgents too, so a later
+  // recomputeNodeAggregate never resurrects a dropped provider's agent count.
   await db
     .update(nodeProviders)
-    .set({ status: 'offline', handlersLive: false, load: 0, lastHeartbeatAt: new Date() })
+    .set({ status: 'offline', handlersLive: false, load: 0, activeAgents: 0, lastHeartbeatAt: new Date() })
     .where(and(eq(nodeProviders.workspaceId, workspaceId), eq(nodeProviders.nodeId, nodeId)));
 
   await db
diff --git a/packages/engine/src/engine/nodeProvider.ts b/packages/engine/src/engine/nodeProvider.ts
@@ -154,9 +154,12 @@ export async function heartbeatProvider(
 }
 
 export async function markProviderOffline(db: Db, workspaceId: string, nodeId: string, name: string): Promise<void> {
+  // A disconnected provider hosts no live agents; zero its activeAgents so
+  // recomputeNodeAggregate (which sums across all providers) stops counting them
+  // on the node. A reconnect/heartbeat repopulates it from the provider's report.
   await db
     .update(nodeProviders)
-    .set({ status: 'offline', handlersLive: false, load: 0, lastHeartbeatAt: new Date() })
+    .set({ status: 'offline', handlersLive: false, load: 0, activeAgents: 0, lastHeartbeatAt: new Date() })
     .where(and(
       eq(nodeProviders.workspaceId, workspaceId),
       eq(nodeProviders.nodeId, nodeId),
diff --git a/packages/engine/src/node-control.ts b/packages/engine/src/node-control.ts
@@ -1,5 +1,16 @@
 export {
   handleNodeControlMessage,
+  // Node-offline lifecycle for out-of-process socket owners (e.g. the
+  // relaycast-cloud NodeDO): when a provider's control socket drops, drive
+  // `handleProviderDisconnect` — it flips only that provider (and its agents)
+  // offline while the node stays online, or, when it was the node's last
+  // connection, marks the whole node offline. `markNodeOffline` is the node-wide
+  // marker for the other lifecycle triggers a socket owner controls (a liveness
+  // alarm lapse, an operator disconnect). Both encapsulate the provider-row,
+  // agent, aggregate, and invocation-reschedule effects so a host never
+  // hand-rolls them.
+  handleProviderDisconnect,
+  markNodeOffline,
   type HandleNodeControlMessageArgs,
   type NodeSocketLike,
 } from './engine/node.js';
