From 3894d4f40684de4c32739c58845baa9def636972 Mon Sep 17 00:00:00 2001
From: ejclaw <phj1081@gmail.com>
Date: Tue, 2 Jun 2026 20:51:21 +0900
Subject: [PATCH] fix: release Codex lease after final output

---
 src/agent-runner-process.ts | 18 +++++++++--
 src/agent-runner.test.ts    | 62 +++++++++++++++++++++++++++++++++++++
 src/agent-runner.ts         | 19 ++++++++++--
 3 files changed, 94 insertions(+), 5 deletions(-)

diff --git a/src/agent-runner-process.ts b/src/agent-runner-process.ts
index 0a56e8e..e2e5ec3 100644
--- a/src/agent-runner-process.ts
+++ b/src/agent-runner-process.ts
@@ -22,6 +22,11 @@ interface RunSpawnedAgentProcessArgs {
   logsDir: string;
   startTime: number;
   onOutput?: (output: AgentOutput) => Promise<void>;
+  onTerminalStreamedOutputFlushed?: (output: AgentOutput) => void;
+}
+
+function isTerminalStreamedOutput(output: AgentOutput): boolean {
+  return (output.phase ?? 'final') !== 'progress';
 }
 
 function parseLegacyAgentOutput(stdout: string): AgentOutput {
@@ -195,9 +200,16 @@ export function runSpawnedAgentProcess(
               'Streamed agent error output',
             );
           }
-          outputChain = outputChain
-            .then(() => onOutput(parsed))
-            .catch((err) => logStreamedOutputDeliveryError(err, group, input));
+          outputChain = outputChain.then(async () => {
+            try {
+              await onOutput(parsed);
+              if (isTerminalStreamedOutput(parsed)) {
+                args.onTerminalStreamedOutputFlushed?.(parsed);
+              }
+            } catch (err) {
+              logStreamedOutputDeliveryError(err, group, input);
+            }
+          });
         } catch (err) {
           logger.warn(
             {
diff --git a/src/agent-runner.test.ts b/src/agent-runner.test.ts
index 8428462..cf826e8 100644
--- a/src/agent-runner.test.ts
+++ b/src/agent-runner.test.ts
@@ -235,6 +235,68 @@ describe('agent-runner timeout behavior', () => {
     expect(result.newSessionId).toBe('session-456');
   });
 
+  it('releases Codex auth lease after final streamed output flushes even before process close', async () => {
+    const releaseLease = vi.fn();
+    const prepareGroupEnvironmentSpy = vi
+      .spyOn(agentRunnerEnvironment, 'prepareGroupEnvironment')
+      .mockReturnValueOnce({
+        env: {
+          EJCLAW_IPC_DIR: '/tmp/ejclaw-test-data/ipc/test-group',
+        },
+        groupDir: '/tmp/ejclaw-test-groups/test-group',
+        runnerDir: '/tmp/ejclaw-test-runners/codex-runner',
+        codexSessionAuth: {
+          canonicalAuthPath: '/tmp/codex-account/auth.json',
+          sessionAuthPath: '/tmp/codex-session/auth.json',
+          accountIndex: 0,
+          lease: {
+            accountIndex: 0,
+            authPath: '/tmp/codex-account/auth.json',
+            release: releaseLease,
+          },
+        },
+      });
+    const onOutput = vi.fn(async () => {});
+    const codexGroup: RegisteredGroup = {
+      ...testGroup,
+      agentType: 'codex',
+    };
+    let resultPromise: Promise<AgentOutput> | undefined;
+
+    try {
+      resultPromise = runAgentProcess(
+        codexGroup,
+        {
+          ...testInput,
+          runId: 'run-codex-lease-final',
+        },
+        () => {},
+        onOutput,
+      );
+
+      emitOutputMarker(fakeProc, {
+        status: 'success',
+        result: 'TASK_DONE\n작업 완료',
+        phase: 'final',
+        newSessionId: 'codex-session-lease-final',
+      });
+      await vi.advanceTimersByTimeAsync(10);
+
+      expect(onOutput).toHaveBeenCalledWith(
+        expect.objectContaining({
+          result: 'TASK_DONE\n작업 완료',
+          phase: 'final',
+        }),
+      );
+      expect(releaseLease).toHaveBeenCalledTimes(1);
+    } finally {
+      fakeProc.emit('close', 0);
+      await vi.advanceTimersByTimeAsync(10);
+      await resultPromise?.catch(() => undefined);
+      prepareGroupEnvironmentSpy.mockRestore();
+    }
+  });
+
   it('preserves streamed progress phase metadata', async () => {
     const onOutput = vi.fn(async () => {});
     const resultPromise = runAgentProcess(
diff --git a/src/agent-runner.ts b/src/agent-runner.ts
index 7d7a234..3b7a7dd 100644
--- a/src/agent-runner.ts
+++ b/src/agent-runner.ts
@@ -108,6 +108,17 @@ function finalizeCodexAuthSession(
   }
 }
 
+function createCodexAuthSessionFinalizer(
+  getAuth: () => PreparedCodexSessionAuth | null | undefined,
+): () => void {
+  let finalized = false;
+  return () => {
+    if (finalized) return;
+    finalized = true;
+    finalizeCodexAuthSession(getAuth());
+  };
+}
+
 export async function runAgentProcess(
   group: RegisteredGroup,
   input: AgentInput,
@@ -175,6 +186,9 @@ export async function runAgentProcess(
   const safeName = group.folder.replace(/[^a-zA-Z0-9-]/g, '-');
   const processSuffix = input.runId || `${Date.now()}`;
   const processName = `ejclaw-${safeName}-${processSuffix}`;
+  const finalizeCodexAuthSessionOnce = createCodexAuthSessionFinalizer(
+    () => codexSessionAuth,
+  );
 
   // Check if runner is built
   const distEntry = path.join(runnerDir, 'dist', 'index.js');
@@ -231,13 +245,14 @@ export async function runAgentProcess(
       logsDir,
       startTime,
       onOutput,
+      onTerminalStreamedOutputFlushed: finalizeCodexAuthSessionOnce,
     })
       .then((output) => {
-        finalizeCodexAuthSession(codexSessionAuth);
+        finalizeCodexAuthSessionOnce();
         resolve(output);
       })
       .catch((err: unknown) => {
-        finalizeCodexAuthSession(codexSessionAuth);
+        finalizeCodexAuthSessionOnce();
         logger.error(
           { err, processName, chatJid: input.chatJid, runId: input.runId },
           'Spawned agent process runner failed',