From aeece77eade57c359ce55e1cfeae8cfb7ee72bf5 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Tue, 11 Aug 2026 08:10:24 -0700 Subject: [PATCH] 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 --- ...ver.sessions.worker-original-order.test.ts | 1 + .../placement-dispatch-reclaim.test.ts | 8 +-- .../placement-dispatch-test-fixtures.ts | 11 +++- .../placement-dispatch-test-harness.ts | 5 +- .../placement-dispatch.test.ts | 52 +++++++++++++++++- .../worker-environments/placement-dispatch.ts | 17 +++--- .../worker-session-tool-executor.test.ts | 55 ++++++++++++++++++- 7 files changed, 131 insertions(+), 18 deletions(-) diff --git a/src/gateway/server.sessions.worker-original-order.test.ts b/src/gateway/server.sessions.worker-original-order.test.ts index b5152be2c978..ffe0e49ce544 100644 --- a/src/gateway/server.sessions.worker-original-order.test.ts +++ b/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`, diff --git a/src/gateway/worker-environments/placement-dispatch-reclaim.test.ts b/src/gateway/worker-environments/placement-dispatch-reclaim.test.ts index 006bb11dd9a2..fc95aa44f776 100644 --- a/src/gateway/worker-environments/placement-dispatch-reclaim.test.ts +++ b/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 () => { diff --git a/src/gateway/worker-environments/placement-dispatch-test-fixtures.ts b/src/gateway/worker-environments/placement-dispatch-test-fixtures.ts index 100cbb3e224d..a54864afa36c 100644 --- a/src/gateway/worker-environments/placement-dispatch-test-fixtures.ts +++ b/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", diff --git a/src/gateway/worker-environments/placement-dispatch-test-harness.ts b/src/gateway/worker-environments/placement-dispatch-test-harness.ts index 6a6ddf0cc129..a2e3b31ab5d5 100644 --- a/src/gateway/worker-environments/placement-dispatch-test-harness.ts +++ b/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 () => { diff --git a/src/gateway/worker-environments/placement-dispatch.test.ts b/src/gateway/worker-environments/placement-dispatch.test.ts index c59c789d4f70..dd94ad72fdc1 100644 --- a/src/gateway/worker-environments/placement-dispatch.test.ts +++ b/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(); diff --git a/src/gateway/worker-environments/placement-dispatch.ts b/src/gateway/worker-environments/placement-dispatch.ts index 5f0674fae7a3..b75c48c2d2c9 100644 --- a/src/gateway/worker-environments/placement-dispatch.ts +++ b/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, diff --git a/src/gateway/worker-environments/worker-session-tool-executor.test.ts b/src/gateway/worker-environments/worker-session-tool-executor.test.ts index 0bfe9dcd28de..b13c4ec90d0f 100644 --- a/src/gateway/worker-environments/worker-session-tool-executor.test.ts +++ b/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((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((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, {