Multi-instance job control: route reconnects to the owner, cancel across machines (#4715)

Author: georgi

docs/websocket-api.md

@@ -38,6 +38,54 @@ same format. Every message is a map/object with at least a `type` field.
    reference client retries up to 10 times starting at 1 s).
 4. **Close** — call `socket.close()` or let the server close the connection.
 
+## Multi-Instance Deployments
+
+A run's replay buffer and its cancel/stream hooks live in the one server
+process executing it, so on a deployment with more than one instance both
+`reconnect_job` and `cancel_job` have to reach that process. Two environment
+variables drive this, and with neither set the whole mechanism is inert:
+
+- `NODETOOL_INSTANCE_ID` — this instance's identity.
+- `FLY_MACHINE_ID` — the fallback, set by Fly on every machine. It is also the
+  value `fly-replay: instance=<id>` addresses.
+
+The instance executing a run stamps its id on the job row (`runner_instance`).
+Two things follow.
+
+**Resuming lands on the owner.** A reconnecting client appends
+`?resume_job=<job_id>` to the handshake URL. If that job is non-terminal and
+owned by another instance, the server answers the upgrade with
+`fly-replay: instance=<owner>` instead of accepting it, and Fly's proxy
+re-issues the whole handshake there. A request the proxy already replayed
+(`fly-replay-src` present) is never replayed again, so this cannot ping-pong.
+The hint names one job — with runs in flight on several instances the rest
+reconnect wherever they land and fall back to `reconnect_job`'s persisted-row
+answer: the right status, without the replayed frames.
+
+The client retires the hint after two consecutive failed connects, so the third
+attempt goes out bare. Without that, a deploy that retires the owning machine
+while its row still reads `running` would have every reconnect replayed at a
+machine that no longer exists — and since the browser shares one socket across
+chat and every other consumer, all of them would stay dark. A successful
+connect resets the count.
+
+**Cancel travels through the row.** `cancel_job` for a run this process does
+not hold, whose row names a *different* instance, marks the row cancelled with
+a conditional update — only while it is still non-terminal, so it cannot
+overwrite the owner's own outcome. Every instance re-reads its own running
+runs on a timer (`NODETOOL_JOB_CANCEL_POLL_MS`, default 15000, `0` disables)
+and cancels any whose row now reads `cancelled`: one indexed query per tick,
+bounded by that instance's concurrency.
+
+The row is the only transport, so a cross-instance cancel takes up to a poll
+interval to land — the trade for having exactly one signal, the durable one. A
+cancel on the machine that *does* hold the run does not go through any of this;
+it reaches the session's hooks directly and is immediate.
+
+A row with no `runner_instance` (an HTTP, trigger, or MCP run — nothing holds a
+session for those anywhere) is left alone and still answers "Job not found or
+already completed".
+
 ## Client → Server Commands
 
 All client messages contain `command` and `data` fields.

packages/models/src/db.ts

@@ -271,6 +271,7 @@ const TABLE_COLUMNS: Record<string, Record<string, string>> = {
     suspension_metadata_json: "text",
     execution_strategy: "text",
     execution_id: "text",
+    runner_instance: "text",
     metadata_json: "text",
     created_at: "text",
     updated_at: "text"
@@ -648,6 +649,7 @@ function getCreateSchemaSql(): string {
       "suspension_metadata_json" text,
       "execution_strategy" text,
       "execution_id" text,
+      "runner_instance" text,
       "metadata_json" text,
       "created_at" text NOT NULL,
       "updated_at" text NOT NULL

packages/models/src/job.ts

@@ -4,7 +4,7 @@
  * Port of Python's `nodetool.models.job`.
  */
 
-import { eq, and, desc, lt } from "drizzle-orm";
+import { eq, and, desc, lt, inArray, notInArray } from "drizzle-orm";
 import { DBModel, createTimeOrderedUuid } from "./base-model.js";
 import { getDb } from "./db.js";
 import { jobs } from "./schema/jobs.js";
@@ -52,6 +52,12 @@ export class Job extends DBModel {
   declare suspension_metadata_json: Record<string, unknown> | null;
   declare execution_strategy: string | null;
   declare execution_id: string | null;
+  /**
+   * The server instance executing this run, when the deployment has more than
+   * one (see `getInstanceId` in the websocket package). Null means
+   * single-machine: nothing routes by it.
+   */
+  declare runner_instance: string | null;
   declare metadata_json: Record<string, unknown> | null;
   declare created_at: string;
   declare updated_at: string;
@@ -85,6 +91,7 @@ export class Job extends DBModel {
     this.suspension_metadata_json ??= null;
     this.execution_strategy ??= null;
     this.execution_id ??= null;
+    this.runner_instance ??= null;
     this.metadata_json ??= null;
     this.name ??= "";
   }
@@ -228,6 +235,55 @@ export class Job extends DBModel {
 
   // ── Static queries ───────────────────────────────────────────────
 
+  /**
+   * Cancel a run without reading it first.
+   *
+   * A cancel arriving on an instance that does not own the run races the
+   * owner's own terminal write. Loading the row, mutating it and calling
+   * `save()` would send a full-row upsert built from a snapshot taken before
+   * that race — resurrecting a `completed` job as `cancelled` and overwriting
+   * the cost and timestamps the owner had just written. This touches only the
+   * two columns a cancel owns, and only while the row is still active.
+   *
+   * Returns whether a row actually changed: false means the run was already
+   * terminal (or is not this user's), which is the caller's cue that there was
+   * nothing to cancel.
+   */
+  static async markCancelledIfActive(
+    jobId: string,
+    userId: string
+  ): Promise<boolean> {
+    const db = getDb();
+    const now = new Date().toISOString();
+    const updated = await db
+      .update(jobs)
+      .set({ status: "cancelled", finished_at: now, updated_at: now })
+      .where(
+        and(
+          eq(jobs.id, jobId),
+          eq(jobs.user_id, userId),
+          notInArray(jobs.status, ["completed", "failed", "cancelled"])
+        )
+      )
+      .returning({ id: jobs.id });
+    return updated.length > 0;
+  }
+
+  /**
+   * Which of these ids are now cancelled. The poller's one query per tick —
+   * indexed on the primary key, and bounded by how many runs an instance is
+   * actually executing.
+   */
+  static async cancelledAmong(jobIds: string[]): Promise<string[]> {
+    if (jobIds.length === 0) return [];
+    const db = getDb();
+    const rows = await db
+      .select({ id: jobs.id })
+      .from(jobs)
+      .where(and(inArray(jobs.id, jobIds), eq(jobs.status, "cancelled")));
+    return rows.map((row: { id: string }) => row.id);
+  }
+
   /** Find a job by id, scoped to the user. */
   static async find(userId: string, jobId: string): Promise<Job | null> {
     const job = await Job.get<Job>(jobId);

packages/models/src/migrations/versions.ts

@@ -2497,6 +2497,29 @@ export const migrations: MigrationDef[] = [
       await db.execute("DROP INDEX IF EXISTS idx_tsv_timeline");
       await db.execute("DROP TABLE IF EXISTS timeline_sequence_versions");
     }
+  },
+
+  // ── Stamp the machine executing a run onto its job row ─────────────
+  // With more than one server instance behind the Fly proxy, a reconnecting
+  // client lands on a random machine while the run lives on exactly one. The
+  // row records which, so the upgrade can be replayed to the owner and a
+  // cancel can be addressed at it. Null everywhere on a single-machine
+  // deployment, where the column is inert.
+  {
+    version: "20260805_000000",
+    name: "add_runner_instance_to_jobs",
+    createsTables: [],
+    modifiesTables: ["nodetool_jobs"],
+    async up(db) {
+      if (!(await db.tableExists("nodetool_jobs"))) return;
+      if (await db.columnExists("nodetool_jobs", "runner_instance")) return;
+      await db.execute(
+        "ALTER TABLE nodetool_jobs ADD COLUMN runner_instance TEXT"
+      );
+    },
+    async down() {
+      // no-op: dropping columns is unsafe across dialects and versions
+    }
   }
 ];
 

packages/models/src/schema-pg/jobs.ts

@@ -35,6 +35,9 @@ export const jobs = pgTable(
     ),
     execution_strategy: text("execution_strategy"),
     execution_id: text("execution_id"),
+    // The server instance executing this run (Fly machine id). Null on a
+    // single-machine deployment; see packages/websocket/src/lib/instance-id.ts.
+    runner_instance: text("runner_instance"),
     metadata_json: jsonText<Record<string, unknown>>()("metadata_json"),
     created_at: text("created_at").notNull(),
     updated_at: text("updated_at").notNull()

packages/models/src/schema/jobs.ts

@@ -41,6 +41,9 @@ export const jobs = sqliteTable(
     ),
     execution_strategy: text("execution_strategy"),
     execution_id: text("execution_id"),
+    // The server instance executing this run (Fly machine id). Null on a
+    // single-machine deployment; see packages/websocket/src/lib/instance-id.ts.
+    runner_instance: text("runner_instance"),
     metadata_json: jsonText<Record<string, unknown>>()("metadata_json"),
     created_at: text("created_at").notNull(),
     updated_at: text("updated_at").notNull()

packages/models/tests/job-cancel.test.ts

@@ -0,0 +1,87 @@
+/**
+ * The two conditional job queries the cross-instance cancel path depends on.
+ *
+ * Both exist to avoid a read-modify-write: a cancel arriving on an instance
+ * that does not own the run races the owner's own terminal write, and a full
+ * row `save()` built from a pre-race snapshot would overwrite whatever the
+ * owner had just recorded.
+ */
+import { describe, it, expect, beforeEach } from "vitest";
+
+import { initTestDb } from "../src/db.js";
+import { Job } from "../src/job.js";
+
+const create = (id: string, status: string, userId = "1") =>
+  Job.create({
+    id,
+    workflow_id: "wf",
+    user_id: userId,
+    status,
+    params: {},
+    graph: { nodes: [], edges: [] }
+  });
+
+describe("Job.markCancelledIfActive", () => {
+  beforeEach(() => {
+    initTestDb();
+  });
+
+  it("cancels an active row and says so", async () => {
+    await create("active", "running");
+
+    expect(await Job.markCancelledIfActive("active", "1")).toBe(true);
+
+    const row = await Job.get<Job>("active");
+    expect(row?.status).toBe("cancelled");
+    expect(row?.finished_at).toBeTruthy();
+  });
+
+  it("leaves a terminal row untouched and says nothing changed", async () => {
+    for (const status of ["completed", "failed", "cancelled"]) {
+      await create(status, status);
+      expect(await Job.markCancelledIfActive(status, "1")).toBe(false);
+      expect((await Job.get<Job>(status))?.status).toBe(status);
+    }
+  });
+
+  it("preserves what the owner wrote when it loses the race", async () => {
+    await create("raced", "running");
+
+    // The owner finishes first, cost and all.
+    const owner = await Job.get<Job>("raced");
+    owner!.markCompleted();
+    owner!.cost = 0.42;
+    await owner!.save();
+
+    expect(await Job.markCancelledIfActive("raced", "1")).toBe(false);
+
+    const row = await Job.get<Job>("raced");
+    expect(row?.status).toBe("completed");
+    expect(row?.cost).toBe(0.42);
+  });
+
+  it("will not cancel another user's run", async () => {
+    await create("theirs", "running", "2");
+
+    expect(await Job.markCancelledIfActive("theirs", "1")).toBe(false);
+    expect((await Job.get<Job>("theirs"))?.status).toBe("running");
+  });
+});
+
+describe("Job.cancelledAmong", () => {
+  beforeEach(() => {
+    initTestDb();
+  });
+
+  it("returns only the cancelled ids", async () => {
+    await create("a", "running");
+    await create("b", "cancelled");
+    await create("c", "completed");
+
+    expect(await Job.cancelledAmong(["a", "b", "c", "missing"])).toEqual(["b"]);
+  });
+
+  it("queries nothing for an empty list", async () => {
+    expect(await Job.cancelledAmong([])).toEqual([]);
+  });
+});

packages/models/tests/migrations.test.ts

@@ -423,7 +423,7 @@ describe("MigrationRunner", () => {
 // ── Built-in migrations smoke test ───────────────────────────────────
 
 describe("Built-in migrations", () => {
-  const EXPECTED_BUILT_IN_MIGRATION_COUNT = 57;
+  const EXPECTED_BUILT_IN_MIGRATION_COUNT = 58;
 
   it("should have correct count of migrations", () => {
     expect(migrations.length).toBe(EXPECTED_BUILT_IN_MIGRATION_COUNT);
@@ -479,6 +479,11 @@ describe("Built-in migrations", () => {
       await adapter.columnExists("nodetool_messages", "provider_session")
     ).toBe(true);
 
+    // The owning-instance stamp multi-machine routing reads.
+    expect(
+      await adapter.columnExists("nodetool_jobs", "runner_instance")
+    ).toBe(true);
+
     // The trigger safety counters are added by migration onto a table an
     // earlier migration created.
     for (const column of [