Fix agent follow-up turns and wait lifecycle

This commit is contained in:
Mohamed Boudra
2026-04-09 00:24:58 +07:00
parent 8ff94dc03e
commit ffcc35485d
6 changed files with 88 additions and 118 deletions
@@ -17,6 +17,7 @@ import type {
AgentMode,
AgentPermissionRequest,
AgentPermissionResponse,
AgentPermissionResult,
AgentPersistenceHandle,
AgentPromptInput,
AgentProvider,
@@ -1156,11 +1157,10 @@ export class AgentManager {
agent.lastError = undefined;
const self = this;
const pendingRun = self.createPendingForegroundRun();
self.pendingForegroundRuns.set(agentId, pendingRun);
const streamForwarder = (async function* streamForwarder() {
const pendingRun = self.createPendingForegroundRun();
self.pendingForegroundRuns.set(agentId, pendingRun);
let turnId: string;
let waiter: ForegroundTurnWaiter | null = null;
try {
@@ -1440,12 +1440,12 @@ export class AgentManager {
agentId: string,
requestId: string,
response: AgentPermissionResponse,
): Promise<void> {
): Promise<AgentPermissionResult | void> {
const agent = this.requireAgent(agentId);
agent.inFlightPermissionResponses.add(requestId);
try {
await agent.session.respondToPermission(requestId, response);
const result = await agent.session.respondToPermission(requestId, response);
agent.pendingPermissions.delete(requestId);
try {
@@ -1463,6 +1463,8 @@ export class AgentManager {
agent.bufferedPermissionResolutions.delete(requestId);
this.dispatchStream(agent.id, bufferedResolution);
}
return result;
} finally {
agent.inFlightPermissionResponses.delete(requestId);
agent.bufferedPermissionResolutions.delete(requestId);
@@ -1625,8 +1627,9 @@ export class AgentManager {
throw new Error(`Agent ${agentId} not found`);
}
const pendingForegroundRun = this.getPendingForegroundRun(agentId);
const hasForegroundTurn =
Boolean(snapshot.activeForegroundTurnId) || this.hasPendingForegroundRun(agentId);
Boolean(snapshot.activeForegroundTurnId) || Boolean(pendingForegroundRun);
const immediatePermission = this.peekPendingPermission(snapshot);
if (immediatePermission) {
@@ -1668,7 +1671,10 @@ export class AgentManager {
}
let currentStatus: AgentLifecycleStatus = initialStatus;
let hasStarted = initialBusy || hasForegroundTurn;
let hasStarted =
isAgentBusy(initialStatus) ||
Boolean(snapshot.activeForegroundTurnId) ||
Boolean(pendingForegroundRun?.started);
let terminalStatusOverride: AgentLifecycleStatus | null = null;
// Bug #3 Fix: Declare unsubscribe and abortHandler upfront so cleanup can reference them
@@ -441,6 +441,14 @@ export interface AgentLaunchContext {
env?: Record<string, string>;
}
/**
* Returned by respondToPermission when the permission resolution requires
* a follow-up turn (e.g. Codex plan approval → implementation).
*/
export interface AgentPermissionResult {
followUpPrompt?: AgentPromptInput;
}
export interface AgentSession {
readonly provider: AgentProvider;
readonly id: string | null;
@@ -455,7 +463,10 @@ export interface AgentSession {
getCurrentMode(): Promise<string | null>;
setMode(modeId: string): Promise<void>;
getPendingPermissions(): AgentPermissionRequest[];
respondToPermission(requestId: string, response: AgentPermissionResponse): Promise<void>;
respondToPermission(
requestId: string,
response: AgentPermissionResponse,
): Promise<AgentPermissionResult | void>;
describePersistence(): AgentPersistenceHandle | null;
interrupt(): Promise<void>;
close(): Promise<void>;
@@ -466,7 +466,7 @@ export class ACPAgentClient implements AgentClient {
async isAvailable(): Promise<boolean> {
try {
this.resolveLaunchCommand();
await this.resolveLaunchCommand();
return true;
} catch {
return false;
@@ -623,17 +623,13 @@ describe("Codex app-server provider", () => {
});
});
test("approving a synthetic Codex plan permission disables plan and fast mode and starts implementation", async () => {
test("approving a synthetic Codex plan permission disables plan and fast mode and returns follow-up prompt", async () => {
const session = createSession({
featureValues: { plan_mode: true, fast_mode: true },
});
const events: AgentStreamEvent[] = [];
session.subscribe((event) => events.push(event));
const startTurnSpy = vi
.spyOn(session, "startTurn")
.mockResolvedValue({ turnId: "follow-up-turn" });
(session as any).handleNotification("turn/started", {
turn: { id: "turn-plan-2" },
});
@@ -653,7 +649,7 @@ describe("Codex app-server provider", () => {
throw new Error("Expected synthetic plan approval permission");
}
await session.respondToPermission(request.request.id, {
const result = await session.respondToPermission(request.request.id, {
behavior: "allow",
selectedActionId: "implement",
});
@@ -664,7 +660,10 @@ describe("Codex app-server provider", () => {
plan_mode: false,
fast_mode: false,
});
expect(startTurnSpy).toHaveBeenCalledWith(
// The session returns the follow-up prompt instead of calling startTurn directly.
// The caller (session/agent-manager) is responsible for sending it through streamAgent.
expect(result).toBeDefined();
expect(result!.followUpPrompt).toEqual(
expect.stringContaining("The user approved the plan. Implement it now."),
);
expect(events.at(-1)).toEqual({
@@ -677,56 +676,4 @@ describe("Codex app-server provider", () => {
},
});
});
test("failed synthetic Codex plan implementation keeps the permission pending for retry", async () => {
const session = createSession({
featureValues: { plan_mode: true, fast_mode: true },
});
const events: AgentStreamEvent[] = [];
session.subscribe((event) => events.push(event));
const startTurnSpy = vi
.spyOn(session, "startTurn")
.mockRejectedValueOnce(new Error("follow-up failed"));
(session as any).handleNotification("turn/started", {
turn: { id: "turn-plan-retry" },
});
(session as any).handleNotification("turn/plan/updated", {
plan: [{ step: "Implement the retriable flow", status: "pending" }],
});
(session as any).handleNotification("turn/completed", {
turn: { status: "completed", error: null },
});
const request = events.find(
(event): event is Extract<AgentStreamEvent, { type: "permission_requested" }> =>
event.type === "permission_requested" && event.request.kind === "plan",
);
expect(request).toBeDefined();
if (!request) {
throw new Error("Expected synthetic plan approval permission");
}
await expect(
session.respondToPermission(request.request.id, {
behavior: "allow",
selectedActionId: "implement",
}),
).rejects.toThrow("follow-up failed");
expect(startTurnSpy).toHaveBeenCalledTimes(1);
expect((session as any).planModeEnabled).toBe(true);
expect((session as any).config.featureValues).toEqual({
plan_mode: true,
fast_mode: true,
});
expect(session.getPendingPermissions()).toEqual([request.request]);
expect(
events.some(
(event) =>
event.type === "permission_resolved" && event.requestId === request.request.id,
),
).toBe(false);
});
});
@@ -9,6 +9,7 @@ import type {
McpServerConfig,
AgentPermissionRequest,
AgentPermissionResponse,
AgentPermissionResult,
AgentPromptContentBlock,
AgentPromptInput,
AgentRunOptions,
@@ -2673,22 +2674,19 @@ class CodexAppServerAgentSession implements AgentSession {
this.emitEvent({ type: "permission_requested", provider: CODEX_PROVIDER, request });
}
private async handleApprovedPlanPermission(params: { planText?: unknown }): Promise<void> {
/**
* Prepare the session for plan implementation by disabling plan/fast mode
* and returning the implementation prompt. The caller is responsible for
* starting the turn through the normal streamAgent path.
*/
private preparePlanImplementation(params: { planText?: unknown }): string {
const planText =
typeof params.planText === "string" ? normalizePlanMarkdown(params.planText) : "";
const previousPlanMode = this.planModeEnabled;
const previousFastMode = this.serviceTier === "fast";
this.applyFeatureValue("plan_mode", false);
this.applyFeatureValue("fast_mode", false);
try {
await this.startTurn(buildCodexPlanImplementationPrompt(planText));
} catch (error) {
this.applyFeatureValue("plan_mode", previousPlanMode);
this.applyFeatureValue("fast_mode", previousFastMode);
throw error;
}
return buildCodexPlanImplementationPrompt(planText);
}
private registerRequestHandlers(): void {
@@ -3102,7 +3100,10 @@ class CodexAppServerAgentSession implements AgentSession {
return Array.from(this.pendingPermissions.values());
}
async respondToPermission(requestId: string, response: AgentPermissionResponse): Promise<void> {
async respondToPermission(
requestId: string,
response: AgentPermissionResponse,
): Promise<AgentPermissionResult | void> {
const pending = this.pendingPermissionHandlers.get(requestId);
if (!pending) {
throw new Error(`No pending Codex app-server permission request with id '${requestId}'`);
@@ -3110,8 +3111,9 @@ class CodexAppServerAgentSession implements AgentSession {
const pendingRequest = this.pendingPermissions.get(requestId) ?? null;
if (pending.kind === "plan") {
let followUpPrompt: string | undefined;
if (response.behavior === "allow") {
await this.handleApprovedPlanPermission({
followUpPrompt = this.preparePlanImplementation({
planText: pending.planText ?? pendingRequest?.metadata?.planText,
});
}
@@ -3125,6 +3127,9 @@ class CodexAppServerAgentSession implements AgentSession {
requestId,
resolution: response,
});
if (followUpPrompt) {
return { followUpPrompt };
}
return;
}
+39 -38
View File
@@ -2836,7 +2836,7 @@ export class Session {
images?: Array<{ data: string; mimeType: string }>,
runOptions?: AgentRunOptions,
options?: { spokenInput?: boolean },
): Promise<void> {
): Promise<{ ok: true } | { ok: false; error: string }> {
this.sessionLogger.info(
{ agentId, textPreview: text.substring(0, 50), imageCount: images?.length ?? 0 },
`Sending text to agent ${agentId}${images && images.length > 0 ? ` with ${images.length} image attachment(s)` : ""}`,
@@ -2848,7 +2848,10 @@ export class Session {
await this.ensureAgentLoaded(agentId);
} catch (error) {
this.handleAgentRunError(agentId, error, "Failed to initialize agent before sending prompt");
return;
return {
ok: false,
error: error instanceof Error ? error.message : String(error),
};
}
try {
@@ -2866,7 +2869,7 @@ export class Session {
const promptText = options?.spokenInput ? wrapSpokenInput(text) : text;
const prompt = this.buildAgentPrompt(promptText, images);
this.startAgentStream(agentId, prompt, runOptions);
return this.startAgentStream(agentId, prompt, runOptions);
}
/**
@@ -2914,6 +2917,29 @@ export class Session {
const snapshot = await this.agentManager.createAgent(sessionConfig, undefined, { labels });
await this.forwardAgentUpdate(snapshot);
if (trimmedPrompt) {
scheduleAgentMetadataGeneration({
agentManager: this.agentManager,
agentId: snapshot.id,
cwd: snapshot.cwd,
initialPrompt: trimmedPrompt,
explicitTitle,
paseoHome: this.paseoHome,
logger: this.sessionLogger,
});
const started = await this.handleSendAgentMessage(
snapshot.id,
trimmedPrompt,
resolveClientMessageId(clientMessageId),
images,
outputSchema ? { outputSchema } : undefined,
);
if (!started.ok) {
throw new Error(started.error);
}
}
if (requestId) {
const agentPayload = await this.getAgentPayloadById(snapshot.id);
if (!agentPayload) {
@@ -2930,40 +2956,6 @@ export class Session {
});
}
if (trimmedPrompt) {
scheduleAgentMetadataGeneration({
agentManager: this.agentManager,
agentId: snapshot.id,
cwd: snapshot.cwd,
initialPrompt: trimmedPrompt,
explicitTitle,
paseoHome: this.paseoHome,
logger: this.sessionLogger,
});
void this.handleSendAgentMessage(
snapshot.id,
trimmedPrompt,
resolveClientMessageId(clientMessageId),
images,
outputSchema ? { outputSchema } : undefined,
).catch((promptError) => {
this.sessionLogger.error(
{ err: promptError, agentId: snapshot.id },
`Failed to run initial prompt for agent ${snapshot.id}`,
);
this.emit({
type: "activity_log",
payload: {
id: uuidv4(),
timestamp: new Date(),
type: "error",
content: `Initial prompt failed: ${(promptError as Error)?.message ?? promptError}`,
},
});
});
}
if (worktreeConfig) {
void runAsyncWorktreeBootstrap({
agentId: snapshot.id,
@@ -3924,8 +3916,16 @@ export class Session {
);
try {
await this.agentManager.respondToPermission(agentId, requestId, response);
const result = await this.agentManager.respondToPermission(agentId, requestId, response);
this.sessionLogger.debug({ agentId }, `Permission response forwarded to agent ${agentId}`);
if (result?.followUpPrompt) {
this.sessionLogger.debug(
{ agentId },
"Permission response requires follow-up turn, starting agent stream",
);
this.startAgentStream(agentId, result.followUpPrompt);
}
} catch (error: any) {
this.sessionLogger.error(
{ err: error, agentId, requestId },
@@ -6602,6 +6602,7 @@ export class Session {
try {
let result = await this.agentManager.waitForAgentEvent(agentId, {
signal: abortController.signal,
waitForActive: true,
});
let final = await this.getAgentPayloadById(agentId);
if (!final) {