From f23495d5474d5107fd2c76dd4981c4c7e53733c4 Mon Sep 17 00:00:00 2001
From: AgentRun Codex <agentrun-codex@users.noreply.github.com>
Date: Tue, 21 Jul 2026 01:42:02 +0200
Subject: [PATCH] fix: stop Kafka replay consumer cleanly

---
 src/debug/stdio-event-reconstruction.ts       | 20 ++++++--
 .../cases/36-stdio-event-reconstruction.ts    | 51 +++++++++++++++++++
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/src/debug/stdio-event-reconstruction.ts b/src/debug/stdio-event-reconstruction.ts
index 332c010..328308a 100644
--- a/src/debug/stdio-event-reconstruction.ts
+++ b/src/debug/stdio-event-reconstruction.ts
@@ -278,15 +278,19 @@ export function reconstructAgentRunEventsFromStdioFrames(
   };
 }
 
-export async function readStdioKafkaFrames(request: StdioFrameReadRequest): Promise<StdioFrameReadResult> {
+export async function readStdioKafkaFrames(
+  request: StdioFrameReadRequest,
+  dependencies: { consumer?: Consumer } = {},
+): Promise<StdioFrameReadResult> {
   const startedAt = Date.now();
   const kafka = new Kafka({ clientId: `${request.config.clientId}-stdio-regenerate`, brokers: request.config.brokers, logLevel: logLevel.NOTHING });
-  const consumer = kafka.consumer({ groupId: request.groupId, allowAutoTopicCreation: false });
+  const consumer = dependencies.consumer ?? kafka.consumer({ groupId: request.groupId, allowAutoTopicCreation: false });
   const records: StdioKafkaFrameRecord[] = [];
   let scanned = 0;
   let invalidJson = 0;
   let timedOut = false;
   await consumer.connect();
+  let runPromise: Promise<void> | null = null;
   try {
     await consumer.subscribe({ topic: request.topic, fromBeginning: true });
     await new Promise<void>((resolve, reject) => {
@@ -301,7 +305,7 @@ export async function readStdioKafkaFrames(request: StdioFrameReadRequest): Prom
         finish();
       }, request.timeoutMs);
       timer.unref?.();
-      consumer.run({
+      runPromise = consumer.run({
         eachMessage: async ({ topic, partition, message }) => {
           scanned += 1;
           const valueText = message.value?.toString("utf8") ?? "";
@@ -330,11 +334,21 @@ export async function readStdioKafkaFrames(request: StdioFrameReadRequest): Prom
       });
     });
   } finally {
+    await stopQuietly(consumer);
+    await runPromise;
     await disconnectQuietly(consumer);
   }
   return { sourceKind: "kafka", records, scanned, invalidJson, timedOut, elapsedMs: Date.now() - startedAt };
 }
 
+async function stopQuietly(consumer: Consumer): Promise<void> {
+  try {
+    await consumer.stop();
+  } catch {
+    // Disconnect remains the final cleanup boundary when Kafka is already unavailable.
+  }
+}
+
 export async function publishReconstructedAgentRunRecords(records: ReconstructedAgentRunKafkaRecord[], config: AgentRunKafkaConfig): Promise<void> {
   await publishAgentRunKafkaMessagesOnce(config, records.map((record) => ({
     topic: record.topic,
diff --git a/src/selftest/cases/36-stdio-event-reconstruction.ts b/src/selftest/cases/36-stdio-event-reconstruction.ts
index 4a68e8c..359c727 100644
--- a/src/selftest/cases/36-stdio-event-reconstruction.ts
+++ b/src/selftest/cases/36-stdio-event-reconstruction.ts
@@ -20,6 +20,7 @@ import {
   parseStdioKafkaFrameJsonl,
   reconstructAgentRunEventsFromStdioFrames,
   regenerateAgentRunEventsFromStdio,
+  readStdioKafkaFrames,
   type ReconstructedAgentRunKafkaRecord,
 } from "../../debug/stdio-event-reconstruction.js";
 import type { SelfTestCase } from "../harness.js";
@@ -110,6 +111,7 @@ const selfTest: SelfTestCase = async (context) => {
   assert.match(String(result.nextHint), /kafka tail agentrun --topic agentrun\.event\.debug\.v1/u);
 
   await assertOneShotPublisherOwnership(frames);
+  await assertKafkaReaderStopsBeforeDisconnect(frames[0]!);
 
   let productPublisherCalls = 0;
   await assert.rejects(
@@ -197,6 +199,7 @@ const selfTest: SelfTestCase = async (context) => {
       "jsonl-offline-source",
       "injected-reader-publisher",
       "one-shot-producer-disconnect",
+      "kafka-reader-stop-before-disconnect",
       "manager-producer-isolation",
       "partial-reconstruction-provenance",
       "debug-topic-group-defaults",
@@ -209,6 +212,54 @@ const selfTest: SelfTestCase = async (context) => {
   };
 };
 
+async function assertKafkaReaderStopsBeforeDisconnect(frame: ReturnType<typeof parseStdioKafkaFrameJsonl>[number]): Promise<void> {
+  const lifecycle: string[] = [];
+  let stopped = false;
+  let releaseRun: (() => void) | null = null;
+  const runFinished = new Promise<void>((resolve) => { releaseRun = resolve; });
+  const consumer = {
+    connect: async () => { lifecycle.push("connect"); },
+    subscribe: async () => { lifecycle.push("subscribe"); },
+    run: async ({ eachMessage }: { eachMessage: (input: JsonRecord) => Promise<void> }) => {
+      lifecycle.push("run");
+      await eachMessage({
+        topic: frame.topic,
+        partition: frame.partition ?? 0,
+        message: {
+          offset: frame.offset ?? "0",
+          key: frame.key === null ? null : Buffer.from(frame.key),
+          timestamp: frame.timestamp,
+          value: Buffer.from(JSON.stringify(frame.value)),
+        },
+      });
+      await runFinished;
+      assert.equal(stopped, true, "consumer.run must only settle after consumer.stop");
+    },
+    stop: async () => { lifecycle.push("stop"); stopped = true; releaseRun?.(); },
+    disconnect: async () => { lifecycle.push("disconnect"); },
+  };
+  const result = await readStdioKafkaFrames({
+    config: {
+      brokers: ["127.0.0.1:9092"],
+      clientId: "agentrun-selftest",
+      agentrunEventTopic: "agentrun.event.v1",
+      codexStdioTopic: "codex-stdio.raw.v1",
+      dnsServers: [],
+      dnsSearchDomains: [],
+      enabled: true,
+      valuesPrinted: false,
+    },
+    topic: frame.topic,
+    groupId: "agentrun-selftest-reader",
+    sessionId,
+    traceId,
+    limit: 1,
+    timeoutMs: 250,
+  }, { consumer: consumer as never });
+  assert.equal(result.records.length, 1);
+  assert.deepEqual(lifecycle, ["connect", "subscribe", "run", "stop", "disconnect"]);
+}
+
 async function assertOneShotPublisherOwnership(frames: ReturnType<typeof parseStdioKafkaFrameJsonl>): Promise<void> {
   const producers: Array<{ connectCalls: number; sendCalls: number; disconnectCalls: number }> = [];
   await setAgentRunKafkaProducerFactoryForSelfTest(() => {