feat(server): add activation commands and dispatch guard

Author: pallyoung

packages/server/src/__tests__/activation-commands.test.ts

@@ -0,0 +1,50 @@
+import { describe, expect, it } from "vitest";
+import { ActivationManager } from "../ws/activation.js";
+import { type CommandContext, dispatch } from "../ws/dispatch.js";
+import "../commands/activation.js";
+
+describe("activation commands", () => {
+  it("returns generation data from activation.claim", async () => {
+    const ctx = {
+      activationMgr: new ActivationManager(),
+    } as unknown as CommandContext;
+
+    const result = await dispatch(
+      {
+        kind: "command",
+        id: "00000000-0000-4000-8000-000000000001",
+        op: "activation.claim",
+        args: { clientInstanceId: "client-a" },
+      },
+      ctx,
+      "ws-a"
+    );
+
+    expect(result.ok).toBe(true);
+    expect(result.data).toMatchObject({
+      active: true,
+      generation: 1,
+      recoveryMode: "fresh",
+    });
+  });
+
+  it("rejects non-activation commands when activation is missing", async () => {
+    const ctx = {
+      activationMgr: new ActivationManager(),
+    } as unknown as CommandContext;
+
+    const result = await dispatch(
+      {
+        kind: "command",
+        id: "00000000-0000-4000-8000-000000000002",
+        op: "workspace.list",
+        args: {},
+      },
+      ctx,
+      "ws-a"
+    );
+
+    expect(result.ok).toBe(false);
+    expect(result.error?.code).toBe("activation_required");
+  });
+});

packages/server/src/commands/activation.ts

@@ -0,0 +1,27 @@
+import { z } from "zod";
+import { registerCommand } from "../ws/dispatch.js";
+
+registerCommand(
+  "activation.claim",
+  z.object({ clientInstanceId: z.string().min(1) }),
+  async (args, ctx, clientId) => {
+    return ctx.activationMgr.claim(args.clientInstanceId, clientId!);
+  }
+);
+
+registerCommand(
+  "activation.heartbeat",
+  z.object({ clientInstanceId: z.string(), generation: z.number().int().positive() }),
+  async (args, ctx) => {
+    return { ok: ctx.activationMgr.heartbeat(args.clientInstanceId, args.generation) };
+  }
+);
+
+registerCommand(
+  "activation.release",
+  z.object({ clientInstanceId: z.string(), generation: z.number().int().positive() }),
+  async (args, ctx) => {
+    ctx.activationMgr.release(args.clientInstanceId, args.generation);
+    return { ok: true };
+  }
+);

packages/server/src/commands/index.ts

@@ -6,6 +6,7 @@
 
 import "./workspace.js";
 import "./workspace-activity.js";
+import "./activation.js";
 import "./connection.js";
 import "./session.js";
 import "./terminal.js";

packages/server/src/server.ts

@@ -38,6 +38,7 @@ import type { TerminalDatabase } from "./terminal/types.js";
 import { deleteWorkspaceUploads, runStartupGc } from "./uploads/cleanup.js";
 import { STARTUP_GC_DELAY_MS } from "./uploads/constants.js";
 import { WorkspaceManager } from "./workspace/manager.js";
+import { ActivationManager } from "./ws/activation.js";
 import type { CommandContext } from "./ws/dispatch.js";
 import { dispatch } from "./ws/dispatch.js";
 import { FencingManager } from "./ws/fencing.js";
@@ -66,6 +67,7 @@ export async function createServer(
 
   const db = openDatabase(config.dataDir);
   const eventBus = new EventBus();
+  const activationMgr = new ActivationManager();
   const fencingMgr = new FencingManager();
   const wsHub = new WsHub({ eventBus, commandContext: null, config, fencingMgr });
   let workspaceMgr: WorkspaceManager;
@@ -211,6 +213,7 @@ export async function createServer(
     autoFetch,
     providerRuntimeDeps,
     providerInstallMgr,
+    activationMgr,
   };
 
   wsHub.setCommandContext(commandContext);

packages/server/src/ws/activation.ts

@@ -0,0 +1,60 @@
+export interface ActivationLease {
+  clientInstanceId: string;
+  wsClientId: string;
+  generation: number;
+}
+
+export interface ActivationClaimResult {
+  active: true;
+  generation: number;
+  recoveryMode: "fresh" | "grace_recover" | "takeover";
+}
+
+export class ActivationManager {
+  private lease: ActivationLease | null = null;
+  private generation = 0;
+
+  claim(clientInstanceId: string, wsClientId: string): ActivationClaimResult {
+    const current = this.lease;
+
+    if (current && current.clientInstanceId === clientInstanceId) {
+      current.wsClientId = wsClientId;
+      return {
+        active: true,
+        generation: current.generation,
+        recoveryMode: "grace_recover",
+      };
+    }
+
+    this.generation += 1;
+    this.lease = {
+      clientInstanceId,
+      wsClientId,
+      generation: this.generation,
+    };
+
+    return {
+      active: true,
+      generation: this.lease.generation,
+      recoveryMode: current ? "takeover" : "fresh",
+    };
+  }
+
+  heartbeat(clientInstanceId: string, generation: number): boolean {
+    return (
+      this.lease?.clientInstanceId === clientInstanceId && this.lease.generation === generation
+    );
+  }
+
+  release(clientInstanceId: string, generation: number): void {
+    if (this.lease?.clientInstanceId !== clientInstanceId || this.lease.generation !== generation) {
+      return;
+    }
+
+    this.lease = null;
+  }
+
+  getLease(): ActivationLease | null {
+    return this.lease;
+  }
+}

packages/server/src/ws/dispatch.ts

@@ -15,6 +15,7 @@ import type { Database } from "../storage/database.js";
 import type { SupervisorManager } from "../supervisor/manager.js";
 import type { TerminalManager } from "../terminal/manager.js";
 import type { WorkspaceManager } from "../workspace/manager.js";
+import type { ActivationManager } from "./activation.js";
 import type { FencingManager } from "./fencing.js";
 import type { Broadcaster } from "./hub.js";
 
@@ -34,6 +35,7 @@ export interface CommandContext {
   autoFetch: AutoFetchRuntime;
   providerRuntimeDeps?: RuntimeStatusDeps;
   providerInstallMgr?: ProviderInstallManager;
+  activationMgr: ActivationManager;
 }
 
 /**
@@ -56,6 +58,12 @@ const handlers = new Map<string, CommandHandler>();
  * Registry of all command schemas
  */
 const schemas = new Map<string, CommandSchema>();
+const ACTIVATION_ALLOWLIST = new Set([
+  "activation.claim",
+  "activation.heartbeat",
+  "activation.release",
+  "connection.probe",
+]);
 
 /**
  * Register a command handler
@@ -77,6 +85,21 @@ export async function dispatch(
   ctx: CommandContext,
   clientId?: string
 ): Promise<Result> {
+  if (!ACTIVATION_ALLOWLIST.has(msg.op)) {
+    const lease = ctx.activationMgr.getLease();
+    if (!clientId || !lease || lease.wsClientId !== clientId) {
+      return {
+        kind: "result",
+        id: msg.id,
+        ok: false,
+        error: {
+          code: "activation_required",
+          message: "This tab is no longer the active session",
+        },
+      };
+    }
+  }
+
   const handler = handlers.get(msg.op);
 
   if (!handler) {