Harden agent refresh rescue path

This commit is contained in:
Mohamed Boudra
2026-04-21 17:22:50 +07:00
parent cea97c83b3
commit f3cb19529c
4 changed files with 319 additions and 20 deletions
@@ -462,6 +462,140 @@ describe("AgentManager", () => {
expect(reloaded.currentModeId).toBe("full-access");
});
test("reloadAgentSession completes when the previous session close hangs", async () => {
const workdir = mkdtempSync(join(tmpdir(), "agent-manager-reload-close-timeout-"));
const storagePath = join(workdir, "agents");
const storage = new AgentStorage(storagePath, logger);
class HangingCloseSession extends TestAgentSession {
closeCalled = false;
override async close(): Promise<void> {
this.closeCalled = true;
await new Promise(() => {});
}
}
class HangingCloseClient extends TestAgentClient {
readonly firstSession = new HangingCloseSession({
provider: "codex",
cwd: workdir,
});
resumeSessionCalls = 0;
override async createSession(): Promise<AgentSession> {
return this.firstSession;
}
override async resumeSession(
_handle: AgentPersistenceHandle,
config?: Partial<AgentSessionConfig>,
): Promise<AgentSession> {
this.resumeSessionCalls += 1;
return new TestAgentSession({
provider: "codex",
cwd: config?.cwd ?? workdir,
});
}
}
const client = new HangingCloseClient();
const manager = new AgentManager({
clients: {
codex: client,
},
registry: storage,
logger,
rescueTimeouts: { reloadSessionCloseMs: 10 },
idFactory: () => "00000000-0000-4000-8000-000000000302",
});
try {
const snapshot = await manager.createAgent({
provider: "codex",
cwd: workdir,
});
const reloaded = await manager.reloadAgentSession(snapshot.id);
expect(reloaded.id).toBe(snapshot.id);
expect(client.firstSession.closeCalled).toBe(true);
expect(client.resumeSessionCalls).toBe(1);
} finally {
rmSync(workdir, { recursive: true, force: true });
}
});
test("cancelAgentRun completes when provider interrupt hangs", async () => {
const workdir = mkdtempSync(join(tmpdir(), "agent-manager-interrupt-timeout-"));
const storagePath = join(workdir, "agents");
const storage = new AgentStorage(storagePath, logger);
class HangingInterruptSession extends TestAgentSession {
interruptCalled = false;
override async interrupt(): Promise<void> {
this.interruptCalled = true;
await new Promise(() => {});
}
}
class HangingInterruptClient extends TestAgentClient {
readonly session = new HangingInterruptSession({
provider: "codex",
cwd: workdir,
});
override async createSession(): Promise<AgentSession> {
return this.session;
}
}
const client = new HangingInterruptClient();
const manager = new AgentManager({
clients: {
codex: client,
},
registry: storage,
logger,
rescueTimeouts: { interruptSessionMs: 10 },
idFactory: () => "00000000-0000-4000-8000-000000000303",
});
try {
const snapshot = await manager.createAgent({
provider: "codex",
cwd: workdir,
});
await new Promise<void>((resolve) => {
const unsubscribe = manager.subscribe(
(event) => {
if (
event.type === "agent_state" &&
event.agent.id === snapshot.id &&
event.agent.lifecycle === "running"
) {
unsubscribe();
resolve();
}
},
{ agentId: snapshot.id, replayState: false },
);
client.session.pushEvent({
type: "turn_started",
provider: "codex",
turnId: "hanging-interrupt-turn",
});
});
await expect(manager.cancelAgentRun(snapshot.id)).resolves.toBe(true);
expect(client.session.interruptCalled).toBe(true);
} finally {
rmSync(workdir, { recursive: true, force: true });
}
});
test("listProviderAvailability uses registered client keys, including custom providers", async () => {
const customClient: AgentClient = {
provider: "zai",
+104 -10
View File
@@ -50,6 +50,17 @@ import {
} from "./agent-stream-coalescer.js";
import { getAgentProviderDefinition } from "./provider-manifest.js";
const RELOAD_SESSION_CLOSE_TIMEOUT_MS = 3_000;
const INTERRUPT_SESSION_TIMEOUT_MS = 2_000;
type TimeoutResult = "completed" | "timed_out";
type TimeoutOptions = {
operation: Promise<void>;
timeoutMs: number;
onLateError?: (error: unknown) => void;
};
export { AGENT_LIFECYCLE_STATUSES, type AgentLifecycleStatus };
export type {
AgentTimelineCursor,
@@ -93,6 +104,11 @@ export type ProviderAvailability = {
error: string | null;
};
type AgentManagerRescueTimeouts = {
reloadSessionCloseMs?: number;
interruptSessionMs?: number;
};
export type AgentManagerOptions = {
clients?: Partial<Record<AgentProvider, AgentClient>>;
idFactory?: () => string;
@@ -102,6 +118,7 @@ export type AgentManagerOptions = {
terminalManager?: TerminalManager | null;
mcpBaseUrl?: string;
agentStreamCoalesceWindowMs?: number;
rescueTimeouts?: AgentManagerRescueTimeouts;
logger: Logger;
};
@@ -319,6 +336,7 @@ export class AgentManager {
private mcpBaseUrl: string | null;
private onAgentAttention?: AgentAttentionCallback;
private logger: Logger;
private readonly rescueTimeouts: Required<AgentManagerRescueTimeouts>;
constructor(options: AgentManagerOptions) {
this.idFactory = options?.idFactory ?? (() => randomUUID());
@@ -327,6 +345,12 @@ export class AgentManager {
this.onAgentAttention = options?.onAgentAttention;
this.mcpBaseUrl = options?.mcpBaseUrl ?? null;
this.logger = options.logger.child({ module: "agent", component: "agent-manager" });
this.rescueTimeouts = {
reloadSessionCloseMs:
options.rescueTimeouts?.reloadSessionCloseMs ?? RELOAD_SESSION_CLOSE_TIMEOUT_MS,
interruptSessionMs:
options.rescueTimeouts?.interruptSessionMs ?? INTERRUPT_SESSION_TIMEOUT_MS,
};
this.agentStreamCoalescer = new AgentStreamCoalescer({
windowMs: options.agentStreamCoalesceWindowMs ?? AGENT_STREAM_COALESCE_DEFAULT_WINDOW_MS,
timers: { setTimeout, clearTimeout },
@@ -725,11 +749,7 @@ export class AgentManager {
}
existing.foregroundTurnWaiters.clear();
this.settlePendingForegroundRun(agentId);
try {
await existing.session.close();
} catch (error) {
this.logger.warn({ err: error, agentId }, "Failed to close previous session during refresh");
}
await this.closeReloadedSession(existing.session, agentId);
// Preserve existing labels and timeline during reload.
return this.registerSession(session, normalizedConfig, agentId, {
@@ -744,6 +764,60 @@ export class AgentManager {
});
}
private async closeReloadedSession(session: AgentSession, agentId: string): Promise<void> {
try {
const result = await this.waitWithTimeout({
operation: session.close(),
timeoutMs: this.rescueTimeouts.reloadSessionCloseMs,
onLateError: (error) => {
this.logger.warn(
{ err: error, agentId },
"Previous session close failed after refresh timeout",
);
},
});
if (result === "timed_out") {
this.logger.warn(
{ agentId, timeoutMs: this.rescueTimeouts.reloadSessionCloseMs },
"Timed out closing previous session during refresh",
);
}
} catch (error) {
this.logger.warn({ err: error, agentId }, "Failed to close previous session during refresh");
}
}
private async waitWithTimeout(options: TimeoutOptions): Promise<TimeoutResult> {
let didTimeOut = false;
let timer: NodeJS.Timeout | null = null;
const operation = options.operation
.then((): TimeoutResult => "completed")
.catch((error) => {
if (didTimeOut) {
options.onLateError?.(error);
return "timed_out" as const;
}
throw error;
});
try {
return await Promise.race([
operation,
new Promise<TimeoutResult>((resolve) => {
timer = setTimeout(() => {
didTimeOut = true;
resolve("timed_out");
}, options.timeoutMs);
}),
]);
} finally {
if (timer) {
clearTimeout(timer);
}
}
}
async closeAgent(agentId: string): Promise<void> {
const agent = this.requireAgent(agentId);
this.logger.trace(
@@ -1454,11 +1528,7 @@ export class AgentManager {
return false;
}
try {
await agent.session.interrupt();
} catch (error) {
this.logger.error({ err: error, agentId }, "Failed to interrupt session");
}
await this.interruptSession(agent.session, agentId);
// The interrupt will produce a turn_canceled/turn_failed event via subscribe(),
// which flows through the session event dispatcher and settles the foreground turn waiter.
@@ -1542,6 +1612,30 @@ export class AgentManager {
return true;
}
private async interruptSession(session: AgentSession, agentId: string): Promise<void> {
try {
const result = await this.waitWithTimeout({
operation: session.interrupt(),
timeoutMs: this.rescueTimeouts.interruptSessionMs,
onLateError: (error) => {
this.logger.warn(
{ err: error, agentId },
"Session interrupt failed after timeout during cancel",
);
},
});
if (result === "timed_out") {
this.logger.warn(
{ agentId, timeoutMs: this.rescueTimeouts.interruptSessionMs },
"Timed out interrupting session during cancel",
);
}
} catch (error) {
this.logger.error({ err: error, agentId }, "Failed to interrupt session");
}
}
getPendingPermissions(agentId: string): AgentPermissionRequest[] {
const agent = this.requireSessionAgent(agentId);
return Array.from(agent.pendingPermissions.values());
@@ -1,8 +1,11 @@
import { describe, expect, test, vi } from "vitest";
import type { ChildProcessWithoutNullStreams } from "node:child_process";
import { EventEmitter } from "node:events";
import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs";
import { mkdtemp } from "node:fs/promises";
import { tmpdir } from "node:os";
import path from "node:path";
import { PassThrough } from "node:stream";
import type {
AgentLaunchContext,
@@ -46,6 +49,31 @@ function createSession(configOverrides: Partial<AgentSessionConfig> = {}) {
}
describe("Codex app-server provider", () => {
test("disposes an unresponsive app-server child with SIGKILL", async () => {
vi.useFakeTimers();
const child = new EventEmitter() as ChildProcessWithoutNullStreams;
child.stdin = new PassThrough() as ChildProcessWithoutNullStreams["stdin"];
child.stdout = new PassThrough() as ChildProcessWithoutNullStreams["stdout"];
child.stderr = new PassThrough() as ChildProcessWithoutNullStreams["stderr"];
child.exitCode = null;
child.signalCode = null;
child.kill = vi.fn(() => true) as ChildProcessWithoutNullStreams["kill"];
const client = new __codexAppServerInternals.CodexAppServerClient(child, createTestLogger());
try {
const disposePromise = client.dispose();
expect(child.kill).toHaveBeenCalledWith("SIGTERM");
await vi.advanceTimersByTimeAsync(2_000);
expect(child.kill).toHaveBeenCalledWith("SIGKILL");
await vi.advanceTimersByTimeAsync(1_000);
await expect(disposePromise).resolves.toBeUndefined();
} finally {
vi.useRealTimers();
}
});
test("lists repo skills using WorkspaceGitService repo-root resolution", async () => {
const tempDir = await mkdtemp(path.join(tmpdir(), "codex-skills-"));
const cwd = path.join(tempDir, "repo", "packages", "app");
@@ -63,6 +63,9 @@ import type { WorkspaceGitService } from "../../workspace-git-service.js";
const DEFAULT_TIMEOUT_MS = 14 * 24 * 60 * 60 * 1000;
const TURN_START_TIMEOUT_MS = 90 * 1000;
const INTERRUPT_TIMEOUT_MS = 2_000;
const APP_SERVER_GRACEFUL_SHUTDOWN_TIMEOUT_MS = 2_000;
const APP_SERVER_FORCE_SHUTDOWN_TIMEOUT_MS = 1_000;
const CODEX_PROVIDER = "codex" as const;
const CODEX_IMAGE_ATTACHMENT_DIR = "paseo-attachments";
const CODEX_PLAN_IMPLEMENTATION_PROMPT_PREFIX =
@@ -664,8 +667,40 @@ class CodexAppServerClient {
} catch {
// ignore
}
terminateChildProcessTree(this.child);
await this.exitPromise;
signalChildProcessTree(this.child, "SIGTERM");
if (await this.waitForExit(APP_SERVER_GRACEFUL_SHUTDOWN_TIMEOUT_MS)) {
return;
}
this.logger.warn(
{ timeoutMs: APP_SERVER_GRACEFUL_SHUTDOWN_TIMEOUT_MS },
"Codex app-server did not exit after SIGTERM; sending SIGKILL",
);
signalChildProcessTree(this.child, "SIGKILL");
if (await this.waitForExit(APP_SERVER_FORCE_SHUTDOWN_TIMEOUT_MS)) {
return;
}
this.logger.warn(
{ timeoutMs: APP_SERVER_FORCE_SHUTDOWN_TIMEOUT_MS },
"Codex app-server did not report exit after SIGKILL",
);
}
private async waitForExit(timeoutMs: number): Promise<boolean> {
let timer: NodeJS.Timeout | null = null;
try {
return await Promise.race([
this.exitPromise.then(() => true),
new Promise<boolean>((resolve) => {
timer = setTimeout(() => resolve(false), timeoutMs);
}),
]);
} finally {
if (timer) {
clearTimeout(timer);
}
}
}
private async handleLine(line: string): Promise<void> {
@@ -717,14 +752,17 @@ class CodexAppServerClient {
}
}
function terminateChildProcessTree(child: ChildProcessWithoutNullStreams): void {
if (child.killed) {
function signalChildProcessTree(
child: ChildProcessWithoutNullStreams,
signal: NodeJS.Signals,
): void {
if (child.exitCode !== null || child.signalCode !== null) {
return;
}
if (process.platform !== "win32" && typeof child.pid === "number" && child.pid > 0) {
try {
process.kill(-child.pid, "SIGTERM");
process.kill(-child.pid, signal);
return;
} catch {
// Fall back to the direct child when no separate process group exists.
@@ -732,7 +770,7 @@ function terminateChildProcessTree(child: ChildProcessWithoutNullStreams): void
}
try {
child.kill("SIGTERM");
child.kill(signal);
} catch {
// ignore
}
@@ -3274,10 +3312,14 @@ class CodexAppServerAgentSession implements AgentSession {
async interrupt(): Promise<void> {
if (!this.client || !this.currentThreadId || !this.currentTurnId) return;
try {
await this.client.request("turn/interrupt", {
threadId: this.currentThreadId,
turnId: this.currentTurnId,
});
await this.client.request(
"turn/interrupt",
{
threadId: this.currentThreadId,
turnId: this.currentTurnId,
},
INTERRUPT_TIMEOUT_MS,
);
} catch (error) {
this.logger.warn({ error }, "Failed to interrupt Codex turn");
}
@@ -4278,6 +4320,7 @@ export class CodexAppServerAgentClient implements AgentClient {
export const __codexAppServerInternals = {
buildCodexAppServerEnv,
CodexAppServerClient,
codexModelSupportsFastMode,
CodexAppServerAgentSession,
formatCodexQuestionPrompts,