From 3131270446cb4bcb39f36cbdfc261c9c4d7803e1 Mon Sep 17 00:00:00 2001 From: Bek Date: Wed, 12 Aug 2026 23:36:21 -0400 Subject: [PATCH] fix(codex): revoke child relay after failed yield --- .../codex/src/app-server/native-hook-relay.ts | 10 +- .../src/app-server/run-attempt-finalize.ts | 6 +- ...ttempt.native-hook-relay-retention.test.ts | 114 ++++++++++++++++++ 3 files changed, 128 insertions(+), 2 deletions(-) diff --git a/extensions/codex/src/app-server/native-hook-relay.ts b/extensions/codex/src/app-server/native-hook-relay.ts index f2b1d074f9bf..13fe160fb88c 100644 --- a/extensions/codex/src/app-server/native-hook-relay.ts +++ b/extensions/codex/src/app-server/native-hook-relay.ts @@ -61,6 +61,7 @@ export type CodexNativePreToolUseFailure = { }; export type CodexNativeHookRelay = NativeHookRelayRegistrationHandle & { + authorizeRetentionAfterSuccessfulYield: () => void; claimDirectChild: (threadId: string) => () => void; rejectPendingDirectChild: (threadId: string, reason: string) => void; }; @@ -174,6 +175,7 @@ export function createCodexNativeHookRelay(params: { } >(); let foregroundClosed = false; + let successfulYieldRetentionAuthorized = false; const assertClaim = (threadId: string, claim: symbol) => () => directChildClaims.get(threadId) === claim; const rejectPendingAdmissions = (reason: string) => { @@ -214,7 +216,10 @@ export function createCodexNativeHookRelay(params: { assertActive: params.hostCapabilities.assertActive, retention: { readClaim: readCodexNativeChildThreadId, - shouldRetainAfterForegroundClose: () => directChildClaims.size > 0, + // A child claim identifies the subject; successful parent finalization + // separately authorizes its lifetime beyond foreground closure. + shouldRetainAfterForegroundClose: () => + successfulYieldRetentionAuthorized && directChildClaims.size > 0, allowPreToolUse: (childThreadId) => directChildClaims.has(childThreadId), awaitForegroundAdmission: (childThreadId) => { if (foregroundClosed) { @@ -262,6 +267,9 @@ export function createCodexNativeHookRelay(params: { return { ...relay, unregister, + authorizeRetentionAfterSuccessfulYield: () => { + successfulYieldRetentionAuthorized = true; + }, rejectPendingDirectChild: (threadIdInput, reason) => { const threadId = threadIdInput.trim(); const pending = threadId ? pendingDirectChildAdmissions.get(threadId) : undefined; diff --git a/extensions/codex/src/app-server/run-attempt-finalize.ts b/extensions/codex/src/app-server/run-attempt-finalize.ts index 154ced4c00a4..3bfc31f792eb 100644 --- a/extensions/codex/src/app-server/run-attempt-finalize.ts +++ b/extensions/codex/src/app-server/run-attempt-finalize.ts @@ -479,7 +479,7 @@ export async function finalizeCodexAttempt( }), }, ); - return { + const finalizedResult: EmbeddedRunAttemptResult = { ...result, terminal: attemptTerminal.normalize({ timedOut: effectiveTimedOut, @@ -499,4 +499,8 @@ export async function finalizeCodexAttempt( : {}), systemPromptReport, }; + if (turnSucceeded && toolState.yieldDetected && !runAbortController.signal.aborted) { + resourceState.nativeHookRelay?.authorizeRetentionAfterSuccessfulYield(); + } + return finalizedResult; } diff --git a/extensions/codex/src/app-server/run-attempt.native-hook-relay-retention.test.ts b/extensions/codex/src/app-server/run-attempt.native-hook-relay-retention.test.ts index ebae19fc50c5..ba4a33421214 100644 --- a/extensions/codex/src/app-server/run-attempt.native-hook-relay-retention.test.ts +++ b/extensions/codex/src/app-server/run-attempt.native-hook-relay-retention.test.ts @@ -333,4 +333,118 @@ describe("runCodexAppServerAttempt native hook relay retention", () => { } }, ); + + it("revokes a claimed child when the parent fails after sessions_yield", async () => { + const childThreadId = "child-failed-parent"; + const sessionFile = path.join(tempDir, `${childThreadId}-session.jsonl`); + const workspaceDir = path.join(tempDir, `${childThreadId}-workspace`); + const harness = createStartedThreadHarness(); + const params = createParams(sessionFile, workspaceDir); + params.disableTools = false; + params.runtimePlan = createCodexRuntimePlanFixture(); + setCodexTestModelSupportsTools(params, true); + const fixture = await createAdmittedHostCapabilityTestFixture(params); + params.hostCapabilities = fixture.hostCapabilities; + + const beforeToolCall = vi.fn(async () => undefined); + initializeGlobalHookRunner( + createMockPluginRegistry([{ hookName: "before_tool_call", handler: beforeToolCall }]), + ); + + const run = runCodexAppServerAttempt(params, { + nativeHookRelay: { enabled: true, events: ["pre_tool_use"] }, + }); + try { + await harness.waitForMethod("turn/start"); + const startRequest = harness.requests.find((request) => request.method === "thread/start"); + const relayId = extractRelayIdFromThreadRequest(startRequest?.params); + await harness.notify({ + method: "thread/started", + params: { + thread: { + id: childThreadId, + parentThreadId: "thread-1", + source: { + subAgent: { thread_spawn: { parent_thread_id: "thread-1", depth: 1 } }, + }, + }, + }, + } as CodexServerNotification); + await harness.notify({ + method: "item/completed", + params: { + threadId: "thread-1", + turnId: "turn-1", + item: { + type: "collabAgentToolCall", + tool: "spawnAgent", + status: "completed", + senderThreadId: "thread-1", + receiverThreadIds: [childThreadId], + }, + }, + } as unknown as CodexServerNotification); + const childPayload = { + hook_event_name: "PreToolUse", + agent_id: childThreadId, + cwd: workspaceDir, + tool_name: "Bash", + tool_use_id: `${childThreadId}-tool`, + tool_input: { command: "allow-child" }, + }; + await expect( + invokeNativeHookRelay({ + provider: "codex", + relayId, + event: "pre_tool_use", + rawPayload: childPayload, + }), + ).resolves.toMatchObject({ exitCode: 0 }); + expect(beforeToolCall).toHaveBeenCalledOnce(); + + await expect( + harness.handleServerRequest({ + id: "request-sessions-yield-before-failure", + method: "item/tool/call", + params: { + threadId: "thread-1", + turnId: "turn-1", + callId: "yield-before-failure", + namespace: null, + tool: "sessions_yield", + arguments: { message: "Waiting for child" }, + }, + }), + ).resolves.toMatchObject({ success: true }); + await harness.notify({ + method: "turn/completed", + params: { + threadId: "thread-1", + turnId: "turn-1", + turn: { + id: "turn-1", + status: "failed", + error: { message: "parent failed after yielding" }, + }, + }, + } as CodexServerNotification); + + const result = await run; + expect(readAttemptTerminal(result).promptError).toContain("parent failed after yielding"); + expect( + nativeHookRelayTesting.getNativeHookRelayRegistrationForTests(relayId), + ).toBeUndefined(); + await expect( + invokeNativeHookRelay({ + provider: "codex", + relayId, + event: "pre_tool_use", + rawPayload: childPayload, + }), + ).rejects.toThrow(/not found|inactive/); + } finally { + fixture.closeHost(); + fixture.closeAdmission(); + } + }); });