diff --git a/packages/protocol/src/messages.ts b/packages/protocol/src/messages.ts index ed36eeb1ad..c469592cf5 100644 --- a/packages/protocol/src/messages.ts +++ b/packages/protocol/src/messages.ts @@ -1075,8 +1075,38 @@ export const WorkspaceRecoveryInspectRequestSchema = z.object({ requestId: z.string(), }); +const AgentRequestOperationSchema = z.object({ + key: z.string().min(1).max(512), + deadlineAt: z.string().datetime().optional(), +}); + +export const AgentRequestsCancelRequestSchema = z.object({ + type: z.literal("agent.requests.cancel.request"), + requestId: z.string(), + key: z.string().min(1).max(512), + creationKey: z.string().min(1).max(512), +}); + +export const AgentRequestsInspectRequestSchema = AgentRequestsCancelRequestSchema.extend({ + type: z.literal("agent.requests.inspect.request"), +}); + +export const AgentRequestsCancelResponseSchema = z.object({ + type: z.literal("agent.requests.cancel.response"), + payload: z.object({ + requestId: z.string(), + agentId: z.string().nullable(), + outcome: z.enum(["settled", "pending", "unknown"]), + }), +}); + +export const AgentRequestsInspectResponseSchema = AgentRequestsCancelResponseSchema.extend({ + type: z.literal("agent.requests.inspect.response"), +}); + export const WorkspaceRecoveryRestoreRequestSchema = z.object({ type: z.literal("workspace.recovery.restore.request"), + operation: AgentRequestOperationSchema.optional(), workspaceId: z.string(), requestId: z.string(), }); @@ -1394,6 +1424,7 @@ export const FetchAgentRequestMessageSchema = z.object({ export const SendAgentMessageRequestSchema = z.object({ type: z.literal("send_agent_message_request"), + operation: AgentRequestOperationSchema.optional(), requestId: z.string(), /** Accepts full ID, unique prefix, or exact full title (server resolves). */ agentId: z.string(), @@ -1680,6 +1711,7 @@ export type CreateAgentWorktreeTarget = z.infer { + const client = new HeldAgentCreationClient(); + const manager = new AgentManager({ clients: { codex: client }, logger }); + const controller = new AbortController(); + let outcome = "pending"; + const creation = manager + .createAgent({ provider: "codex", cwd: process.cwd() }, undefined, { + workspaceId: undefined, + signal: controller.signal, + }) + .then( + () => { + outcome = "created"; + return; + }, + () => { + outcome = "canceled"; + }, + ); + await client.waitForCreationToStart(); + try { + controller.abort(new Error("startup expired")); + await expect.poll(() => outcome, { timeout: 500 }).toBe("canceled"); + client.finishCreating(); + await expect.poll(() => client.createdSessionClosed).toBe(true); + expect(manager.listAgents()).toEqual([]); + } finally { + client.finishCreating(); + await creation; + manager.prepareForShutdown(); + await manager.flushForShutdown(); + } +}); + test("does not register a session that finishes starting after shutdown begins", async () => { const client = new HeldAgentCreationClient(); const manager = new AgentManager({ diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index d5dc1ab52b..6630324d46 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -278,6 +278,7 @@ type ProviderEnabledMap = Partial>; type ProviderClientMap = Partial>; export interface CreateAgentOptions { + signal?: AbortSignal; labels?: Record; initialPrompt?: string; env?: Record; @@ -1227,8 +1228,16 @@ export class AgentManager { ); const providerLaunchConfig = this.resolveProviderLaunchConfig(launchConfig, launchContext); const createOptions = this.buildCreateSessionOptions(options); - const session = await client.createSession(providerLaunchConfig, launchContext, createOptions); + options.signal?.throwIfAborted(); + const session = await this.waitForCreatedSession( + client.createSession(providerLaunchConfig, launchContext, createOptions), + options.signal, + ); await this.requireExternalMcpSupport(session, storedConfig); + if (options.signal?.aborted) { + await this.closeUnregisteredSession(session); + options.signal.throwIfAborted(); + } const agent = await this.registerSession(session, storedConfig, resolvedAgentId, { labels: options.labels, initialTitle: options.initialTitle, @@ -1244,6 +1253,38 @@ export class AgentManager { return agent; } + private async waitForCreatedSession( + operation: Promise, + signal?: AbortSignal, + ): Promise { + if (!signal) return operation; + return new Promise((accept, reject) => { + const abort = () => { + signal.removeEventListener("abort", abort); + reject(signal.reason); + }; + signal.addEventListener("abort", abort, { once: true }); + if (signal.aborted) abort(); + void operation + .then( + async (session) => { + signal.removeEventListener("abort", abort); + if (signal.aborted) { + await this.closeUnregisteredSession(session); + } else accept(session); + return; + }, + (error: unknown) => { + signal.removeEventListener("abort", abort); + reject(error); + }, + ) + .catch((error: unknown) => + this.logger.error({ err: error }, "Failed to close canceled provider creation"), + ); + }); + } + private buildCreateSessionOptions(options?: { persistSession?: boolean; }): AgentCreateSessionOptions | undefined { diff --git a/packages/server/src/server/agent/agent-prompt.ts b/packages/server/src/server/agent/agent-prompt.ts index 96f5ce5ee3..a87b631303 100644 --- a/packages/server/src/server/agent/agent-prompt.ts +++ b/packages/server/src/server/agent/agent-prompt.ts @@ -182,6 +182,7 @@ export function isSystemInjectedEnvelope(text: string): boolean { } export interface SendPromptToAgentParams { + signal?: AbortSignal; agentManager: AgentManager; agentStorage: AgentStorage; agentId: string; @@ -204,6 +205,7 @@ export interface SendPromptToAgentParams { } export interface StartCreatedAgentInitialPromptParams { + signal?: AbortSignal; agentManager: AgentManager; agentId: string; snapshot?: ManagedAgent; @@ -229,6 +231,7 @@ const AGENT_RUN_START_TIMEOUT_MS = 60_000; export async function waitForAgentRunStartWithTimeout( agentManager: AgentManager, agentId: string, + signal?: AbortSignal, ): Promise { const provider = agentManager.getAgent(agentId)?.provider ?? "provider"; const startAbort = new AbortController(); @@ -243,7 +246,9 @@ export async function waitForAgentRunStartWithTimeout( ); try { - await agentManager.waitForAgentRunStart(agentId, { signal: startAbort.signal }); + await agentManager.waitForAgentRunStart(agentId, { + signal: signal ? AbortSignal.any([signal, startAbort.signal]) : startAbort.signal, + }); } finally { clearTimeout(startTimeout); } @@ -263,6 +268,7 @@ export async function waitForAgentRunStartWithTimeout( export async function sendPromptToAgent( params: SendPromptToAgentParams, ): Promise<{ disposition: PromptDispatchDisposition }> { + params.signal?.throwIfAborted(); const unarchive = params.unarchive ?? true; const record = await params.agentStorage.get(params.agentId); @@ -283,6 +289,7 @@ export async function sendPromptToAgent( await params.agentManager.setAgentMode(params.agentId, params.sessionMode); } + params.signal?.throwIfAborted(); const runOptions = params.messageId ? { ...params.runOptions, clientMessageId: params.messageId } : params.runOptions; @@ -307,6 +314,7 @@ export async function startCreatedAgentInitialPrompt( return currentSnapshot; } + params.signal?.throwIfAborted(); const dispatchResult = await startAgentRun( params.agentManager, params.agentId, @@ -318,7 +326,7 @@ export async function startCreatedAgentInitialPrompt( ); if (dispatchResult.disposition === "turn_started") { - await waitForAgentRunStartWithTimeout(params.agentManager, params.agentId); + await waitForAgentRunStartWithTimeout(params.agentManager, params.agentId, params.signal); } const refreshedSnapshot = params.agentManager.getAgent(params.agentId) ?? params.snapshot ?? null; diff --git a/packages/server/src/server/agent/create-agent/create.ts b/packages/server/src/server/agent/create-agent/create.ts index bd4dcb16d8..86e7f4a635 100644 --- a/packages/server/src/server/agent/create-agent/create.ts +++ b/packages/server/src/server/agent/create-agent/create.ts @@ -55,6 +55,7 @@ export type EnsureWorkspaceForCreate = ( export interface CreateAgentFromSessionInput { kind: "session"; + signal?: AbortSignal; agentId?: string; config: AgentSessionConfig; workspaceId: string; @@ -79,6 +80,7 @@ export interface CreateAgentFromSessionInput { export interface CreateAgentFromMcpInput { kind: "mcp"; + signal?: AbortSignal; provider: string; title: string; initialPrompt?: string; @@ -180,10 +182,12 @@ export async function createAgentCommand( ? await resolveSessionCreateAgent(dependencies, input) : await resolveMcpCreateAgent(dependencies, input); + const signal = input.signal; + signal?.throwIfAborted(); const snapshot = await dependencies.agentManager.createAgent( resolved.config, input.kind === "session" ? input.agentId : undefined, - resolved.createOptions, + { ...resolved.createOptions, signal }, ); resolved.setupContinuation?.startAfterAgentCreate({ @@ -197,7 +201,7 @@ export async function createAgentCommand( input.onCreated?.({ agentId: snapshot.id, createdWorktree: resolved.createdWorktree ?? null }); } if (resolved.prompt !== undefined) { - const sendResult = await sendInitialPrompt(dependencies, resolved, snapshot); + const sendResult = await sendInitialPrompt(dependencies, resolved, snapshot, signal); initialPromptStarted = sendResult.started; liveSnapshot = sendResult.liveSnapshot; initialPromptError = sendResult.error ?? null; @@ -449,6 +453,7 @@ async function sendInitialPrompt( dependencies: CreateAgentCommandDependencies, resolved: ResolvedCreateAgent, snapshot: ManagedAgent, + signal?: AbortSignal, ): Promise<{ started: boolean; liveSnapshot: ManagedAgent; error?: unknown }> { try { const prompt = resolved.prompt; @@ -458,6 +463,7 @@ async function sendInitialPrompt( const liveSnapshot = await startCreatedAgentInitialPrompt({ agentManager: dependencies.agentManager, agentId: snapshot.id, + signal, snapshot, prompt, runOptions: resolved.runOptions, diff --git a/packages/server/src/server/agent/requests/index.test.ts b/packages/server/src/server/agent/requests/index.test.ts index de3b010148..d066cda8a4 100644 --- a/packages/server/src/server/agent/requests/index.test.ts +++ b/packages/server/src/server/agent/requests/index.test.ts @@ -141,3 +141,194 @@ test("failed local message preparation does not leave an ambiguous receipt", asy await requests.send(input); expect(sends).toBe(1); }); + +test("canceling startup rejects its waiter and fences later phases and retries", async () => { + const { requests, directory } = await fixture(); + const context = { key: "arrival", deadlineAt: new Date(Date.now() + 60_000).toISOString() }; + let entered!: () => void; + const started = new Promise((resolve) => { + entered = resolve; + }); + const pending = requests.run(context, async (signal) => { + if (!signal) throw new Error("missing operation signal"); + entered(); + await new Promise((_, reject) => + signal.addEventListener("abort", () => reject(signal.reason), { once: true }), + ); + }); + const failed = expect(pending).rejects.toThrow("agent_request_canceled"); + await started; + await requests.cancel(context.key); + await failed; + let sent = false; + await expect( + new AgentRequests(directory).run(context, async () => { + sent = true; + }), + ).rejects.toThrow("agent_request_canceled"); + expect(sent).toBe(false); + expect(await requests.run({ ...context, key: "other-arrival" }, async () => "ok")).toBe("ok"); +}); + +test("an expired startup deadline cannot launch even without a cancel frame", async () => { + const { requests } = await fixture(); + let launched = false; + await expect( + requests.run({ key: "expired", deadlineAt: "2000-01-01T00:00:00.000Z" }, async () => { + launched = true; + }), + ).rejects.toThrow("agent_request_canceled"); + expect(launched).toBe(false); +}); + +test("a noncooperative operation leaves explicit pending cleanup without trapping cancel callers", async () => { + const { requests } = await fixture(); + let entered!: () => void; + const started = new Promise((resolve) => { + entered = resolve; + }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const pending = requests.run( + { key: "hung", deadlineAt: new Date(Date.now() + 60_000).toISOString() }, + async () => { + entered(); + await gate; + }, + ); + const failed = expect(pending).rejects.toThrow("agent_request_canceled"); + await started; + expect(await requests.cancelOperation("hung", "conversation")).toEqual({ + agentId: null, + outcome: "pending", + }); + await failed; + release(); + await expect + .poll(() => requests.cancelOperation("hung", "conversation")) + .toEqual({ agentId: null, outcome: "settled" }); +}); + +test("a reconstructed journal reports interrupted startup as unknown and never launches a canceled phase", async () => { + const { requests, directory } = await fixture(); + let entered!: () => void; + const started = new Promise((resolve) => { + entered = resolve; + }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const context = { key: "interrupted", deadlineAt: new Date(Date.now() + 60_000).toISOString() }; + const pending = requests.run(context, async () => { + entered(); + await gate; + }); + await started; + const restarted = new AgentRequests(directory); + const replays = await Promise.allSettled([ + restarted.run(context, async () => "must not launch"), + restarted.run(context, async () => "must not launch"), + ]); + expect(replays).toEqual([ + { status: "rejected", reason: new Error("agent_request_outcome_unknown") }, + { status: "rejected", reason: new Error("agent_request_outcome_unknown") }, + ]); + expect(await restarted.cancelOperation(context.key, "conversation")).toEqual({ + agentId: null, + outcome: "unknown", + }); + await expect(restarted.run(context, async () => "must not run")).rejects.toThrow( + "agent_request_outcome_unknown", + ); + release(); + await pending; +}); + +test("a live deadline settles a hung waiter without a cancel frame and preserves late cleanup", async () => { + const { requests } = await fixture(); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + let entered = false; + const pending = requests.run( + { key: "deadline", deadlineAt: new Date(Date.now() + 200).toISOString() }, + async () => { + entered = true; + await gate; + }, + ); + await expect(pending).rejects.toThrow("agent_request_canceled"); + expect(entered).toBe(true); + expect(await requests.inspectOperation("deadline", "conversation")).toMatchObject({ + outcome: "pending", + }); + await expect(requests.run({ key: "deadline" }, async () => "must not launch")).rejects.toThrow( + "agent_request_canceled", + ); + release(); + await expect + .poll(() => requests.inspectOperation("deadline", "conversation")) + .toMatchObject({ outcome: "settled" }); + expect(await requests.run({ key: "next" }, async () => "ready")).toBe("ready"); +}); + +test("retransmitted phases serialize while other operation keys continue", async () => { + const { requests } = await fixture(); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + let entered!: () => void; + const started = new Promise((resolve) => { + entered = resolve; + }); + const order: string[] = []; + const first = requests.run({ key: "arrival" }, async () => { + order.push("first"); + entered(); + await gate; + order.push("settled"); + }); + await started; + const second = requests.run({ key: "arrival" }, async () => { + order.push("second"); + }); + await requests.run({ key: "other" }, async () => { + order.push("other"); + }); + expect(order).toEqual(["first", "other"]); + release(); + await Promise.all([first, second]); + expect(order).toEqual(["first", "other", "settled", "second"]); +}); + +test("successful control receipts survive restart and reject changed targets", async () => { + const { requests, directory } = await fixture(); + let archives = 0; + const action = async () => { + archives++; + }; + await requests.control({ key: "cleanup" }, { workspaceId: "one" }, action); + const restarted = new AgentRequests(directory); + await restarted.control({ key: "cleanup" }, { workspaceId: "one" }, action); + expect(archives).toBe(1); + await expect( + restarted.control({ key: "cleanup" }, { workspaceId: "two" }, action), + ).rejects.toThrow("agent_request_key_conflict"); +}); + +test("refused provider cancellation remains unknown after restart", async () => { + const { requests, directory } = await fixture(); + await expect( + requests.run({ key: "refused" }, async () => { + throw new Error("agent_request_outcome_unknown"); + }), + ).rejects.toThrow("agent_request_outcome_unknown"); + expect( + await new AgentRequests(directory).inspectOperation("refused", "conversation"), + ).toMatchObject({ outcome: "unknown" }); +}); diff --git a/packages/server/src/server/agent/requests/index.ts b/packages/server/src/server/agent/requests/index.ts index ad64331c35..dfb0a61184 100644 --- a/packages/server/src/server/agent/requests/index.ts +++ b/packages/server/src/server/agent/requests/index.ts @@ -13,10 +13,218 @@ type Receipt = z.infer; /** One daemon-owned request journal, shared by all of its socket sessions. */ export class AgentRequests { - private readonly pending = new Map>(); + private readonly pending = new Map>(); + private readonly operations = new Map< + string, + { controllers: Set; initialized: Promise; unknown: boolean } + >(); + private readonly operationWrites = new Map>(); constructor(private readonly directory: string) {} + /** Runtime cancellation belongs to the same journal as idempotent create/send receipts. */ + async run( + context: { key: string; deadlineAt?: string } | undefined, + operation: (signal?: AbortSignal) => Promise, + ): Promise { + if (!context) return operation(); + let active = this.operations.get(context.key); + if (!active) { + const previousWrite = this.operationWrites.get(context.key); + active = { + controllers: new Set(), + unknown: false, + initialized: (async () => { + await previousWrite; + if (await this.wasInterrupted(context.key)) + throw new Error("agent_request_outcome_unknown"); + })(), + }; + this.operations.set(context.key, active); + } + const operationState = active; + const controller = new AbortController(); + const controllers = active.controllers; + controllers.add(controller); + let timer: ReturnType | undefined; + let abort: (() => void) | undefined; + let started = false; + let recorded = false; + try { + await active.initialized; + const remaining = + context.deadlineAt === undefined ? undefined : Date.parse(context.deadlineAt) - Date.now(); + if ( + (remaining !== undefined && (!Number.isFinite(remaining) || remaining <= 0)) || + (await this.wasCanceled(context.key)) + ) { + throw new Error("agent_request_canceled"); + } + controller.signal.throwIfAborted(); + if (remaining !== undefined) + timer = setTimeout(() => controller.abort(new Error("agent_request_canceled")), remaining); + await this.recordOperationState(context.key); + recorded = true; + controller.signal.throwIfAborted(); + started = true; + const operationKey = digest(["operation", context.key]); + const previous = this.pending.get(operationKey) ?? Promise.resolve(); + const tracked = previous + .catch(() => undefined) + .then(() => { + controller.signal.throwIfAborted(); + return operation(controller.signal); + }) + .catch((error: unknown) => { + if (error instanceof Error && error.message === "agent_request_outcome_unknown") + operationState.unknown = true; + throw error; + }) + .finally(async () => { + controllers.delete(controller); + if (controllers.size === 0) this.operations.delete(context.key); + await this.recordOperationState(context.key, operationState.unknown); + if (this.pending.get(operationKey) === tracked) this.pending.delete(operationKey); + }); + this.pending.set(operationKey, tracked); + return await Promise.race([ + tracked, + new Promise((_, reject) => { + abort = () => reject(controller.signal.reason); + controller.signal.addEventListener("abort", abort, { once: true }); + if (controller.signal.aborted) abort(); + }), + ]); + } finally { + clearTimeout(timer); + if (abort) controller.signal.removeEventListener("abort", abort); + if (controller.signal.aborted) await this.recordCancellation(context.key); + if (!started) { + controllers.delete(controller); + if (controllers.size === 0) this.operations.delete(context.key); + if (recorded) await this.recordOperationState(context.key); + } + } + } + + async cancel(key: string): Promise { + // Persist before acknowledging: a delayed request or a new socket must also be fenced. + await this.recordCancellation(key); + for (const controller of this.operations.get(key)?.controllers ?? []) { + controller.abort(new Error("agent_request_canceled")); + } + } + + async control( + context: { key: string; deadlineAt?: string } | undefined, + request: unknown, + operation: () => Promise, + ): Promise { + await this.run(context, async () => { + if (!context) { + await operation(); + return; + } + const file = path.join(this.directory, `control-${digest(context.key)}.json`); + const fingerprint = digest(request); + try { + const receipt = z + .object({ fingerprint: z.string() }) + .parse(JSON.parse(await readFile(file, "utf8"))); + if (receipt.fingerprint !== fingerprint) throw new Error("agent_request_key_conflict"); + return; + } catch (error) { + if (!(error instanceof Error && "code" in error && error.code === "ENOENT")) throw error; + } + await operation(); + await writeJsonFileAtomic(file, { fingerprint }); + }); + } + + async cancelAndSettle(key: string): Promise { + await this.cancel(key); + await this.pending.get(digest(["operation", key]))?.catch(() => undefined); + await this.operationWrites.get(key); + if (await this.wasInterrupted(key)) throw new Error("agent_request_outcome_unknown"); + } + + async cancelOperation( + key: string, + creationKey: string, + ): Promise<{ + agentId: string | null; + outcome: "settled" | "pending" | "unknown"; + }> { + await this.cancel(key); + return this.inspectOperation(key, creationKey); + } + + async inspectOperation( + key: string, + creationKey: string, + ): Promise<{ + agentId: string | null; + outcome: "settled" | "pending" | "unknown"; + }> { + await this.operationWrites.get(key); + const active = this.operations.has(key); + const receipt = await readReceipt( + path.join(this.directory, `${digest(["create", creationKey])}.json`), + ); + let outcome: "pending" | "unknown" | "settled" = "settled"; + if (receipt?.state === "pending" || (await this.wasInterrupted(key))) outcome = "unknown"; + if (active) outcome = "pending"; + return { agentId: receipt?.agentId ?? null, outcome }; + } + + private recordOperationState(key: string, unknown = false): Promise { + const previous = this.operationWrites.get(key) ?? Promise.resolve(); + const write = previous + .catch(() => undefined) + .then(() => + writeJsonFileAtomic(path.join(this.directory, `operation-${digest(key)}.json`), { + state: unknown || this.operations.has(key) ? "pending" : "settled", + }), + ); + this.operationWrites.set(key, write); + void write + .finally(() => { + if (this.operationWrites.get(key) === write) this.operationWrites.delete(key); + }) + .catch(() => undefined); + return write; + } + + private async wasInterrupted(key: string): Promise { + try { + const state = z + .object({ state: z.enum(["pending", "settled"]) }) + .parse( + JSON.parse( + await readFile(path.join(this.directory, `operation-${digest(key)}.json`), "utf8"), + ), + ); + return state.state === "pending"; + } catch (error) { + if (error instanceof Error && "code" in error && error.code === "ENOENT") return false; + throw error; + } + } + + private recordCancellation(key: string): Promise { + return writeJsonFileAtomic(path.join(this.directory, `canceled-${digest(key)}.json`), true); + } + + private async wasCanceled(key: string): Promise { + try { + await readFile(path.join(this.directory, `canceled-${digest(key)}.json`)); + return true; + } catch (error) { + if (error instanceof Error && "code" in error && error.code === "ENOENT") return false; + throw error; + } + } + create(input: { key: string; request: unknown; diff --git a/packages/server/src/server/authorization/operation-permissions.ts b/packages/server/src/server/authorization/operation-permissions.ts index ec6c4db2dd..75896fac4d 100644 --- a/packages/server/src/server/authorization/operation-permissions.ts +++ b/packages/server/src/server/authorization/operation-permissions.ts @@ -6,6 +6,8 @@ type OutboundOperation = SessionOutboundMessage["type"]; export type PermissionRequirement = DaemonPermission | readonly DaemonPermission[] | null; const INBOUND_PERMISSION = { + "agent.requests.inspect.request": ["workspace.write", "hub.execute"], + "agent.requests.cancel.request": ["workspace.write", "hub.execute"], abort_request: "workspace.write", "agent.config.apply.request": ["workspace.write", "hub.execute"], "agent.detach.request": "workspace.write", @@ -206,6 +208,8 @@ const INBOUND_PERMISSION = { } as const satisfies Record; const OUTBOUND_PERMISSION = { + "agent.requests.inspect.response": ["workspace.write", "hub.execute"], + "agent.requests.cancel.response": ["workspace.write", "hub.execute"], activity_log: "workspace.read", "agent.config.apply.response": ["workspace.write", "hub.execute"], "agent.detach.response": "workspace.write", diff --git a/packages/server/src/server/bootstrap.ts b/packages/server/src/server/bootstrap.ts index 81aee8aa7f..32ff6bd052 100644 --- a/packages/server/src/server/bootstrap.ts +++ b/packages/server/src/server/bootstrap.ts @@ -1,3 +1,4 @@ +import { AgentRequests } from "./agent/requests/index.js"; import { describeHookWorkspace } from "./plugins/lifecycle/index.js"; import express from "express"; import { createServer as createHTTPServer, type IncomingMessage, type ServerResponse } from "http"; @@ -856,6 +857,7 @@ export async function createPaseoDaemon( serviceProxyListenTarget = parseListenString(config.serviceProxy.standaloneListen); } + const agentRequests = new AgentRequests(path.join(config.paseoHome, "agent-requests")); const agentStorage = new AgentStorage(config.agentStoragePath, logger); const projectRegistry = new FileBackedProjectRegistry( path.join(config.paseoHome, "projects", "projects.json"), @@ -1257,6 +1259,7 @@ export async function createPaseoDaemon( createExecutionAgents: (daemonId) => new DaemonExecutions({ daemonId, + agentRequests, agentManager, agentStorage, createAgent, @@ -1717,6 +1720,7 @@ export async function createPaseoDaemon( pluginRuntime, orchestrationSkills, workspaceLabelService, + agentRequests, ); pluginRuntime.bindPaseoSessionHost(wsServer); await pluginRuntime.start(); diff --git a/packages/server/src/server/hub/daemon-executions.test.ts b/packages/server/src/server/hub/daemon-executions.test.ts index a96dc8cfdd..33fcb94d2b 100644 --- a/packages/server/src/server/hub/daemon-executions.test.ts +++ b/packages/server/src/server/hub/daemon-executions.test.ts @@ -326,3 +326,45 @@ test("failed create never archives a reused worktree", async () => { }); expect(await hub.worktreeState(worktreeCwd!)).toEqual({ exists: true, listed: true }); }); + +test("control cancels a held legacy provider create and fences a delayed replay", async () => { + const hub = await launchRelationship(); + hub.holdAgentCreation(); + hub.beginOwnedCreate("held-create", "expired-create"); + await hub.agentCreationAttempts(1); + const creation = hub.ownedCreateResult("held-create"); + let settled = false; + const control = hub.controlExecution("expired-create", "archive").then((result) => { + settled = true; + return result; + }); + try { + await expect.poll(() => settled, { timeout: 500 }).toBe(true); + expect(await control).toMatchObject({ success: true }); + expect(await creation).toMatchObject({ payload: { success: false } }); + } finally { + hub.finishAgentCreation(); + } + hub.beginOwnedCreate("late-create", "expired-create"); + expect(await hub.ownedCreateResult("late-create")).toMatchObject({ payload: { success: false } }); +}); + +test("the legacy absolute deadline cancels provider creation without a control frame", async () => { + const hub = await launchRelationship(); + hub.holdAgentCreation(); + hub.beginOwnedCreate("deadline-create", "deadline-create", { + deadlineAt: new Date(Date.now() + 1_000).toISOString(), + }); + try { + await hub.agentCreationAttempts(1); + expect(await hub.ownedCreateResult("deadline-create")).toMatchObject({ + payload: { success: false }, + }); + } finally { + hub.finishAgentCreation(); + } + hub.beginOwnedCreate("deadline-replay", "deadline-create"); + expect(await hub.ownedCreateResult("deadline-replay")).toMatchObject({ + payload: { success: false }, + }); +}); diff --git a/packages/server/src/server/hub/daemon-executions.ts b/packages/server/src/server/hub/daemon-executions.ts index a619be18a8..09f8f12538 100644 --- a/packages/server/src/server/hub/daemon-executions.ts +++ b/packages/server/src/server/hub/daemon-executions.ts @@ -1,3 +1,4 @@ +import type { AgentRequests } from "../agent/requests/index.js"; import type { AgentSnapshotPayload, AgentStreamEventPayload, @@ -17,6 +18,7 @@ import { daemonExecutionKey, type DaemonAgentOwner } from "../agent/agent-owner. export interface HubExecutionAgentCreateInput { executionId: string; + deadlineAt?: string; provider: string; cwd: string; prompt: string; @@ -53,6 +55,7 @@ export type OwnedAgentEvent = interface DaemonExecutionsOptions { daemonId: string; + agentRequests: AgentRequests; agentManager: AgentManager; agentStorage: AgentStorage; createAgent: BoundCreateAgentCommand; @@ -122,6 +125,8 @@ export class DaemonExecutions implements HubExecutionAgents { const pending = this.pendingControlActions.get(actionKey); if (pending) return pending; + const canceled = this.options.agentRequests.cancelAndSettle(executionKey); + void canceled.catch(() => undefined); const previous = this.controlTails.get(executionKey) ?? this.pendingCreates.get(executionKey)?.then(() => undefined) ?? @@ -129,6 +134,7 @@ export class DaemonExecutions implements HubExecutionAgents { const authorityGeneration = this.authorityGeneration; const control = previous .catch(() => undefined) + .then(() => canceled) .then(() => this.controlOwnedExecution(owner, input, authorityGeneration)); this.pendingControlActions.set(actionKey, control); this.controlTails.set(executionKey, control); @@ -148,6 +154,9 @@ export class DaemonExecutions implements HubExecutionAgents { this.authorityActive = false; this.authorityGeneration++; await Promise.allSettled([ + ...[...this.pendingCreates.keys()].map((key) => + this.options.agentRequests.cancelAndSettle(key), + ), ...this.pendingCreates.values(), ...this.pendingControlActions.values(), ]); @@ -180,12 +189,25 @@ export class DaemonExecutions implements HubExecutionAgents { requireHubMcpNamespace(input.mcpServers); requireToolPolicyServers(input.toolPolicy, input.mcpServers); + return this.options.agentRequests.run( + { key: daemonExecutionKey(owner), deadlineAt: input.deadlineAt }, + (signal) => this.createNewOwnedAgent(owner, input, authorityGeneration, signal), + ); + } + + private async createNewOwnedAgent( + owner: DaemonAgentOwner, + input: HubExecutionAgentCreateInput, + authorityGeneration: number, + signal?: AbortSignal, + ): Promise { let createdWorktree: CreatePaseoWorktreeWorkflowResult | null = null; let createdAgentId: string | null = null; let result: Awaited>; try { result = await this.createAgentCommand({ kind: "mcp", + signal, provider: input.model ? `${input.provider}/${input.model}` : input.provider, title: input.prompt, initialPrompt: input.prompt, @@ -215,6 +237,7 @@ export class DaemonExecutions implements HubExecutionAgents { createdAgentId = created.agentId; }, }); + signal?.throwIfAborted(); this.requireAuthority(authorityGeneration); requireExecutionWorkspaceId(result.liveSnapshot); } catch (error) { diff --git a/packages/server/src/server/hub/execution-controller.ts b/packages/server/src/server/hub/execution-controller.ts index 025b05a48b..32e8b94acc 100644 --- a/packages/server/src/server/hub/execution-controller.ts +++ b/packages/server/src/server/hub/execution-controller.ts @@ -149,6 +149,7 @@ export class HubExecutionController { if (!isAbsolute(message.cwd)) throw new Error("Hub agent cwd must be absolute"); const result = await this.agents.create({ executionId: message.executionId, + deadlineAt: message.deadlineAt, provider: message.provider, cwd: message.cwd, prompt: message.prompt, diff --git a/packages/server/src/server/hub/execution-session.websocket.test.ts b/packages/server/src/server/hub/execution-session.websocket.test.ts index 6e9f4036ed..6697846e60 100644 --- a/packages/server/src/server/hub/execution-session.websocket.test.ts +++ b/packages/server/src/server/hub/execution-session.websocket.test.ts @@ -251,7 +251,7 @@ test("Hub interrupts an owned running execution idempotently", async () => { expect(hub.ownedAgentIsRunning(created.payload.agentId!)).toBe(false); }); -test("Hub control waits for an in-flight create of the same execution", async () => { +test("Hub control cancels an in-flight create of the same execution", async () => { const hub = await launchRelationship(); hub.holdAgentCreation(); hub.beginOwnedCreate("pending-control-create", "execution-pending-control", { @@ -264,10 +264,8 @@ test("Hub control waits for an in-flight create of the same execution", async () const created = await hub.ownedCreateResult("pending-control-create"); const archived = await hub.executionControlResult("pending-control-archive"); - expect(created).toMatchObject({ payload: { success: true, agentId: expect.any(String) } }); + expect(created).toMatchObject({ payload: { success: false, agentId: null } }); expect(archived).toMatchObject({ success: true, error: null, action: "archive" }); - expect(created.payload.agent?.workspaceId).toEqual(expect.any(String)); - expect(await hub.ownedAgentArchivedAt(created.payload.agentId!)).toEqual(expect.any(String)); }); test("Hub archives an execution workspace on a local checkout", async () => { diff --git a/packages/server/src/server/hub/test-utils/relationship-harness.ts b/packages/server/src/server/hub/test-utils/relationship-harness.ts index 3113b543ad..57d6d57738 100644 --- a/packages/server/src/server/hub/test-utils/relationship-harness.ts +++ b/packages/server/src/server/hub/test-utils/relationship-harness.ts @@ -1,3 +1,4 @@ +import { AgentRequests } from "../../agent/requests/index.js"; import { execFile } from "node:child_process"; import { execFileSync } from "node:child_process"; import { EventEmitter } from "node:events"; @@ -1558,6 +1559,7 @@ export class HubRelationshipHarness { private executionsForReconstruction(manager: AgentManager, storage: AgentStorage) { return new DaemonExecutions({ + agentRequests: new AgentRequests(path.join(this.paseoHome, "agent-requests")), daemonId: this.relationshipFile()!.relationship.daemonId, agentManager: manager, agentStorage: storage, diff --git a/packages/server/src/server/session.agent-startup.e2e.test.ts b/packages/server/src/server/session.agent-startup.e2e.test.ts new file mode 100644 index 0000000000..06456d648b --- /dev/null +++ b/packages/server/src/server/session.agent-startup.e2e.test.ts @@ -0,0 +1,255 @@ +import { randomUUID } from "node:crypto"; +import { expect, test } from "vitest"; +import { WebSocket } from "ws"; +import { z } from "zod"; +import { createTestPaseoDaemon } from "./test-utils/paseo-daemon.js"; +import { createTestAgentClient } from "./test-utils/fake-agent-client.js"; + +const Envelope = z.object({ + type: z.literal("session"), + message: z.object({ type: z.string(), payload: z.record(z.string(), z.unknown()) }), +}); + +function exchange( + socket: WebSocket, + frame: unknown, + accepts: (type: string, payload: Record) => boolean, +) { + return new Promise>((resolve, reject) => { + const timer = setTimeout(() => { + cleanup(); + reject(new Error("Daemon response missing")); + }, 10_000); + const receive = (data: WebSocket.RawData) => { + const parsed = Envelope.safeParse(JSON.parse(data.toString())); + if (parsed.success && accepts(parsed.data.message.type, parsed.data.message.payload)) { + cleanup(); + resolve(parsed.data.message.payload); + } + }; + const cleanup = () => { + clearTimeout(timer); + socket.off("message", receive); + }; + socket.on("message", receive); + socket.send(JSON.stringify(frame)); + }); +} +function request(socket: WebSocket, message: Record) { + const requestId = randomUUID(); + return exchange( + socket, + { type: "session", message: { ...message, requestId } }, + (_type, payload) => payload["requestId"] === requestId, + ); +} +async function connect(port: number) { + const socket = new WebSocket(`ws://127.0.0.1:${port}/ws`); + await new Promise((resolve, reject) => { + socket.once("open", resolve); + socket.once("error", reject); + }); + const info = await exchange( + socket, + { type: "hello", clientId: randomUUID(), clientType: "browser", protocolVersion: 1 }, + (_type, payload) => payload["status"] === "server_info", + ); + expect(info["features"]).toMatchObject({ agentRequestCancellation: true }); + return socket; +} + +test("socket cancellation settles a held create, disposes its late provider session, and fences a reconnect replay", async () => { + let started!: () => void; + const entered = new Promise((resolve) => { + started = resolve; + }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + let creates = 0; + let closes = 0; + const provider = createTestAgentClient("codex", { + beforeCreateSession: async () => { + creates++; + started(); + await gate; + }, + closeSession: async () => { + closes++; + }, + }); + const host = await createTestPaseoDaemon({ agentClients: { codex: provider } }); + let socket = await connect(host.port); + const key = randomUUID(); + const creationKey = randomUUID(); + const config = { provider: "codex", cwd: host.paseoHome }; + const operation = { key, deadlineAt: new Date(Date.now() + 60_000).toISOString() }; + try { + const unkeyed = await request(socket, { + type: "create_agent_request", + config, + operation, + initialPrompt: "must not launch", + }); + expect(unkeyed).toMatchObject({ status: "agent_create_failed" }); + expect(creates).toBe(0); + const creation = request(socket, { + type: "create_agent_request", + config, + idempotencyKey: creationKey, + operation, + }); + await entered; + const canceled = await request(socket, { + type: "agent.requests.cancel.request", + key, + creationKey, + }); + expect(["pending", "settled"]).toContain(canceled["outcome"]); + expect(await creation).not.toMatchObject({ status: "agent_created" }); + release(); + await expect.poll(() => closes).toBe(1); + await expect + .poll(() => request(socket, { type: "agent.requests.inspect.request", key, creationKey })) + .toMatchObject({ outcome: "settled", agentId: null }); + expect(host.daemon.agentManager.listAgents()).toEqual([]); + + socket.close(); + socket = await connect(host.port); + const replay = await request(socket, { + type: "create_agent_request", + config, + idempotencyKey: creationKey, + operation, + }); + expect(replay).not.toMatchObject({ status: "agent_created" }); + expect(creates).toBe(1); + const next = await request(socket, { + type: "create_agent_request", + config, + idempotencyKey: creationKey, + operation: { ...operation, key: randomUUID() }, + }); + expect(next).toMatchObject({ status: "agent_created" }); + const repeated = await request(socket, { + type: "create_agent_request", + config, + idempotencyKey: creationKey, + operation: { ...operation, key: randomUUID() }, + }); + expect(repeated["agentId"]).toBe(next["agentId"]); + expect(creates).toBe(2); + expect(host.daemon.agentManager.listAgents()).toHaveLength(1); + } finally { + release(); + socket.close(); + await host.close(); + } +}, 30_000); + +test("expiry during provider send settles the socket caller and retains cleanup ownership until cancellation finishes", async () => { + let entered!: () => void; + const started = new Promise((resolve) => { + entered = resolve; + }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + let turns = 0; + const provider = createTestAgentClient("codex", { + beforeStartTurn: async () => { + entered(); + await gate; + }, + onStartTurn: () => { + turns++; + }, + }); + const host = await createTestPaseoDaemon({ agentClients: { codex: provider } }); + const socket = await connect(host.port); + const creationKey = randomUUID(); + const key = randomUUID(); + try { + const created = await request(socket, { + type: "create_agent_request", + config: { provider: "codex", cwd: host.paseoHome }, + idempotencyKey: creationKey, + }); + const agentId = z.string().parse(created["agentId"]); + const messageId = randomUUID(); + const operation = { key, deadlineAt: new Date(Date.now() + 60_000).toISOString() }; + const sending = request(socket, { + type: "send_agent_message_request", + agentId, + messageId, + text: "hello", + operation, + }); + await started; + await request(socket, { type: "agent.requests.cancel.request", key, creationKey }); + expect(await sending).not.toMatchObject({ accepted: true }); + // The provider has not acknowledged start. Cancellation cannot claim this work settled yet. + expect( + await request(socket, { type: "agent.requests.inspect.request", key, creationKey }), + ).toMatchObject({ outcome: "pending", agentId }); + release(); + await expect + .poll(() => request(socket, { type: "agent.requests.inspect.request", key, creationKey })) + .toMatchObject({ outcome: "settled", agentId }); + expect(host.daemon.agentManager.hasInFlightRun(agentId)).toBe(false); + await request(socket, { + type: "send_agent_message_request", + agentId, + messageId, + text: "hello", + operation, + }); + expect(turns).toBe(1); + const next = await request(socket, { + type: "send_agent_message_request", + agentId, + messageId: randomUUID(), + text: "next", + operation: { ...operation, key: randomUUID() }, + }); + expect(next).toMatchObject({ accepted: true }); + expect(turns).toBe(2); + } finally { + release(); + socket.close(); + await host.close(); + } +}, 30_000); + +test("a delayed duplicate archive acknowledgement cannot archive a restored continuation again", async () => { + const host = await createTestPaseoDaemon({ + agentClients: { codex: createTestAgentClient("codex") }, + }); + const socket = await connect(host.port); + try { + const created = await request(socket, { + type: "create_agent_request", + config: { provider: "codex", cwd: host.paseoHome }, + idempotencyKey: randomUUID(), + }); + const workspaceId = z.object({ workspaceId: z.string() }).parse(created["agent"]).workspaceId; + const archive = { + type: "archive_workspace_request", + workspaceId, + operation: { key: randomUUID() }, + }; + expect(await request(socket, archive)).toMatchObject({ error: null }); + expect( + await request(socket, { type: "workspace.recovery.restore.request", workspaceId }), + ).toMatchObject({ error: null }); + expect(await request(socket, archive)).toMatchObject({ error: null }); + expect( + await request(socket, { type: "workspace.recovery.inspect.request", workspaceId }), + ).toMatchObject({ state: { kind: "unavailable", reason: "workspace_not_archived" } }); + } finally { + socket.close(); + await host.close(); + } +}, 30_000); diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index c1e149ee5e..2de6556c7b 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -464,7 +464,10 @@ export interface SessionOptions { worktreesRoot?: string; agentManager: AgentManager; agentStorage: AgentStorage; - agentRequests: Pick; + agentRequests: Pick< + AgentRequests, + "create" | "send" | "run" | "control" | "cancelOperation" | "inspectOperation" + >; projectRegistry: ProjectRegistry; workspaceRegistry: WorkspaceRegistry; directorySync?: DirectorySyncService; @@ -754,7 +757,10 @@ export class Session { private readonly daemonSession: DaemonSession; private readonly hubExecutionController: HubExecutionController | null; private readonly workspaceScripts: WorkspaceScriptsService; - private readonly agentRequests: Pick; + private readonly agentRequests: Pick< + AgentRequests, + "create" | "send" | "run" | "control" | "cancelOperation" | "inspectOperation" + >; private readonly createAgentLifecycleDispatch: CreateAgentLifecycleDispatch; constructor(options: SessionOptions) { @@ -2002,6 +2008,7 @@ export class Session { this.dispatchAgentRelationshipMessage(msg) ?? this.dispatchAgentTimelineMessage(msg, source) ?? this.dispatchHubExecutionMessage(msg) ?? + this.dispatchAgentRequestMessage(msg) ?? this.dispatchAgentLifecycleMessage(msg) ?? this.dispatchAgentConfigMessage(msg) ?? this.dispatchCheckoutMessage(msg) ?? @@ -2385,6 +2392,27 @@ export class Session { return undefined; } + private dispatchAgentRequestMessage(msg: SessionInboundMessage): Promise | undefined { + switch (msg.type) { + case "agent.requests.inspect.request": + return this.agentRequests.inspectOperation(msg.key, msg.creationKey).then((result) => { + return this.emit({ + type: "agent.requests.inspect.response", + payload: { requestId: msg.requestId, ...result }, + }); + }); + case "agent.requests.cancel.request": + return this.agentRequests.cancelOperation(msg.key, msg.creationKey).then((result) => { + return this.emit({ + type: "agent.requests.cancel.response", + payload: { requestId: msg.requestId, ...result }, + }); + }); + default: + return undefined; + } + } + private dispatchAgentLifecycleMessage(msg: SessionInboundMessage): Promise | undefined { switch (msg.type) { case "fetch_agents_request": @@ -2408,11 +2436,15 @@ export class Session { case "project.icon.set.request": return this.handleProjectIconSetRequest(msg); case "send_agent_message_request": - return this.handleSendAgentMessageRequest(msg); + return this.agentRequests.run(msg.operation, (signal) => + this.handleSendAgentMessageRequest(msg, signal), + ); case "wait_for_finish_request": return this.handleWaitForFinish(msg.agentId, msg.requestId, msg.timeoutMs); case "create_agent_request": - return this.handleCreateAgentRequest(msg); + return this.agentRequests.run(msg.operation, (signal) => + this.handleCreateAgentRequest(msg, signal), + ); case "resume_agent_request": return this.handleResumeAgentRequest(msg); case "import_agent_request": @@ -2420,7 +2452,7 @@ export class Session { case "refresh_agent_request": return this.handleRefreshAgentRequest(msg); case "cancel_agent_request": - return this.handleCancelAgentRequest(msg.agentId, msg.requestId); + return this.handleCancelAgentRequest(msg.agentId, msg.requestId, msg.operation); case "agent_permission_response": return this.handleAgentPermissionResponse(msg.agentId, msg.requestId, msg.response); case "clear_agent_attention": @@ -2660,7 +2692,9 @@ export class Session { case "workspace.recovery.inspect.request": return this.handleWorkspaceRecoveryInspectRequest(msg); case "workspace.recovery.restore.request": - return this.handleWorkspaceRecoveryRestoreRequest(msg); + return this.agentRequests.run(msg.operation, (signal) => + this.handleWorkspaceRecoveryRestoreRequest(msg, signal), + ); case "workspace.clear_attention.request": return this.handleWorkspaceClearAttentionRequest(msg); case "workspace.mark_unread.request": @@ -3481,9 +3515,10 @@ export class Session { private async handleWorkspaceRecoveryRestoreRequest( request: Extract, + signal?: AbortSignal, ): Promise { try { - await this.restoreWorkspaceAndEmit(request.workspaceId); + await this.restoreWorkspaceAndEmit(request.workspaceId, signal); this.emit({ type: "workspace.recovery.restore.response", payload: { @@ -3568,14 +3603,20 @@ export class Session { /** * Handle create agent request */ - private async handleCreateAgentRequest(msg: CreateAgentRequestMessage): Promise { + private async handleCreateAgentRequest( + msg: CreateAgentRequestMessage, + signal?: AbortSignal, + ): Promise { try { + if (msg.operation && msg.idempotencyKey === undefined) { + throw new Error("Cancelable creation requires an idempotencyKey"); + } let agent: AgentSnapshotPayload; if (msg.idempotencyKey !== undefined) { if (msg.initialPrompt !== undefined) { throw new Error("Idempotent creation requires sending the initial prompt separately"); } - const { requestId: _requestId, idempotencyKey, ...request } = msg; + const { requestId: _requestId, idempotencyKey, operation: _operation, ...request } = msg; const id = await this.agentRequests.create({ key: idempotencyKey, request, @@ -3583,14 +3624,14 @@ export class Session { this.agentManager.getAgent(agentId) != null || (await this.agentStorage.get(agentId)) !== null, create: async (agentId) => { - await this.createSessionAgent(msg, agentId); + await this.createSessionAgent(msg, agentId, signal); }, }); const record = await this.agentStorage.get(id); if (!record) throw new Error("Previously created agent no longer exists"); agent = this.buildStoredAgentPayload(record); } else { - agent = await this.createSessionAgent(msg); + agent = await this.createSessionAgent(msg, undefined, signal); } this.emit({ type: "status", @@ -3625,9 +3666,20 @@ export class Session { } } + private async cleanupCanceledAgentCreate( + agentId: string | null, + signal?: AbortSignal, + ): Promise { + if (!signal?.aborted || !agentId) return agentId; + await this.agentManager.closeAgent(agentId); + await this.agentManager.deleteAgentState(agentId); + return null; + } + private async createSessionAgent( msg: CreateAgentRequestMessage, agentId?: string, + signal?: AbortSignal, ): Promise { const { config, @@ -3676,6 +3728,7 @@ export class Session { hasLegacyGitOptions: Boolean(git), }); createdWorktreeForCleanup = createdWorktree; + signal?.throwIfAborted(); const resolvedIntent = await this.resolveSessionCreateAgentIntent({ request: msg, createdWorktree, @@ -3697,6 +3750,7 @@ export class Session { }, { kind: "session", + signal, agentId, config: resolvedIntent.config, workspaceId: resolvedIntent.intent.workspaceId, @@ -3716,6 +3770,7 @@ export class Session { }, ); createdAgentId = snapshot.id; + signal?.throwIfAborted(); await this.agentUpdates.forwardLiveAgent(snapshot); if (resolvedIntent.createdDirectoryWorkspace && trimmedPrompt) { this.workspaceAutoName.scheduleForDirectory( @@ -3738,6 +3793,7 @@ export class Session { ); return this.buildAgentPayload(liveSnapshot); } catch (error) { + createdAgentId = await this.cleanupCanceledAgentCreate(createdAgentId, signal); await this.createAgentLifecycleDispatch.cleanupCreatedWorktreeAfterFailedAgentCreate({ createdWorktree: createdWorktreeForCleanup, createdAgentId, @@ -4031,13 +4087,19 @@ export class Session { } } - private async handleCancelAgentRequest(agentId: string, requestId?: string): Promise { + private async handleCancelAgentRequest( + agentId: string, + requestId?: string, + operation?: { key: string; deadlineAt?: string }, + ): Promise { this.sessionLogger.info({ agentId }, `Cancel request received for agent ${agentId}`); try { - await cancelAgentRunCommand( - { agentManager: this.agentManager, logger: this.sessionLogger }, - agentId, + await this.agentRequests.control(operation, { action: "interrupt", agentId }, () => + cancelAgentRunCommand( + { agentManager: this.agentManager, logger: this.sessionLogger }, + agentId, + ), ); if (requestId) { const agent = this.agentManager.getAgent(agentId); @@ -5269,8 +5331,8 @@ export class Session { }; } - private async restoreWorkspaceAndEmit(workspaceId: string): Promise { - await this.workspaceRecovery.restore(workspaceId); + private async restoreWorkspaceAndEmit(workspaceId: string, signal?: AbortSignal): Promise { + await this.workspaceRecovery.restore(workspaceId, signal); const workspace = await this.workspaceRegistry.get(workspaceId); if (!workspace) { throw new Error(`Recovered workspace record not found: ${workspaceId}`); @@ -6834,34 +6896,39 @@ export class Session { throw new Error(`Workspace not found: ${request.workspaceId}`); } - await archiveByScope( - { - paseoHome: this.paseoHome, - paseoWorktreesBaseRoot: this.worktreesRoot, - github: this.github, - workspaceGitService: this.workspaceGitService, - agentManager: this.agentManager, - agentStorage: this.agentStorage, - findWorkspaceIdForCwd: (cwd) => this.findWorkspaceIdForCwd(cwd), - getWorkspace: (workspaceId) => this.workspaceRegistry.get(workspaceId), - listActiveWorkspaces: () => this.listActiveWorkspaceRefs(), - archiveWorkspaceRecord: (workspaceId) => this.archiveWorkspaceRecord(workspaceId), - emitWorkspaceUpdatesForWorkspaceIds: (workspaceIds) => - this.emitWorkspaceUpdatesForWorkspaceIds(workspaceIds), - markWorkspaceArchiving: (workspaceIds, archivingAt) => - this.markWorkspaceArchiving(workspaceIds, archivingAt), - clearWorkspaceArchiving: (workspaceIds) => this.clearWorkspaceArchiving(workspaceIds), - assertWorkspaceAutomationAllowed: (workspaceId) => - assertWorkspaceAutomationAllowedForWorkspace(this.workspaceRegistry, workspaceId), - killTerminalsForWorkspace: (workspaceId) => - this.terminalController.killTerminalsForWorkspace(workspaceId), - stopWorkspaceSetup: (workspaceId) => this.workspaceSetupRuntime.stop(workspaceId), - sessionLogger: this.sessionLogger, - }, - { - scope: { kind: "workspace", workspaceId: existing.workspaceId }, - requestId: request.requestId, - }, + await this.agentRequests.control( + request.operation, + { action: "archive", workspaceId: request.workspaceId }, + () => + archiveByScope( + { + paseoHome: this.paseoHome, + paseoWorktreesBaseRoot: this.worktreesRoot, + github: this.github, + workspaceGitService: this.workspaceGitService, + agentManager: this.agentManager, + agentStorage: this.agentStorage, + findWorkspaceIdForCwd: (cwd) => this.findWorkspaceIdForCwd(cwd), + getWorkspace: (workspaceId) => this.workspaceRegistry.get(workspaceId), + listActiveWorkspaces: () => this.listActiveWorkspaceRefs(), + archiveWorkspaceRecord: (workspaceId) => this.archiveWorkspaceRecord(workspaceId), + emitWorkspaceUpdatesForWorkspaceIds: (workspaceIds) => + this.emitWorkspaceUpdatesForWorkspaceIds(workspaceIds), + markWorkspaceArchiving: (workspaceIds, archivingAt) => + this.markWorkspaceArchiving(workspaceIds, archivingAt), + clearWorkspaceArchiving: (workspaceIds) => this.clearWorkspaceArchiving(workspaceIds), + assertWorkspaceAutomationAllowed: (workspaceId) => + assertWorkspaceAutomationAllowedForWorkspace(this.workspaceRegistry, workspaceId), + killTerminalsForWorkspace: (workspaceId) => + this.terminalController.killTerminalsForWorkspace(workspaceId), + stopWorkspaceSetup: (workspaceId) => this.workspaceSetupRuntime.stop(workspaceId), + sessionLogger: this.sessionLogger, + }, + { + scope: { kind: "workspace", workspaceId: existing.workspaceId }, + requestId: request.requestId, + }, + ), ); const archivedWorkspace = await this.workspaceRegistry.get(request.workspaceId); @@ -7569,6 +7636,7 @@ export class Session { private async handleSendAgentMessageRequest( msg: Extract, + signal?: AbortSignal, ): Promise { const resolved = await this.resolveAgentIdentifier(msg.agentId); if (!resolved.ok) { @@ -7584,7 +7652,20 @@ export class Session { return; } + let cancellation: Promise | undefined; + const abort = () => { + if (this.agentManager.getAgent(resolved.agentId)) { + cancellation ??= this.agentManager.cancelAgentRun(resolved.agentId).then((result) => { + if (result.status === "refused") throw new Error("agent_request_outcome_unknown"); + return result; + }); + void cancellation.catch(() => undefined); + } + }; + signal?.addEventListener("abort", abort, { once: true }); + if (signal?.aborted) abort(); try { + signal?.throwIfAborted(); const agentId = resolved.agentId; const prompt = buildAgentPrompt(msg.text, msg.images, msg.attachments); @@ -7598,9 +7679,11 @@ export class Session { "agent.session.send_agent_message", ); const send = async () => { + signal?.throwIfAborted(); const result = await sendPromptToAgent({ agentManager: this.agentManager, agentStorage: this.agentStorage, + signal, agentId, prompt, messageId: msg.messageId, @@ -7608,8 +7691,12 @@ export class Session { clearPendingPermissions: true, logger: this.sessionLogger, }); + if (signal?.aborted) { + await this.agentManager.cancelAgentRun(agentId); + signal.throwIfAborted(); + } if (result.disposition === "turn_started") { - await waitForAgentRunStartWithTimeout(this.agentManager, agentId); + await waitForAgentRunStartWithTimeout(this.agentManager, agentId, signal); } }; if (msg.messageId) { @@ -7650,6 +7737,9 @@ export class Session { error: errorToFriendlyMessage(error), }, }); + } finally { + signal?.removeEventListener("abort", abort); + await cancellation; } } diff --git a/packages/server/src/server/session/workspace-recovery/workspace-recovery-service.test.ts b/packages/server/src/server/session/workspace-recovery/workspace-recovery-service.test.ts index 91631a5631..8061eba69f 100644 --- a/packages/server/src/server/session/workspace-recovery/workspace-recovery-service.test.ts +++ b/packages/server/src/server/session/workspace-recovery/workspace-recovery-service.test.ts @@ -117,6 +117,50 @@ describe("workspace recovery", () => { expect(unarchived).toEqual([workspace.workspaceId]); }); + test("a repeated restore of an already active workspace succeeds without another unarchive", async () => { + const workspace = createWorkspace({ archivedAt: null }); + const { service, unarchived } = createHarness({ workspace }); + await expect(service.restore(workspace.workspaceId)).resolves.toEqual({ + workspaceId: workspace.workspaceId, + action: "unarchive", + }); + expect(unarchived).toEqual([]); + }); + + test("expiry during restore preparation cannot unarchive the workspace", async () => { + const workspace = createWorkspace({ kind: "directory", branch: null }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + let entered!: () => void; + const started = new Promise((resolve) => { + entered = resolve; + }); + let unarchived = false; + const service = createWorkspaceRecoveryService({ + paseoHome: "/paseo-home", + worktreesRoot: "/worktrees", + getWorkspace: async () => workspace, + getProject: async () => createProject(), + isDirectory: async () => { + entered(); + await gate; + return true; + }, + unarchiveWorkspace: async () => { + unarchived = true; + }, + }); + const controller = new AbortController(); + const restoring = service.restore(workspace.workspaceId, controller.signal); + await started; + controller.abort(new Error("expired")); + release(); + await expect(restoring).rejects.toThrow("expired"); + expect(unarchived).toBe(false); + }); + test("does not offer recovery for a missing non-worktree directory", async () => { const workspace = createWorkspace({ kind: "directory", branch: null }); const { service } = createHarness({ workspace }); diff --git a/packages/server/src/server/session/workspace-recovery/workspace-recovery-service.ts b/packages/server/src/server/session/workspace-recovery/workspace-recovery-service.ts index 1d98a9128c..63801f8625 100644 --- a/packages/server/src/server/session/workspace-recovery/workspace-recovery-service.ts +++ b/packages/server/src/server/session/workspace-recovery/workspace-recovery-service.ts @@ -40,7 +40,10 @@ export type WorkspaceRecoveryState = export interface WorkspaceRecoveryService { inspect(workspaceId: string): Promise; - restore(workspaceId: string): Promise<{ workspaceId: string; action: WorkspaceRecoveryAction }>; + restore( + workspaceId: string, + signal?: AbortSignal, + ): Promise<{ workspaceId: string; action: WorkspaceRecoveryAction }>; } type RecoveryPlan = @@ -140,15 +143,20 @@ export function createWorkspaceRecoveryService(deps: { async function restore( workspaceId: string, + signal?: AbortSignal, ): Promise<{ workspaceId: string; action: WorkspaceRecoveryAction }> { + signal?.throwIfAborted(); const resolved = await resolveRecovery(workspaceId); + signal?.throwIfAborted(); if (resolved.kind === "unavailable") { + if (resolved.reason === "workspace_not_archived") return { workspaceId, action: "unarchive" }; throw new Error(resolved.message); } if (resolved.kind === "restore") { await recreateArchivedWorktree(resolved.workspace, resolved.sourceRepoRoot); } + signal?.throwIfAborted(); await deps.unarchiveWorkspace(resolved.workspace); return { workspaceId, action: resolved.kind }; } diff --git a/packages/server/src/server/test-utils/fake-agent-client.ts b/packages/server/src/server/test-utils/fake-agent-client.ts index 4d16f0302e..627dff60c8 100644 --- a/packages/server/src/server/test-utils/fake-agent-client.ts +++ b/packages/server/src/server/test-utils/fake-agent-client.ts @@ -58,11 +58,14 @@ interface FakeAgentSessionOptions { sessionId?: string; memoryMarker?: string | null; closeSession?: () => Promise; + beforeStartTurn?: () => Promise; onStartTurn?: (prompt: AgentPromptInput) => void; } export interface TestAgentClientOptions { + beforeCreateSession?: () => Promise; closeSession?: () => Promise; + beforeStartTurn?: () => Promise; onStartTurn?: (prompt: AgentPromptInput) => void; supportsMcpServers?: boolean; } @@ -336,6 +339,7 @@ class FakeAgentSession implements AgentSession { private activeForegroundTurnId: string | null = null; private readonly closeSession: (() => Promise) | undefined; + private readonly beforeStartTurn: (() => Promise) | undefined; private readonly onStartTurn: ((prompt: AgentPromptInput) => void) | undefined; constructor(options: FakeAgentSessionOptions) { @@ -349,6 +353,7 @@ class FakeAgentSession implements AgentSession { this.memoryMarker = options.memoryMarker ?? null; this.closeSession = options.closeSession; this.onStartTurn = options.onStartTurn; + this.beforeStartTurn = options.beforeStartTurn; this.historyPath = path.join( tmpdir(), "paseo-fake-provider-history", @@ -433,6 +438,7 @@ class FakeAgentSession implements AgentSession { } async startTurn(prompt: AgentPromptInput): Promise<{ turnId: string }> { + await this.beforeStartTurn?.(); if (this.activeForegroundTurnId) { throw new Error("A foreground turn is already active"); } @@ -1202,12 +1208,14 @@ class FakeAgentClient implements AgentClient { config: AgentSessionConfig, _launchContext?: AgentLaunchContext, ): Promise { + await this.options.beforeCreateSession?.(); return new FakeAgentSession({ providerName: this.provider, config: { ...config }, supportsMcpServers: this.options.supportsMcpServers, closeSession: this.options.closeSession, onStartTurn: this.options.onStartTurn, + beforeStartTurn: this.options.beforeStartTurn, }); } @@ -1233,6 +1241,7 @@ class FakeAgentClient implements AgentClient { memoryMarker: typeof marker === "string" ? marker : null, closeSession: this.options.closeSession, onStartTurn: this.options.onStartTurn, + beforeStartTurn: this.options.beforeStartTurn, }); } diff --git a/packages/server/src/server/test-utils/session-stubs.ts b/packages/server/src/server/test-utils/session-stubs.ts index 15ed7ecc09..df9bbd0741 100644 --- a/packages/server/src/server/test-utils/session-stubs.ts +++ b/packages/server/src/server/test-utils/session-stubs.ts @@ -286,6 +286,12 @@ export function createAgentRequestsStub(): SessionOptions["agentRequests"] { return agentId; }, send: (input) => input.send(), + run: (_startup, operation) => operation(), + control: async (_context, _request, operation) => { + await operation(); + }, + cancelOperation: async () => ({ agentId: null, outcome: "settled" }), + inspectOperation: async () => ({ agentId: null, outcome: "settled" }), }; } diff --git a/packages/server/src/server/websocket-server.ts b/packages/server/src/server/websocket-server.ts index 30fce35b13..af8932ae49 100644 --- a/packages/server/src/server/websocket-server.ts +++ b/packages/server/src/server/websocket-server.ts @@ -650,6 +650,7 @@ export class VoiceAssistantWebSocketServer { pluginRuntime?: SessionOptions["pluginRuntime"], orchestrationSkills?: SessionOptions["orchestrationSkills"], workspaceLabelService?: WorkspaceLabelService, + agentRequests?: AgentRequests, ) { this.logger = logger.child({ module: "websocket-server" }); this.workspaceSetupRuntime = workspaceSetupRuntime; @@ -668,7 +669,7 @@ export class VoiceAssistantWebSocketServer { this.orchestrationSkills = orchestrationSkills; this.agentManager = agentManager; this.agentStorage = agentStorage; - this.agentRequests = new AgentRequests(join(paseoHome, "agent-requests")); + this.agentRequests = agentRequests ?? new AgentRequests(join(paseoHome, "agent-requests")); this.projectRegistry = projectRegistry ?? createNoopProjectRegistry(); this.workspaceRegistry = workspaceRegistry ?? createNoopWorkspaceRegistry(); this.workspaceLabelService = workspaceLabelService ?? null; @@ -1633,6 +1634,7 @@ export class VoiceAssistantWebSocketServer { ...(this.serverCapabilities ? { capabilities: this.serverCapabilities } : {}), features: { agentRequestReceipts: true, + agentRequestCancellation: true, hubAgentRpc: true, // COMPAT(directorySync): added in v0.3.x, remove gate after 2027-02-12. directorySync: true, diff --git a/public-docs/hub/configuration/hub-yml.md b/public-docs/hub/configuration/hub-yml.md index 0bc82abf16..3e292ee336 100644 --- a/public-docs/hub/configuration/hub-yml.md +++ b/public-docs/hub/configuration/hub-yml.md @@ -192,7 +192,7 @@ An environment or dynamic named-agent expression must have a finite set of possi | --------------- | -------- | ---------------------------------------------------------------------------------------- | | `id` | yes | Unique within the workflow. | | `environment` | yes | Literal environment name or finite expression resolving to one. | -| `max_runtime` | yes | Step hard limit. | +| `max_runtime` | yes | Absolute step limit, including agent startup. | | `idle_timeout` | yes | Idle limit no longer than `max_runtime`. | | `agent` | yes | Named agent, finite expression selecting a named agent, or complete static inline agent. | | `prompt` | yes | Ordered `text` and `include` blocks. | @@ -203,6 +203,12 @@ An environment or dynamic named-agent expression must have a finite set of possi | `auto_archive` | no | Archive the agent after the step ends. | | `github` | no | Explicit GitHub authority for this step. | +`max_runtime` starts when Hub dispatches the step. Preparing a filesystem worktree, creating the provider agent, restoring a continuation workspace, and acknowledging the initial prompt all consume that budget. The deadline does not reset when startup finishes or activity arrives. The trigger's whole-run deadline can end the step sooner. Provider-specific startup failures can also end a step before its deadline. + +`idle_timeout` starts alongside dispatch. An `initializing` or `running` agent status clears the idle deadline; an `idle` status starts it again. Meaningful agent activity refreshes an armed idle deadline. Connection heartbeats do not count as agent activity. Allow enough idle time for preparation before the daemon reports an active agent. + +Continuation requires a daemon that supports cancelable agent requests. Update the daemon if Hub requests an update. When a deadline expires, Hub ends the step and reconciles daemon cleanup before another step can use the same conversation. If a daemon restart leaves the operation's outcome unknown, inspect the agent and workspace in Paseo before retrying with a new conversation key; Hub does not automatically repeat an ambiguous launch or prompt. + An inline agent is static and complete: ```yaml