perf(cloud): reuse attached tunnel during dispatch (#122077)

Amp-Thread-ID: https://ampcode.com/threads/T-019feaaa-c7ed-769e-9f29-a3612bec72e7

Co-authored-by: Amp <amp@ampcode.com>

src/gateway/server.sessions.worker-original-order.test.ts

@@ -450,6 +450,7 @@ test("preserves ordered fallback through restart, workspace sync, and safe sessi
     profileId: PROFILE_ID,
   });
   expect(active).toMatchObject({ state: "active", environmentId: ENVIRONMENT_ID });
+  expect(runner.starts).toHaveLength(1);
   await expect(fs.stat(runner.bootstrapUploadPath)).rejects.toMatchObject({ code: "ENOENT" });
   await expect(fs.readFile(runner.bootstrapReceiptPath, "utf8")).resolves.toBe(
     `${JSON.stringify(RECEIPT)}\n`,

src/gateway/worker-environments/placement-dispatch-reclaim.test.ts

@@ -36,7 +36,7 @@ describe("worker placement dispatch reclaim", () => {
     await fs.rm(root, { recursive: true, force: true });
   });
 
-  it("orders the migration barrier, provisioning, sync, attachment, and activation", async () => {
+  it("attaches before opening one tunnel for workspace sync and activation", async () => {
     const harness = createHarness(placementStore);
 
     await expect(harness.service.dispatch(REQUEST)).resolves.toMatchObject({
@@ -55,14 +55,14 @@ describe("worker placement dispatch reclaim", () => {
       "placement:provisioning",
       "create",
       "placement:syncing",
-      "tunnel:ready",
-      "sync",
-      "placement:starting",
       "attach",
       "tunnel:attached",
+      "sync",
+      "placement:starting",
       "activation",
       "placement:active",
     ]);
+    expect(harness.environments.startTunnel).toHaveBeenCalledOnce();
   });
 
   it("reclaims an unchanged active placement through the fenced teardown lifecycle", async () => {

src/gateway/worker-environments/placement-dispatch-test-fixtures.ts

@@ -21,7 +21,6 @@ export type DispatchStage =
   | "workspace"
   | "preflight"
   | "create"
-  | "tunnel:ready"
   | "sync"
   | "attach"
   | "tunnel:attached"
@@ -36,7 +35,7 @@ export const REQUEST: WorkerDispatchRequest = {
   profileId: "development",
 };
 
-export function seedStartingPlacement(
+export function seedSyncingPlacement(
   store: PlacementStore,
   environmentId: string,
 ): WorkerSessionPlacementRecord {
@@ -55,6 +54,14 @@ export function seedStartingPlacement(
     expectedGeneration: current.generation,
     patch: { workerBundleHash: BUNDLE_HASH },
   });
+  return current;
+}
+
+export function seedStartingPlacement(
+  store: PlacementStore,
+  environmentId: string,
+): WorkerSessionPlacementRecord {
+  let current = seedSyncingPlacement(store, environmentId);
   current = store.transition({
     sessionId: REQUEST.sessionId,
     from: "syncing",

src/gateway/worker-environments/placement-dispatch-test-harness.ts

@@ -308,7 +308,10 @@ export function createHarness(
       return minted;
     }),
     startTunnel: vi.fn(async ({ ownerEpoch }) => {
-      fail(ownerEpoch === 1 ? "tunnel:ready" : "tunnel:attached");
+      fail("tunnel:attached");
+      if (ownerEpoch !== currentEnvironment?.ownerEpoch) {
+        throw new Error("tunnel fixture received a stale owner epoch");
+      }
       return tunnelHandle(ownerEpoch);
     }),
     stopTunnel: vi.fn(async () => {

src/gateway/worker-environments/placement-dispatch.test.ts

@@ -13,6 +13,7 @@ import {
   type DispatchStage,
   type PlacementStore,
   REQUEST,
+  seedSyncingPlacement,
 } from "./placement-dispatch-test-fixtures.js";
 import { createHarness } from "./placement-dispatch-test-harness.js";
 import { createWorkerSessionPlacementStore } from "./placement-store.js";
@@ -72,6 +73,10 @@ describe("worker placement dispatch", () => {
     ).rejects.toThrow("sync failed");
 
     expect(states).toEqual(["requested", "provisioning", "syncing", "failed"]);
+    expect(harness.environments.stopTunnel).toHaveBeenCalledWith(
+      harness.attached.environmentId,
+      harness.attached.ownerEpoch,
+    );
   });
 
   it("recovers a completed turn's durable pending workspace result before stale-claim teardown", async () => {
@@ -506,7 +511,6 @@ describe("worker placement dispatch", () => {
     "barrier",
     "workspace",
     "create",
-    "tunnel:ready",
     "sync",
     "attach",
     "tunnel:attached",
@@ -610,6 +614,25 @@ describe("worker placement dispatch", () => {
     },
   );
 
+  it("tears down the attached owner after restart interrupts workspace sync", async () => {
+    const harness = createHarness(placementStore);
+    const interrupted = seedSyncingPlacement(placementStore, harness.attached.environmentId);
+    harness.markEnvironmentOwnerEpoch(harness.attached.ownerEpoch);
+
+    await harness.service.reconcile();
+
+    expect(harness.placements.current()).toMatchObject({
+      state: "failed",
+      recoveryError: "Worker dispatch interrupted in syncing",
+    });
+    expect(harness.environments.attachSession).not.toHaveBeenCalled();
+    expect(harness.environments.stopTunnel).toHaveBeenCalledWith(
+      interrupted.environmentId,
+      undefined,
+    );
+    expect(harness.environments.destroy).toHaveBeenCalledOnce();
+  });
+
   it("does not fail or tear down a dispatch owned by another invocation", async () => {
     placementStore.startDispatch(REQUEST);
     const harness = createHarness(placementStore);
@@ -814,6 +837,33 @@ describe("worker placement dispatch", () => {
     expect(harness.environments.create).not.toHaveBeenCalled();
   });
 
+  it("resumes a synced starting placement with its existing attached owner", async () => {
+    const harness = createHarness(placementStore);
+    harness.placements.seedStarting();
+    harness.markEnvironmentOwnerEpoch(harness.attached.ownerEpoch);
+    harness.log.length = 0;
+
+    await harness.service.reconcile();
+
+    expect(harness.placements.current()).toMatchObject({
+      state: "active",
+      environmentId: harness.attached.environmentId,
+      activeOwnerEpoch: harness.attached.ownerEpoch,
+    });
+    expect(harness.log).toEqual([
+      "environment:reconcile",
+      "workspace",
+      "tunnel:attached",
+      "activation",
+      "placement:active",
+    ]);
+    expect(harness.environments.attachSession).not.toHaveBeenCalled();
+    expect(harness.environments.startTunnel).toHaveBeenCalledWith({
+      environmentId: harness.attached.environmentId,
+      ownerEpoch: harness.attached.ownerEpoch,
+    });
+  });
+
   it("tears down a starting worker missing execution context instead of resuming it", async () => {
     const harness = createHarness(placementStore);
     harness.placements.seedStarting();

src/gateway/worker-environments/placement-dispatch.ts

@@ -189,8 +189,14 @@ export function createWorkerPlacementDispatchService(options: WorkerPlacementDis
         },
       });
       reportTransition(onTransition, placement);
-      const readyTunnel = await environments.startTunnel({ environmentId, ownerEpoch });
-      const synced = await readyTunnel.syncWorkspace({
+      const credential = await environments.attachSession({
+        environmentId,
+        ownerEpoch,
+        sessionId: request.sessionId,
+      });
+      ownerEpoch = credential.ownerEpoch;
+      const tunnel = await environments.startTunnel({ environmentId, ownerEpoch });
+      const synced = await tunnel.syncWorkspace({
         localPath,
         sessionId: request.sessionId,
         generation: placement.generation,
@@ -206,13 +212,6 @@ export function createWorkerPlacementDispatchService(options: WorkerPlacementDis
         },
       });
       reportTransition(onTransition, placement);
-      const credential = await environments.attachSession({
-        environmentId,
-        ownerEpoch,
-        sessionId: request.sessionId,
-      });
-      ownerEpoch = credential.ownerEpoch;
-      await environments.startTunnel({ environmentId, ownerEpoch });
       const startingPlacement = placement;
       const activePlacement = await options.runActivationBarrier({
         sessionId: request.sessionId,

src/gateway/worker-environments/worker-session-tool-executor.test.ts

@@ -38,7 +38,7 @@ vi.mock("../session-utils.js", async (importOriginal) => {
 vi.mock("../../agents/tools/sessions-send-tool.js", () => ({
   createSessionsSendTool: (options: unknown) => ({
     execute: async (toolCallId: string, args: unknown) => {
-      delivered({ args, options, toolCallId });
+      await delivered({ args, options, toolCallId });
       return {
         content: [{ type: "text", text: "sent" }],
         details: { status: "ok" },
@@ -350,6 +350,36 @@ describe("worker session tool topology", () => {
     expect(replay.resultJson).toBe(first.resultJson);
   });
 
+  it("coalesces concurrent spawn retries into one cloud child", async () => {
+    setEntry(SOURCE.sessionKey, SOURCE.sessionId);
+    const create = gatewayCreate.getMockImplementation();
+    if (!create) {
+      throw new Error("missing session creation fixture");
+    }
+    let finishCreate: (() => void) | undefined;
+    gatewayCreate.mockImplementation(async (request) => {
+      await new Promise<void>((resolve) => {
+        finishCreate = resolve;
+      });
+      return await create(request);
+    });
+    const request = {
+      identity,
+      toolName: "sessions_spawn" as const,
+      request: { toolCallId: "concurrent-spawn", task: "start one child" },
+    };
+
+    const retries = Array.from({ length: 32 }, () => execute(request));
+    await vi.waitFor(() => expect(gatewayCreate).toHaveBeenCalledOnce());
+    finishCreate?.();
+    const results = await Promise.all(retries);
+
+    expect(new Set(results.map((result) => result.resultJson))).toHaveLength(1);
+    expect(gatewayCreate).toHaveBeenCalledOnce();
+    expect(dispatchChild).toHaveBeenCalledOnce();
+    expect(gatewayRequest).toHaveBeenCalledOnce();
+  });
+
   it("recovers a committed child when session creation loses its response", async () => {
     setEntry(SOURCE.sessionKey, SOURCE.sessionId);
     gatewayCreate.mockImplementationOnce(
@@ -655,6 +685,29 @@ describe("worker session tool topology", () => {
     expect(secondKey).not.toBe(firstKey);
   });
 
+  it("coalesces concurrent retries into one message effect", async () => {
+    setEntry(SOURCE.sessionKey, SOURCE.sessionId);
+    setEntry(TARGET.sessionKey, TARGET.sessionId, {
+      sessionKey: SOURCE.sessionKey,
+      sessionId: SOURCE.sessionId,
+    });
+    let finishDelivery: (() => void) | undefined;
+    delivered.mockImplementation(
+      async () =>
+        await new Promise<void>((resolve) => {
+          finishDelivery = resolve;
+        }),
+    );
+
+    const retries = Array.from({ length: 32 }, () => send("concurrent-retry"));
+    await vi.waitFor(() => expect(delivered).toHaveBeenCalledOnce());
+    finishDelivery?.();
+    const results = await Promise.all(retries);
+
+    expect(new Set(results.map((result) => result.resultJson))).toHaveLength(1);
+    expect(delivered).toHaveBeenCalledOnce();
+  });
+
   it("replays a completed send after the target incarnation changes", async () => {
     setEntry(SOURCE.sessionKey, SOURCE.sessionId);
     setEntry(TARGET.sessionKey, TARGET.sessionId, {