mirror of
https://github.com/openclaw/openclaw.git
synced 2026-08-23 19:08:22 -06:00
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:
committed by
GitHub
parent
bf70d5ddb6
commit
3c5e2ff296
@@ -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 };
|
||||
}
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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 };
|
||||
}
|
||||
Reference in New Issue
Block a user