fix(gateway): retire worker placements without sessions (#123785)

* fix(gateway): retire placements without sessions

* fix(gateway): preserve unreadable placement evidence

* fix(gateway): drain placement retirement on shutdown

* test(sessions): auto-clean identity probe temp dirs

* fix(gateway): join placement sidecar shutdown
This commit is contained in:
Peter Steinberger
2026-08-14 14:07:24 -07:00
committed by GitHub
parent bf70d5ddb6
commit 3c5e2ff296
13 changed files with 981 additions and 294 deletions
+5 -2
View File
@@ -190,8 +190,11 @@ stated honestly (revision 1 undersold this):
silent superseded-pairing pruning share one client-invalidation, credential,
environment, and placement reconciliation flow. Explicit RPCs wait for the
credential fence before success returns; periodic reconciliation retries
failed provider or placement cleanup. Unreferenced
terminal environment rows retain
failed provider or placement cleanup. Placement startup, identity-mutation,
and periodic reconciliation also compare each durable placement with the
canonical session entry: confirmed absence force-fences live environments
and exact-CAS retires safe terminal rows, while unreadable session evidence
retains the placement. Unreferenced terminal environment rows retain
seven days of operator diagnostics, then prune in bounded post-reconcile
batches; any surviving placement keeps its environment provenance.
Device-side GC of per-session workspace dirs and superseded bundles is a
@@ -1,6 +1,10 @@
import fs from "node:fs";
import { afterEach, describe, expect, it } from "vitest";
import { cleanupTempDirs, makeTempDir } from "../../../test/helpers/temp-dir.js";
import {
cleanupTempDirs,
makeTempDir,
useAutoCleanupTempDirTracker,
} from "../../../test/helpers/temp-dir.js";
import {
closeOpenClawAgentDatabasesForTest,
isOpenClawAgentDatabaseOpen,
@@ -15,11 +19,13 @@ import {
hasSessionEntriesByStatusReadOnly,
listSessionEntriesCore,
listSessionEntriesReadOnly,
readSessionIdentityEvidence,
resolveTranscriptSessionKeyBySessionId,
upsertSessionEntryCore,
} from "./session-accessor.js";
const tempDirs: string[] = [];
const autoTempDirs = useAutoCleanupTempDirTracker(afterEach);
function countRegisteredAgentDatabases(env: NodeJS.ProcessEnv): number {
const row = openOpenClawStateDatabase({ env })
@@ -138,6 +144,75 @@ describe("session accessor readonly listing", () => {
expect(countRegisteredAgentDatabases(env)).toBe(0);
});
it("probes session identity by exact key and indexed current session id", async () => {
const stateDir = autoTempDirs.make("openclaw-session-readonly-evidence-");
const env = { OPENCLAW_STATE_DIR: stateDir };
const agentId = "worker-1";
const sessionKey = "agent:worker-1:moved";
const sessionId = "session-1";
await upsertSessionEntryCore({ agentId, env, sessionKey }, { sessionId, updatedAt: 1 });
const storePath = resolveOpenClawAgentSqlitePath({ agentId, env });
closeOpenClawAgentDatabasesForTest();
expect(readSessionIdentityEvidence({ agentId, sessionId, sessionKey, storePath })).toEqual({
status: "current",
sessionKey,
});
expect(
readSessionIdentityEvidence({
agentId,
sessionId,
sessionKey: "agent:worker-1:old-key",
storePath,
}),
).toEqual({ status: "current", sessionKey });
expect(
readSessionIdentityEvidence({
agentId,
sessionId: "missing-session",
sessionKey,
storePath,
}),
).toEqual({ status: "absent" });
});
it("reports migration-invalid session evidence as unknown", async () => {
const stateDir = autoTempDirs.make("openclaw-session-readonly-evidence-invalid-");
const env = { OPENCLAW_STATE_DIR: stateDir };
const agentId = "worker-1";
const sessionKey = "agent:worker-1:main";
const sessionId = "session-1";
await upsertSessionEntryCore({ agentId, env, sessionKey }, { sessionId, updatedAt: 1 });
const database = openOpenClawAgentDatabase({ agentId, env });
database.db.exec("PRAGMA user_version = 999;");
const storePath = database.path;
closeOpenClawAgentDatabasesForTest();
expect(readSessionIdentityEvidence({ agentId, sessionId, sessionKey, storePath })).toEqual({
status: "unknown",
reason: "read-failed",
});
});
it("uses the current-session-id index for fallback identity probes", async () => {
const stateDir = autoTempDirs.make("openclaw-session-readonly-evidence-index-");
const env = { OPENCLAW_STATE_DIR: stateDir };
const agentId = "worker-1";
const database = openOpenClawAgentDatabase({ agentId, env });
const detail = database.db
.prepare(
"EXPLAIN QUERY PLAN SELECT session_key FROM session_nodes WHERE current_session_id = ? LIMIT 2",
)
.all("session-1")
.map((row) => {
const rowDetail = (row as { detail?: unknown }).detail;
return typeof rowDetail === "string" ? rowDetail : "";
})
.join(" ");
expect(detail).toContain("idx_agent_session_nodes_current_session_id");
});
it("does not register a populated database during readonly health-style listing", async () => {
const stateDir = makeTempDir(tempDirs, "openclaw-session-readonly-registry-");
const env = { OPENCLAW_STATE_DIR: stateDir };
@@ -1,4 +1,7 @@
import { executeSqliteQueryTakeFirstSync } from "../../infra/kysely-sync.js";
import {
executeSqliteQuerySync,
executeSqliteQueryTakeFirstSync,
} from "../../infra/kysely-sync.js";
import { withOpenClawAgentDatabaseReadOnly } from "../../state/openclaw-agent-db-readonly.js";
import type { ExactSessionEntry, SessionAccessScope } from "./session-accessor.sqlite-contract.js";
import { readExactSessionEntryRowValidated } from "./session-accessor.sqlite-entry-store.js";
@@ -10,6 +13,14 @@ import {
} from "./session-accessor.sqlite-scope.js";
import type { SessionEntry } from "./types.js";
export type SessionIdentityEvidenceResult =
| { status: "current"; sessionKey: string }
| { status: "absent" }
| {
status: "unknown";
reason: "ambiguous" | "read-failed" | "row-invalid" | "schema-missing" | "table-missing";
};
type ExactSessionEntryReadOnlyResult =
| { found: true; value: ExactSessionEntry | undefined }
| {
@@ -70,3 +81,62 @@ export function loadExactSessionEntryReadOnlyResult(
},
};
}
/** Indexed exact-key/session-id probe that preserves unreadable state as unknown. */
export function readSessionIdentityEvidence(params: {
agentId: string;
sessionId: string;
sessionKey: string;
storePath: string;
}): SessionIdentityEvidenceResult {
const resolved = resolveSqliteScope({
agentId: params.agentId,
sessionKey: params.sessionKey,
storePath: params.storePath,
});
let result:
| { found: true; value: SessionIdentityEvidenceResult }
| { found: false; reason: "database-missing" | "schema-missing" | "table-missing" };
try {
result = withOpenClawAgentDatabaseReadOnly((database): SessionIdentityEvidenceResult => {
const exact = readExactSessionEntryRowValidated(database, resolved.sessionKey)?.entry;
if (exact?.sessionId === params.sessionId) {
return { status: "current", sessionKey: resolved.sessionKey };
}
const rows = executeSqliteQuerySync(
database.db,
getSessionKysely(database.db)
.selectFrom("session_nodes")
.select(["session_key", "entry_valid"])
.where("current_session_id", "=", params.sessionId)
.limit(2),
).rows;
if (rows.length === 0) {
return { status: "absent" };
}
if (rows.length !== 1) {
return { status: "unknown", reason: "ambiguous" };
}
const row = rows[0];
if (row?.entry_valid === -1) {
return { status: "absent" };
}
const sessionKey = row?.session_key;
if (!sessionKey || row.entry_valid !== 1) {
return { status: "unknown", reason: "row-invalid" };
}
const entry = readExactSessionEntryRowValidated(database, sessionKey)?.entry;
return entry?.sessionId === params.sessionId
? { status: "current", sessionKey }
: { status: "unknown", reason: "row-invalid" };
}, toDatabaseOptions(resolved));
} catch {
return { status: "unknown", reason: "read-failed" };
}
if (result.found) {
return result.value;
}
return result.reason === "database-missing"
? { status: "absent" }
: { status: "unknown", reason: result.reason };
}
+4
View File
@@ -156,6 +156,10 @@ export {
updateResolvedSessionEntry,
upsertSessionEntryCore,
} from "./session-accessor.entry.js";
export {
readSessionIdentityEvidence,
type SessionIdentityEvidenceResult,
} from "./session-accessor.sqlite-entry-availability.js";
export {
createSessionEntryWithTranscript,
forkSessionEntryFromParentTarget,
@@ -0,0 +1,61 @@
import { afterEach, describe, expect, it } from "vitest";
import { useAutoCleanupTempDirTracker } from "../../test/helpers/temp-dir.js";
import { upsertSessionEntryCore } from "../config/sessions/session-accessor.js";
import {
closeOpenClawAgentDatabasesForTest,
openOpenClawAgentDatabase,
} from "../state/openclaw-agent-db.js";
import { withEnvAsync } from "../test-utils/env.js";
import { resolveWorkerPlacementSessionEvidence } from "./server-worker-placement-session-evidence.js";
import type { WorkerSessionPlacementRecord } from "./worker-environments/placement-record.js";
const tempDirs = useAutoCleanupTempDirTracker(afterEach);
afterEach(() => {
closeOpenClawAgentDatabasesForTest();
});
function localPlacement(
sessionId: string,
sessionKey: string,
): Extract<WorkerSessionPlacementRecord, { state: "local" }> {
return {
sessionId,
sessionKey,
agentId: "main",
state: "local",
generation: 1,
turnClaim: null,
environmentId: null,
activeOwnerEpoch: null,
workspaceBaseManifestRef: null,
remoteWorkspaceDir: null,
workerBundleHash: null,
lastTranscriptAckCursor: null,
lastLiveEventAckCursor: null,
recoveryError: null,
terminalReason: null,
terminalAtMs: null,
createdAtMs: 1,
updatedAtMs: 1,
stateChangedAtMs: 1,
};
}
describe("worker placement session evidence", () => {
it("keeps a placement when its session database is migration-invalid", async () => {
const stateDir = tempDirs.make("openclaw-placement-session-evidence-");
await withEnvAsync({ OPENCLAW_STATE_DIR: stateDir }, async () => {
const sessionId = "session-1";
const sessionKey = "agent:main:main";
await upsertSessionEntryCore({ agentId: "main", sessionKey }, { sessionId, updatedAt: 1 });
const database = openOpenClawAgentDatabase({ agentId: "main" });
database.db.exec("PRAGMA user_version = 999;");
closeOpenClawAgentDatabasesForTest();
await expect(
resolveWorkerPlacementSessionEvidence(localPlacement(sessionId, sessionKey)),
).resolves.toBe("unknown");
});
});
});
@@ -0,0 +1,33 @@
import { getRuntimeConfig } from "../config/config.js";
import { createLazyRuntimeModule } from "../shared/lazy-runtime.js";
import type { WorkerSessionPlacementRecord } from "./worker-environments/placement-record.js";
const loadPlacementSessionEvidenceRuntime = createLazyRuntimeModule(async () => {
const [sessionUtils, sessionAccessor] = await Promise.all([
import("./session-utils.js"),
import("../config/sessions/session-accessor.js"),
]);
return {
readSessionIdentityEvidence: sessionAccessor.readSessionIdentityEvidence,
resolveGatewaySessionStoreTarget: sessionUtils.resolveGatewaySessionStoreTarget,
};
});
/** Resolves authoritative session existence without treating unreadable state as absence. */
export async function resolveWorkerPlacementSessionEvidence(
placement: WorkerSessionPlacementRecord,
): Promise<"current" | "absent" | "unknown"> {
const runtime = await loadPlacementSessionEvidenceRuntime();
const target = runtime.resolveGatewaySessionStoreTarget({
cfg: getRuntimeConfig(),
key: placement.sessionKey,
agentId: placement.agentId,
});
const evidence = runtime.readSessionIdentityEvidence({
agentId: target.agentId,
sessionId: placement.sessionId,
sessionKey: target.canonicalKey,
storePath: target.storePath,
});
return evidence.status;
}
@@ -4,6 +4,7 @@ import { createDeferredCore } from "../shared/deferred.js";
const runtimeFactoryMocks = vi.hoisted(() => ({
createDispatch: vi.fn(),
createDiskSpace: vi.fn(),
resolveSessionEvidence: vi.fn(),
}));
vi.mock("./worker-environments/placement-dispatch.js", async (importOriginal) => {
@@ -15,6 +16,10 @@ vi.mock("./worker-environments/placement-dispatch.js", async (importOriginal) =>
};
});
vi.mock("./server-worker-placement-session-evidence.js", () => ({
resolveWorkerPlacementSessionEvidence: runtimeFactoryMocks.resolveSessionEvidence,
}));
vi.mock("./worker-environments/placement-disk-space.js", async (importOriginal) => {
const actual =
await importOriginal<typeof import("./worker-environments/placement-disk-space.js")>();
@@ -24,152 +29,7 @@ vi.mock("./worker-environments/placement-disk-space.js", async (importOriginal)
};
});
import {
coordinateWorkerPlacementDispatch,
createGatewayWorkerPlacementRuntime,
type GatewayWorkerPlacementRuntime,
} from "./server-worker-placement-startup.js";
import type { WorkerPlacementDispatchRequest } from "./worker-environments/service-contract.js";
type DispatchService = GatewayWorkerPlacementRuntime["dispatchService"];
const REQUEST: WorkerPlacementDispatchRequest = {
sessionId: "session-1",
sessionKey: "agent:main:session-1",
agentId: "main",
profileId: "test",
};
describe("worker placement dispatch coordinator", () => {
it("forwards the optional internal transition observer", async () => {
const observer = vi.fn();
const dispatch = vi.fn().mockResolvedValue({ state: "active" });
const service = {
dispatch,
forceDestroyEnvironment: vi.fn(),
reclaim: vi.fn(),
reconcile: vi.fn(),
reconcileActive: vi.fn(),
} as unknown as DispatchService;
await coordinateWorkerPlacementDispatch(service).dispatch(REQUEST, observer);
expect(dispatch).toHaveBeenCalledWith(REQUEST, observer);
});
it("coalesces an identical dispatch and rejects a conflicting in-flight request", async () => {
const dispatchStarted = createDeferredCore();
const releaseDispatch = createDeferredCore();
const active = { state: "active" };
const dispatch = vi.fn(async () => {
dispatchStarted.resolve();
await releaseDispatch.promise;
return active;
});
const service = {
dispatch,
forceDestroyEnvironment: vi.fn(),
reclaim: vi.fn(),
reconcile: vi.fn(),
reconcileActive: vi.fn(),
} as unknown as DispatchService;
const coordinated = coordinateWorkerPlacementDispatch(service);
const first = coordinated.dispatch(REQUEST);
await dispatchStarted.promise;
await expect(
coordinated.dispatch({ ...REQUEST, profileId: "another-profile" }),
).rejects.toThrow(`Session ${REQUEST.sessionKey} is already dispatching another request`);
await expect(
coordinated.dispatch({
...REQUEST,
inheritedProfile: {
providerId: "fake",
profileSnapshot: { settings: { region: "parent" } },
},
}),
).rejects.toThrow(`Session ${REQUEST.sessionKey} is already dispatching another request`);
const retry = coordinated.dispatch(REQUEST);
releaseDispatch.resolve();
const [firstResult, retryResult] = await Promise.all([first, retry]);
expect(retryResult).toBe(firstResult);
expect(dispatch).toHaveBeenCalledOnce();
await coordinated.dispatch({ ...REQUEST, profileId: "another-profile" });
expect(dispatch).toHaveBeenCalledTimes(2);
});
it("joins a retry before a queued reconciliation after dispatch failure", async () => {
const dispatchStarted = createDeferredCore();
const releaseDispatch = createDeferredCore();
const dispatchError = new Error("provision failed");
const dispatch = vi.fn(async () => {
dispatchStarted.resolve();
await releaseDispatch.promise;
throw dispatchError;
});
const reconcileActive = vi.fn();
const service = {
dispatch,
forceDestroyEnvironment: vi.fn(),
reclaim: vi.fn(),
reconcile: vi.fn(),
reconcileActive,
} as unknown as DispatchService;
const coordinated = coordinateWorkerPlacementDispatch(service);
const first = coordinated.dispatch(REQUEST);
await dispatchStarted.promise;
const reconciliation = coordinated.reconcileActive();
const retry = coordinated.dispatch(REQUEST);
const outcomes = Promise.allSettled([first, retry]);
releaseDispatch.resolve();
expect(await outcomes).toEqual([
{ status: "rejected", reason: dispatchError },
{ status: "rejected", reason: dispatchError },
]);
await reconciliation;
expect(dispatch).toHaveBeenCalledOnce();
expect(reconcileActive).toHaveBeenCalledOnce();
await expect(coordinated.dispatch({ ...REQUEST, profileId: "another-profile" })).rejects.toBe(
dispatchError,
);
expect(dispatch).toHaveBeenCalledTimes(2);
});
it("coalesces full sweeps but runs a fresh targeted pass with its environment id", async () => {
const fullSweepStarted = createDeferredCore();
const releaseFullSweep = createDeferredCore();
const reconcileActive = vi.fn(async (environmentId?: string) => {
if (environmentId === undefined) {
fullSweepStarted.resolve();
await releaseFullSweep.promise;
}
});
const service = {
dispatch: vi.fn(),
forceDestroyEnvironment: vi.fn(),
reclaim: vi.fn(),
reconcile: vi.fn(),
reconcileActive,
} as unknown as DispatchService;
const coordinated = coordinateWorkerPlacementDispatch(service);
const firstFullSweep = coordinated.reconcileActive();
const secondFullSweep = coordinated.reconcileActive();
await fullSweepStarted.promise;
const targetedSweep = coordinated.reconcileActive("worker-target");
expect(reconcileActive).toHaveBeenCalledTimes(1);
releaseFullSweep.resolve();
await Promise.all([firstFullSweep, secondFullSweep, targetedSweep]);
expect(reconcileActive.mock.calls).toEqual([[], ["worker-target"]]);
});
});
import { createGatewayWorkerPlacementRuntime } from "./server-worker-placement-startup.js";
describe("worker placement startup health lifetime", () => {
it("samples disk on schedule while reconciliation is stuck and drains both on stop", async () => {
@@ -205,6 +65,9 @@ describe("worker placement startup health lifetime", () => {
const warn = vi.fn();
const runtime = createGatewayWorkerPlacementRuntime({
placements: {
get: () => undefined,
list: () => [],
retireSessionPlacement: vi.fn(),
pruneOrphanedWorkspaceReconciliations: () => [],
listWorkspaceReconciliationOwners: () => [],
} as never,
@@ -233,7 +96,7 @@ describe("worker placement startup health lifetime", () => {
releaseScheduledHealth.reject(healthError);
await Promise.resolve();
expect(stopSettled).toBe(false);
expect(environments.stop).toHaveBeenCalledOnce();
expect(environments.stop).not.toHaveBeenCalled();
releaseReconcile.resolve();
await stopping;
@@ -244,4 +107,87 @@ describe("worker placement startup health lifetime", () => {
vi.useRealTimers();
}
});
it("drains deferred startup session evidence before stopping environments", async () => {
const evidence = createDeferredCore<"current">();
runtimeFactoryMocks.resolveSessionEvidence.mockImplementation(async () => evidence.promise);
runtimeFactoryMocks.createDiskSpace.mockReturnValue({
read: vi.fn(),
version: vi.fn(() => 0),
sweep: vi.fn().mockResolvedValue(undefined),
});
runtimeFactoryMocks.createDispatch.mockReturnValue({
dispatch: vi.fn(),
forceDestroyEnvironment: vi.fn(),
reclaim: vi.fn(),
reconcile: vi.fn().mockResolvedValue(undefined),
reconcileActive: vi.fn().mockResolvedValue(undefined),
});
const placement = {
sessionId: "session-startup",
sessionKey: "agent:main:startup",
agentId: "main",
state: "local",
generation: 1,
turnClaim: null,
environmentId: null,
activeOwnerEpoch: null,
workspaceBaseManifestRef: null,
remoteWorkspaceDir: null,
workerBundleHash: null,
lastTranscriptAckCursor: null,
lastLiveEventAckCursor: null,
recoveryError: null,
terminalReason: null,
terminalAtMs: null,
createdAtMs: 1,
updatedAtMs: 1,
stateChangedAtMs: 1,
} as const;
const environments = {
start: vi.fn(),
stop: vi.fn().mockResolvedValue(undefined),
};
const runtime = createGatewayWorkerPlacementRuntime({
placements: {
get: () => placement,
list: () => [placement],
retireSessionPlacement: vi.fn(),
pruneOrphanedWorkspaceReconciliations: () => [],
listWorkspaceReconciliationOwners: () => [],
} as never,
environments: environments as never,
admitNewPlacements: true,
revokeSessionAuthority: vi.fn(),
warn: vi.fn(),
});
let closeStarted = false;
let sidecar: { stop: () => Promise<void> } | undefined;
const starting = runtime.startRuntime({
isClosePreludeStarted: () => closeStarted,
registerSidecar: (registered) => {
sidecar = registered;
},
});
await vi.waitFor(() => expect(runtimeFactoryMocks.resolveSessionEvidence).toHaveBeenCalled());
closeStarted = true;
const stopping = sidecar?.stop();
const repeatedStop = sidecar?.stop();
if (!stopping || !repeatedStop) {
throw new Error("startup did not register its placement sidecar");
}
let repeatedStopSettled = false;
void repeatedStop.then(() => {
repeatedStopSettled = true;
});
await Promise.resolve();
expect(repeatedStop).toBe(stopping);
expect(repeatedStopSettled).toBe(false);
expect(environments.stop).not.toHaveBeenCalled();
evidence.resolve("current");
await expect(starting).resolves.toBeNull();
await Promise.all([stopping, repeatedStop]);
expect(environments.stop).toHaveBeenCalledOnce();
});
});
+46 -142
View File
@@ -1,4 +1,3 @@
import { isDeepStrictEqual } from "node:util";
import { installSessionPlacementAdmissionProvider } from "../agents/session-placement-admission.js";
import { clearSessionQueues } from "../auto-reply/reply/queue/cleanup.js";
import { getRuntimeConfig } from "../config/config.js";
@@ -10,16 +9,16 @@ import {
runExclusiveSessionLifecycleMutation,
SESSION_WORK_ADMISSION_DRAIN_TIMEOUT_MS,
} from "../sessions/session-lifecycle-admission.js";
import { onSessionIdentityMutation } from "../sessions/session-lifecycle-events.js";
import { createLazyRuntimeModule } from "../shared/lazy-runtime.js";
import { resolveWorkerPlacementSessionEvidence } from "./server-worker-placement-session-evidence.js";
import { createWorkerPlacementDiskSpaceMonitor } from "./worker-environments/placement-disk-space.js";
import {
createWorkerPlacementDispatchService,
type WorkerPlacementDispatchService,
} from "./worker-environments/placement-dispatch.js";
import { coordinateWorkerPlacementDispatch } from "./worker-environments/placement-dispatch-coordinator.js";
import { createWorkerPlacementDispatchService } from "./worker-environments/placement-dispatch.js";
import { FORCED_WORKER_ABANDONMENT_ERROR } from "./worker-environments/placement-force-abandon.js";
import { createPlacementSessionRetirement } from "./worker-environments/placement-session-retirement.js";
import type { WorkerSessionPlacementStore } from "./worker-environments/placement-store.js";
import { createReclaimedPlacementRedispatch } from "./worker-environments/reclaimed-placement-redispatch.js";
import type { WorkerPlacementDispatchRequest } from "./worker-environments/service-contract.js";
import type { WorkerEnvironmentService } from "./worker-environments/service.js";
import { createWorkerSessionTurnPlacementProvider } from "./worker-environments/worker-turn-launcher.js";
import { createWorkerWorkspaceOperationCoordinator } from "./worker-environments/workspace-operation-coordinator.js";
@@ -110,126 +109,6 @@ function resolveWorkerPlacementSessionTarget(params: {
return { config: params.config, target, entry, worktree };
}
/** Serializes reconciliation sweeps against in-flight dispatches so a sweep never
* observes a placement mid-transition. Dispatches wait out any pending sweep. */
export function coordinateWorkerPlacementDispatch(
service: WorkerPlacementDispatchService,
): WorkerPlacementDispatchService {
let activeDispatchCount = 0;
let reconciliation: Promise<void> | undefined;
const dispatchIdleWaiters = new Set<() => void>();
const waitForDispatchIdle = (): Promise<void> => {
if (activeDispatchCount === 0) {
return Promise.resolve();
}
return new Promise<void>((resolve) => {
dispatchIdleWaiters.add(resolve);
});
};
const runReconciliation = (operation: () => Promise<void>): Promise<void> => {
if (reconciliation) {
return reconciliation;
}
const current = (async () => {
await waitForDispatchIdle();
await operation();
})();
reconciliation = current;
const clearCurrent = () => {
if (reconciliation === current) {
reconciliation = undefined;
}
};
void current.then(clearCurrent, clearCurrent);
return current;
};
const runExclusivePlacementOperation = <T>(operation: () => Promise<T>): Promise<T> => {
const current = (async () => {
const pendingReconciliation = reconciliation;
if (pendingReconciliation) {
await pendingReconciliation.catch(() => undefined);
}
await waitForDispatchIdle();
return await operation();
})();
const barrier = current.then(
() => undefined,
() => undefined,
);
reconciliation = barrier;
return current.finally(() => {
if (reconciliation === barrier) {
reconciliation = undefined;
}
});
};
const runPlacementOperation = async <T>(operation: () => Promise<T>): Promise<T> => {
for (;;) {
const pendingReconciliation = reconciliation;
if (!pendingReconciliation) {
break;
}
await pendingReconciliation.catch(() => undefined);
}
activeDispatchCount += 1;
try {
return await operation();
} finally {
activeDispatchCount -= 1;
if (activeDispatchCount === 0) {
const waiters = [...dispatchIdleWaiters];
dispatchIdleWaiters.clear();
for (const resolve of waiters) {
resolve();
}
}
}
};
const dispatchInFlight = new Map<
string,
{
request: WorkerPlacementDispatchRequest;
operation: ReturnType<WorkerPlacementDispatchService["dispatch"]>;
}
>();
return {
dispatch: async (request, onTransition) => {
const inFlight = dispatchInFlight.get(request.sessionId);
if (inFlight) {
if (
inFlight.request.sessionKey !== request.sessionKey ||
inFlight.request.agentId !== request.agentId ||
inFlight.request.profileId !== request.profileId ||
inFlight.request.deviceId !== request.deviceId ||
!isDeepStrictEqual(inFlight.request.inheritedProfile, request.inheritedProfile)
) {
throw new Error(`Session ${request.sessionKey} is already dispatching another request`);
}
return await inFlight.operation;
}
const operation = runPlacementOperation(() => service.dispatch(request, onTransition));
dispatchInFlight.set(request.sessionId, { request, operation });
try {
return await operation;
} finally {
if (dispatchInFlight.get(request.sessionId)?.operation === operation) {
dispatchInFlight.delete(request.sessionId);
}
}
},
forceDestroyEnvironment: (environmentId, onCleanupError) =>
runExclusivePlacementOperation(() =>
service.forceDestroyEnvironment(environmentId, onCleanupError),
),
reclaim: async (request) => await runPlacementOperation(() => service.reclaim(request)),
reconcile: () => runReconciliation(service.reconcile),
reconcileActive: (environmentId) =>
environmentId === undefined
? runReconciliation(() => service.reconcileActive())
: runExclusivePlacementOperation(() => service.reconcileActive(environmentId)),
};
}
type WorkerPlacementSidecar = { stop: () => Promise<void> };
export type GatewayWorkerPlacementRuntimeParams = {
@@ -521,6 +400,13 @@ export function createGatewayWorkerPlacementRuntime(params: GatewayWorkerPlaceme
workspaceOperations,
}),
);
const sessionRetirement = createPlacementSessionRetirement({
placements: params.placements,
environments: params.environments,
forceDestroyEnvironment: dispatchService.forceDestroyEnvironment,
resolveSessionEvidence: resolveWorkerPlacementSessionEvidence,
warn: params.warn,
});
const admissionProvider = createWorkerSessionTurnPlacementProvider({
environments: params.environments,
placements: params.placements,
@@ -608,7 +494,10 @@ export function createGatewayWorkerPlacementRuntime(params: GatewayWorkerPlaceme
}
return trackOperation(
placementReconcile,
dispatchService.reconcileActive(),
(async () => {
await sessionRetirement.reconcile();
await dispatchService.reconcileActive();
})(),
"Worker placement reconcile sweep failed",
);
};
@@ -626,26 +515,38 @@ export function createGatewayWorkerPlacementRuntime(params: GatewayWorkerPlaceme
// Session-lifetime sampling covers idle placements independently of provider health.
void sweepDiskSpace();
};
const sidecar: WorkerPlacementSidecar = {
stop: async () => {
if (stopped) {
const uninstallSessionIdentityMutation = onSessionIdentityMutation((mutation) => {
const previousSessionId = mutation.previous.sessionId;
const currentSessionId = "current" in mutation ? mutation.current.sessionId : undefined;
if (previousSessionId && previousSessionId !== currentSessionId) {
const pending = placementReconcile.current;
if (!pending) {
void reconcileActivePlacements();
return;
}
void pending.then(reconcileActivePlacements, reconcileActivePlacements);
}
});
let stopPromise: Promise<void> | undefined;
const sidecar: WorkerPlacementSidecar = {
stop: () => {
if (stopPromise) {
return stopPromise;
}
stopped = true;
clearInterval(placementReconcileInterval);
placementReconcileInterval = undefined;
uninstallSessionIdentityMutation();
uninstallPlacementAdmission();
const environmentStop = params.environments.stop();
const stopResults = await Promise.allSettled([
...[placementReconcile.current, diskSpaceSweep.current].filter(
(operation): operation is Promise<void> => operation !== undefined,
),
environmentStop,
]);
const environmentStopResult = stopResults.at(-1);
if (environmentStopResult?.status === "rejected") {
throw environmentStopResult.reason;
}
stopPromise = (async () => {
await Promise.allSettled(
[placementReconcile.current, diskSpaceSweep.current].filter(
(operation): operation is Promise<void> => operation !== undefined,
),
);
await params.environments.stop();
})();
return stopPromise;
},
};
// Close must see the drain handle before reconciliation can yield.
@@ -665,7 +566,10 @@ export function createGatewayWorkerPlacementRuntime(params: GatewayWorkerPlaceme
await sidecar.stop();
return null;
}
const startupReconcile = dispatchService.reconcile();
const startupReconcile = (async () => {
await dispatchService.reconcile();
await sessionRetirement.reconcile();
})();
placementReconcile.current = startupReconcile;
try {
try {
@@ -1,4 +1,5 @@
import { afterEach, expect, test } from "vitest";
import { afterEach, expect, test, vi } from "vitest";
import { deleteSessionEntryLifecycle } from "../config/sessions/session-accessor.js";
import { getFallbackGatewayContext } from "./server-plugin-fallback-context.js";
import { startGatewayServerHarness, type GatewayServerHarness } from "./server.e2e-ws-harness.js";
import { loadSessionEntry } from "./session-utils.js";
@@ -104,6 +105,37 @@ test(
lifecycleRevision: resetLifecycleRevision,
});
expect(placements.get(resetSessionId)).toBeUndefined();
const createdForExternalDelete = await rpcReq<{ key?: string; sessionId?: string }>(
ws,
"sessions.create",
{ agentId: "main", key: "startup-placement-external-delete" },
);
const externalSessionId = createdForExternalDelete.payload?.sessionId;
const externalSessionKey = createdForExternalDelete.payload?.key;
if (!externalSessionId || !externalSessionKey) {
throw new Error("external-delete session creation did not return placement identity");
}
const externalClaim = placements.claimTurn({
sessionId: externalSessionId,
sessionKey: externalSessionKey,
agentId: "main",
owner: { kind: "local" },
claimId: "startup-placement-external-delete-claim",
runId: "startup-placement-external-delete-run",
});
placements.releaseTurn(externalClaim);
const externalTarget = loadSessionEntry(externalSessionKey);
await deleteSessionEntryLifecycle({
archiveTranscript: false,
storePath: externalTarget.storePath,
target: {
canonicalKey: externalTarget.canonicalKey,
storeKeys: externalTarget.storeKeys,
},
});
await vi.waitFor(() => expect(placements.get(externalSessionId)).toBeUndefined());
expect(getFallbackGatewayContext()?.workerEnvironmentService).toBeDefined();
ws.close();
},
@@ -0,0 +1,145 @@
import { describe, expect, it, vi } from "vitest";
import { createDeferredCore } from "../../shared/deferred.js";
import { coordinateWorkerPlacementDispatch } from "./placement-dispatch-coordinator.js";
import type { WorkerPlacementDispatchService } from "./placement-dispatch.js";
import type { WorkerPlacementDispatchRequest } from "./service-contract.js";
type DispatchService = WorkerPlacementDispatchService;
const REQUEST: WorkerPlacementDispatchRequest = {
sessionId: "session-1",
sessionKey: "agent:main:session-1",
agentId: "main",
profileId: "test",
};
describe("worker placement dispatch coordinator", () => {
it("forwards the optional internal transition observer", async () => {
const observer = vi.fn();
const dispatch = vi.fn().mockResolvedValue({ state: "active" });
const service = {
dispatch,
forceDestroyEnvironment: vi.fn(),
reclaim: vi.fn(),
reconcile: vi.fn(),
reconcileActive: vi.fn(),
} as unknown as DispatchService;
await coordinateWorkerPlacementDispatch(service).dispatch(REQUEST, observer);
expect(dispatch).toHaveBeenCalledWith(REQUEST, observer);
});
it("coalesces an identical dispatch and rejects a conflicting in-flight request", async () => {
const dispatchStarted = createDeferredCore();
const releaseDispatch = createDeferredCore();
const active = { state: "active" };
const dispatch = vi.fn(async () => {
dispatchStarted.resolve();
await releaseDispatch.promise;
return active;
});
const service = {
dispatch,
forceDestroyEnvironment: vi.fn(),
reclaim: vi.fn(),
reconcile: vi.fn(),
reconcileActive: vi.fn(),
} as unknown as DispatchService;
const coordinated = coordinateWorkerPlacementDispatch(service);
const first = coordinated.dispatch(REQUEST);
await dispatchStarted.promise;
await expect(
coordinated.dispatch({ ...REQUEST, profileId: "another-profile" }),
).rejects.toThrow(`Session ${REQUEST.sessionKey} is already dispatching another request`);
await expect(
coordinated.dispatch({
...REQUEST,
inheritedProfile: {
providerId: "fake",
profileSnapshot: { settings: { region: "parent" } },
},
}),
).rejects.toThrow(`Session ${REQUEST.sessionKey} is already dispatching another request`);
const retry = coordinated.dispatch(REQUEST);
releaseDispatch.resolve();
const [firstResult, retryResult] = await Promise.all([first, retry]);
expect(retryResult).toBe(firstResult);
expect(dispatch).toHaveBeenCalledOnce();
await coordinated.dispatch({ ...REQUEST, profileId: "another-profile" });
expect(dispatch).toHaveBeenCalledTimes(2);
});
it("joins a retry before a queued reconciliation after dispatch failure", async () => {
const dispatchStarted = createDeferredCore();
const releaseDispatch = createDeferredCore();
const dispatchError = new Error("provision failed");
const dispatch = vi.fn(async () => {
dispatchStarted.resolve();
await releaseDispatch.promise;
throw dispatchError;
});
const reconcileActive = vi.fn();
const service = {
dispatch,
forceDestroyEnvironment: vi.fn(),
reclaim: vi.fn(),
reconcile: vi.fn(),
reconcileActive,
} as unknown as DispatchService;
const coordinated = coordinateWorkerPlacementDispatch(service);
const first = coordinated.dispatch(REQUEST);
await dispatchStarted.promise;
const reconciliation = coordinated.reconcileActive();
const retry = coordinated.dispatch(REQUEST);
const outcomes = Promise.allSettled([first, retry]);
releaseDispatch.resolve();
expect(await outcomes).toEqual([
{ status: "rejected", reason: dispatchError },
{ status: "rejected", reason: dispatchError },
]);
await reconciliation;
expect(dispatch).toHaveBeenCalledOnce();
expect(reconcileActive).toHaveBeenCalledOnce();
await expect(coordinated.dispatch({ ...REQUEST, profileId: "another-profile" })).rejects.toBe(
dispatchError,
);
expect(dispatch).toHaveBeenCalledTimes(2);
});
it("coalesces full sweeps but runs a fresh targeted pass with its environment id", async () => {
const fullSweepStarted = createDeferredCore();
const releaseFullSweep = createDeferredCore();
const reconcileActive = vi.fn(async (environmentId?: string) => {
if (environmentId === undefined) {
fullSweepStarted.resolve();
await releaseFullSweep.promise;
}
});
const service = {
dispatch: vi.fn(),
forceDestroyEnvironment: vi.fn(),
reclaim: vi.fn(),
reconcile: vi.fn(),
reconcileActive,
} as unknown as DispatchService;
const coordinated = coordinateWorkerPlacementDispatch(service);
const firstFullSweep = coordinated.reconcileActive();
const secondFullSweep = coordinated.reconcileActive();
await fullSweepStarted.promise;
const targetedSweep = coordinated.reconcileActive("worker-target");
expect(reconcileActive).toHaveBeenCalledTimes(1);
releaseFullSweep.resolve();
await Promise.all([firstFullSweep, secondFullSweep, targetedSweep]);
expect(reconcileActive.mock.calls).toEqual([[], ["worker-target"]]);
});
});
@@ -0,0 +1,122 @@
import { isDeepStrictEqual } from "node:util";
import type { WorkerPlacementDispatchService } from "./placement-dispatch.js";
import type { WorkerPlacementDispatchRequest } from "./service-contract.js";
/** Serializes reconciliation sweeps against dispatches and deduplicates exact requests. */
export function coordinateWorkerPlacementDispatch(
service: WorkerPlacementDispatchService,
): WorkerPlacementDispatchService {
let activeDispatchCount = 0;
let reconciliation: Promise<void> | undefined;
const dispatchIdleWaiters = new Set<() => void>();
const waitForDispatchIdle = (): Promise<void> => {
if (activeDispatchCount === 0) {
return Promise.resolve();
}
return new Promise<void>((resolve) => {
dispatchIdleWaiters.add(resolve);
});
};
const runReconciliation = (operation: () => Promise<void>): Promise<void> => {
if (reconciliation) {
return reconciliation;
}
const current = (async () => {
await waitForDispatchIdle();
await operation();
})();
reconciliation = current;
const clearCurrent = () => {
if (reconciliation === current) {
reconciliation = undefined;
}
};
void current.then(clearCurrent, clearCurrent);
return current;
};
const runExclusivePlacementOperation = <T>(operation: () => Promise<T>): Promise<T> => {
const current = (async () => {
const pendingReconciliation = reconciliation;
if (pendingReconciliation) {
await pendingReconciliation.catch(() => undefined);
}
await waitForDispatchIdle();
return await operation();
})();
const barrier = current.then(
() => undefined,
() => undefined,
);
reconciliation = barrier;
return current.finally(() => {
if (reconciliation === barrier) {
reconciliation = undefined;
}
});
};
const runPlacementOperation = async <T>(operation: () => Promise<T>): Promise<T> => {
for (;;) {
const pendingReconciliation = reconciliation;
if (!pendingReconciliation) {
break;
}
await pendingReconciliation.catch(() => undefined);
}
activeDispatchCount += 1;
try {
return await operation();
} finally {
activeDispatchCount -= 1;
if (activeDispatchCount === 0) {
const waiters = [...dispatchIdleWaiters];
dispatchIdleWaiters.clear();
for (const resolve of waiters) {
resolve();
}
}
}
};
const dispatchInFlight = new Map<
string,
{
request: WorkerPlacementDispatchRequest;
operation: ReturnType<WorkerPlacementDispatchService["dispatch"]>;
}
>();
return {
dispatch: async (request, onTransition) => {
const inFlight = dispatchInFlight.get(request.sessionId);
if (inFlight) {
if (
inFlight.request.sessionKey !== request.sessionKey ||
inFlight.request.agentId !== request.agentId ||
inFlight.request.profileId !== request.profileId ||
inFlight.request.deviceId !== request.deviceId ||
!isDeepStrictEqual(inFlight.request.inheritedProfile, request.inheritedProfile)
) {
throw new Error(`Session ${request.sessionKey} is already dispatching another request`);
}
return await inFlight.operation;
}
const operation = runPlacementOperation(() => service.dispatch(request, onTransition));
dispatchInFlight.set(request.sessionId, { request, operation });
try {
return await operation;
} finally {
if (dispatchInFlight.get(request.sessionId)?.operation === operation) {
dispatchInFlight.delete(request.sessionId);
}
}
},
forceDestroyEnvironment: (environmentId, onCleanupError) =>
runExclusivePlacementOperation(() =>
service.forceDestroyEnvironment(environmentId, onCleanupError),
),
reclaim: async (request) => await runPlacementOperation(() => service.reclaim(request)),
reconcile: () => runReconciliation(service.reconcile),
reconcileActive: (environmentId) =>
environmentId === undefined
? runReconciliation(() => service.reconcileActive())
: runExclusivePlacementOperation(() => service.reconcileActive(environmentId)),
};
}
@@ -0,0 +1,182 @@
import { describe, expect, it, vi } from "vitest";
import type { WorkerSessionPlacementRecord } from "./placement-record.js";
import { createPlacementSessionRetirement } from "./placement-session-retirement.js";
import type { WorkerSessionPlacementRetirement } from "./placement-store.js";
function localPlacement(
sessionId: string,
): Extract<WorkerSessionPlacementRecord, { state: "local" }> {
return {
sessionId,
sessionKey: `agent:main:${sessionId}`,
agentId: "main",
state: "local",
generation: 1,
turnClaim: null,
environmentId: null,
activeOwnerEpoch: null,
workspaceBaseManifestRef: null,
remoteWorkspaceDir: null,
workerBundleHash: null,
lastTranscriptAckCursor: null,
lastLiveEventAckCursor: null,
recoveryError: null,
terminalReason: null,
terminalAtMs: null,
createdAtMs: 1,
updatedAtMs: 1,
stateChangedAtMs: 1,
};
}
function activePlacement(
sessionId: string,
): Extract<WorkerSessionPlacementRecord, { state: "active" }> {
return {
...localPlacement(sessionId),
state: "active",
generation: 2,
turnClaim: null,
environmentId: `environment:${sessionId}`,
activeOwnerEpoch: 3,
workspaceBaseManifestRef: "manifest",
remoteWorkspaceDir: "/workspace",
workerBundleHash: "a".repeat(64),
};
}
function failedPlacement(
placement: Extract<WorkerSessionPlacementRecord, { state: "active" }>,
): Extract<WorkerSessionPlacementRecord, { state: "failed" }> {
return {
...placement,
state: "failed",
generation: placement.generation + 1,
turnClaim: null,
recoveryError: "forced teardown",
terminalReason: "failed",
terminalAtMs: 2,
};
}
function createHarness(records: WorkerSessionPlacementRecord[]) {
const placements = new Map(records.map((record) => [record.sessionId, record]));
const environments = new Map<
string,
{ environmentId: string; state: "attached" | "destroyed"; leaseId: string | null }
>(
records.flatMap((record) =>
record.environmentId
? [
[
record.environmentId,
{
environmentId: record.environmentId,
state: "attached" as const,
leaseId: `lease:${record.environmentId}`,
},
] as const,
]
: [],
),
);
const retired: WorkerSessionPlacementRetirement[] = [];
const forceDestroyEnvironment = vi.fn(async (environmentId: string) => {
environments.set(environmentId, {
environmentId,
state: "destroyed",
leaseId: null,
});
for (const record of placements.values()) {
if (record.environmentId === environmentId && record.state === "active") {
placements.set(record.sessionId, failedPlacement(record));
}
}
});
const retirement = createPlacementSessionRetirement({
placements: {
get: (sessionId) => placements.get(sessionId),
list: () => [...placements.values()],
retireSessionPlacement: (input) => {
const current = placements.get(input.sessionId);
if (
current?.state !== input.expectedState ||
current.generation !== input.expectedGeneration ||
current.turnClaim
) {
throw new Error("placement changed");
}
placements.delete(input.sessionId);
retired.push(input);
},
},
environments: {
get: (environmentId) => environments.get(environmentId) as never,
},
forceDestroyEnvironment,
resolveSessionEvidence: async () => "absent",
warn: vi.fn(),
});
return { forceDestroyEnvironment, placements, retired, retirement };
}
describe("placement session retirement", () => {
it("retires an exact local placement after its session disappears", async () => {
const harness = createHarness([localPlacement("session-local")]);
await harness.retirement.reconcile();
expect(harness.retired).toEqual([
{
sessionId: "session-local",
expectedState: "local",
expectedGeneration: 1,
},
]);
expect(harness.forceDestroyEnvironment).not.toHaveBeenCalled();
});
it("fences a live environment before retiring its failed placement", async () => {
const harness = createHarness([activePlacement("session-active")]);
await harness.retirement.reconcile();
expect(harness.forceDestroyEnvironment).toHaveBeenCalledWith(
"environment:session-active",
expect.any(Function),
);
expect(harness.retired).toEqual([
{
sessionId: "session-active",
expectedState: "failed",
expectedGeneration: 3,
},
]);
});
it("retains current and unknown session evidence", async () => {
const current = localPlacement("session-current");
const unknown = localPlacement("session-unknown");
const harness = createHarness([current, unknown]);
const retirement = createPlacementSessionRetirement({
placements: {
get: (sessionId) => harness.placements.get(sessionId),
list: () => [...harness.placements.values()],
retireSessionPlacement: () => {
throw new Error("must not retire");
},
},
environments: { get: () => undefined },
forceDestroyEnvironment: async () => {
throw new Error("must not destroy");
},
resolveSessionEvidence: async (placement) =>
placement.sessionId === current.sessionId ? "current" : "unknown",
warn: vi.fn(),
});
await retirement.reconcile();
expect(harness.placements.size).toBe(2);
});
});
@@ -0,0 +1,110 @@
import type { WorkerSessionPlacementRecord } from "./placement-record.js";
import type { WorkerSessionPlacementStore } from "./placement-store.js";
import type { WorkerEnvironmentService } from "./service.js";
import { isFailedWorkerPlacementEnvironmentGone } from "./session-placement-lifecycle.js";
type PlacementSessionEvidence = "current" | "absent" | "unknown";
type PlacementSessionRetirementDeps = {
placements: Pick<WorkerSessionPlacementStore, "get" | "list" | "retireSessionPlacement">;
environments: Pick<WorkerEnvironmentService, "get">;
forceDestroyEnvironment: (
environmentId: string,
onCleanupError?: (error: unknown) => void,
) => Promise<unknown>;
resolveSessionEvidence: (
placement: WorkerSessionPlacementRecord,
) => Promise<PlacementSessionEvidence>;
warn: (message: string) => void;
};
export function createPlacementSessionRetirement(deps: PlacementSessionRetirementDeps) {
const retireCurrent = (placement: WorkerSessionPlacementRecord): boolean => {
if (placement.turnClaim) {
return false;
}
const retirement =
placement.state === "local" || placement.state === "reclaimed"
? {
sessionId: placement.sessionId,
expectedState: placement.state,
expectedGeneration: placement.generation,
}
: placement.state === "failed" &&
isFailedWorkerPlacementEnvironmentGone({
environmentService: deps.environments,
placement,
})
? {
sessionId: placement.sessionId,
expectedState: placement.state,
expectedGeneration: placement.generation,
}
: undefined;
if (!retirement) {
return false;
}
deps.placements.retireSessionPlacement(retirement);
return true;
};
const reconcilePlacement = async (placement: WorkerSessionPlacementRecord): Promise<void> => {
const evidence = await deps.resolveSessionEvidence(placement);
if (evidence !== "absent") {
return;
}
let current = deps.placements.get(placement.sessionId);
if (!current) {
return;
}
try {
if (retireCurrent(current)) {
return;
}
} catch {
return;
}
const environmentId = current.environmentId;
if (!environmentId) {
return;
}
try {
await deps.forceDestroyEnvironment(environmentId, (error) => {
deps.warn(
`Worker placement orphan cleanup deferred for ${current?.sessionId ?? placement.sessionId}: ${String(error)}`,
);
});
} catch (error) {
deps.warn(
`Worker placement orphan teardown failed for ${current.sessionId}: ${String(error)}`,
);
return;
}
current = deps.placements.get(placement.sessionId);
if (!current) {
return;
}
try {
retireCurrent(current);
} catch {
// A concurrent placement transition owns the next reconciliation pass.
}
};
const reconcile = async (): Promise<void> => {
for (const placement of deps.placements.list()) {
try {
await reconcilePlacement(placement);
} catch (error) {
deps.warn(
`Worker placement session evidence check failed for ${placement.sessionId}: ${String(error)}`,
);
}
}
};
return { reconcile };
}