fix: 通过 Kafka SSE 单路径恢复 Workbench 刷新状态

internal/cloud/kafka-event-bridge.test.ts

@@ -6,7 +6,22 @@ import { test } from "bun:test";
 
 import { createCloudApiBunServer } from "./bun-server.ts";
 import { buildCloudApiReadiness } from "./health-contract.ts";
-import { decodeCanonicalAgentRunKafkaMessage, kafkaDnsLookup, kafkaEventBridgeConfig, projectAgentRunKafkaEventToHwlabEvent, publishAgentRunKafkaMessageLive, relayHwlabKafkaOutboxOnce, startHwlabKafkaEventBridge } from "./kafka-event-bridge.ts";
+import { decodeCanonicalAgentRunKafkaMessage, kafkaDnsLookup, kafkaEventBridgeConfig, kafkaPartitionForKey, projectAgentRunKafkaEventToHwlabEvent, publishAgentRunKafkaMessageLive, relayHwlabKafkaOutboxOnce, startHwlabKafkaEventBridge } from "./kafka-event-bridge.ts";
+
+test("Workbench refresh replay resolves the same Kafka partition as the session-keyed producer", () => {
+  const offsets = [
+    { partition: 0, startOffset: "0", endOffset: "100" },
+    { partition: 1, startOffset: "0", endOffset: "200" },
+    { partition: 2, startOffset: "0", endOffset: "300" }
+  ];
+  const sessionId = "ses_refresh_partition_scope";
+  const first = kafkaPartitionForKey(sessionId, "hwlab.event.v1", offsets);
+  const repeated = kafkaPartitionForKey(sessionId, "hwlab.event.v1", [...offsets].reverse());
+
+  assert.ok([0, 1, 2].includes(first));
+  assert.equal(repeated, first);
+  assert.equal(kafkaPartitionForKey(null, "hwlab.event.v1", offsets), null);
+});
 
 const PROJECTOR_ENV = Object.freeze({
   HWLAB_KAFKA_DIRECT_PUBLISH_ENABLED: "false",

internal/cloud/kafka-event-bridge.ts

@@ -7,7 +7,7 @@ import { isIP } from "node:net";
 import net from "node:net";
 import tls from "node:tls";
 
-import { Kafka, logLevel } from "kafkajs";
+import { Kafka, logLevel, Partitioners } from "kafkajs";
 import { emitCodeAgentOtelSpan } from "./otel-trace.ts";
 import { buildWorkbenchProjectionEventFacts } from "./workbench-projection-writer.ts";
 import {
@@ -958,7 +958,16 @@ function requireKafkaProjectorStore(runtimeStore, capabilities = {}) {
   if (missing.length > 0) throw contractError("hwlab_kafka_projector_store_invalid", `Kafka durable capabilities require a runtime store: ${missing.join(", ")}`);
 }
 
-export async function queryKafkaEventStream({ env = process.env, stream = "hwlab", topic = null, traceId = null, sessionId = null, runId = null, commandId = null, limit = DEFAULT_QUERY_LIMIT, scanLimit = null, timeoutMs = DEFAULT_QUERY_TIMEOUT_MS, fromBeginning = true, groupIdPrefix = null, signal = null, kafkaFactory = defaultKafkaFactory } = {}) {
+export async function queryKafkaEventStream({ env = process.env, stream = "hwlab", topic = null, traceId = null, sessionId = null, runId = null, commandId = null, partitionKey = null, limit = DEFAULT_QUERY_LIMIT, scanLimit = null, timeoutMs = DEFAULT_QUERY_TIMEOUT_MS, fromBeginning = true, groupIdPrefix = null, signal = null, kafkaFactory = defaultKafkaFactory } = {}) {
+  const queryStartedAtMs = Date.now();
+  const timing = {
+    endOffsetSnapshotMs: 0,
+    groupOffsetsMs: 0,
+    consumerConnectMs: 0,
+    consumerSubscribeMs: 0,
+    scanMs: 0,
+    cleanupMs: 0
+  };
   const brokers = csv(env.HWLAB_KAFKA_BOOTSTRAP_SERVERS);
   if (brokers.length === 0) throw new Error("HWLAB_KAFKA_BOOTSTRAP_SERVERS is required for Kafka event queries.");
   const clientId = stringValue(env.HWLAB_KAFKA_CLIENT_ID) || DEFAULT_CLIENT_ID;
@@ -975,14 +984,24 @@ export async function queryKafkaEventStream({ env = process.env, stream = "hwlab
     dnsServers: csv(env.HWLAB_KAFKA_DNS_SERVERS),
     dnsSearchDomains: csv(env.HWLAB_KAFKA_DNS_SEARCH_DOMAINS)
   });
+  let stageStartedAtMs = Date.now();
   const endOffsetSnapshot = await kafkaTopicEndOffsetSnapshot(kafka, resolvedTopic, deadlineMs, signal);
+  timing.endOffsetSnapshotMs = Date.now() - stageStartedAtMs;
   const endOffsetByPartition = new Map(endOffsetSnapshot.offsets.map((entry) => [entry.partition, entry.endOffset]));
   const startOffsetByPartition = new Map(endOffsetSnapshot.offsets.map((entry) => [entry.partition, entry.startOffset]));
-  const reachedEndPartitions = new Set(endOffsetSnapshot.offsets.filter(kafkaOffsetPartitionIsEmpty).map((entry) => entry.partition));
+  const targetPartition = kafkaPartitionForKey(partitionKey, resolvedTopic, endOffsetSnapshot.offsets);
+  const skippedPartitions = targetPartition === null ? [] : endOffsetSnapshot.offsets.filter((entry) => entry.partition !== targetPartition).map((entry) => entry.partition);
+  const requiresScopedGroupOffsets = skippedPartitions.length > 0;
+  const reachedEndPartitions = new Set([...skippedPartitions, ...endOffsetSnapshot.offsets.filter(kafkaOffsetPartitionIsEmpty).map((entry) => entry.partition)]);
   const verifiedStartPartitions = new Set(reachedEndPartitions);
   const resolvedGroupIdPrefix = stringValue(groupIdPrefix) || `${clientId}-query`;
   const groupId = `${resolvedGroupIdPrefix}-${Date.now()}-${randomUUID().slice(0, 8)}`;
-  const consumer = kafka.consumer({ groupId, allowAutoTopicCreation: false });
+  const consumer = kafka.consumer({
+    groupId,
+    allowAutoTopicCreation: false,
+    minBytes: 1,
+    maxWaitTimeInMs: 100
+  });
   const events = [];
   const firstScannedOffsetByPartition = new Map();
   const lastScannedOffsetByPartition = new Map();
@@ -1005,11 +1024,28 @@ export async function queryKafkaEventStream({ env = process.env, stream = "hwlab
       completionReason = "retention-start-unavailable";
       return kafkaQueryResult();
     }
+    if (requiresScopedGroupOffsets) {
+      stageStartedAtMs = Date.now();
+      await prepareKafkaQueryGroupOffsets(kafka, {
+        groupId,
+        topic: resolvedTopic,
+        targetPartition,
+        offsets: endOffsetSnapshot.offsets,
+        deadlineMs,
+        signal
+      });
+      timing.groupOffsetsMs = Date.now() - stageStartedAtMs;
+    }
+    stageStartedAtMs = Date.now();
     await withinKafkaQueryBudget(consumer.connect(), deadlineMs, "consumer-connect", signal);
-    await withinKafkaQueryBudget(consumer.subscribe({ topic: resolvedTopic, fromBeginning: fromBeginning !== false }), deadlineMs, "consumer-subscribe", signal);
+    timing.consumerConnectMs = Date.now() - stageStartedAtMs;
+    stageStartedAtMs = Date.now();
+    await withinKafkaQueryBudget(consumer.subscribe({ topic: resolvedTopic, fromBeginning: !requiresScopedGroupOffsets && fromBeginning !== false }), deadlineMs, "consumer-subscribe", signal);
+    timing.consumerSubscribeMs = Date.now() - stageStartedAtMs;
     const alreadyAtBoundedEnd = fromBeginning !== false && endOffsetSnapshot.available && reachedEndPartitions.size === endOffsetByPartition.size;
     if (alreadyAtBoundedEnd) completionReason = "end-offset";
     if (alreadyAtBoundedEnd) return kafkaQueryResult();
+    stageStartedAtMs = Date.now();
     await new Promise((resolve, reject) => {
       let finished = false;
       const finish = (reason = "timeout") => {
@@ -1027,64 +1063,79 @@ export async function queryKafkaEventStream({ env = process.env, stream = "hwlab
       }
       running = true;
       consumer.run({
-        eachMessage: async ({ topic: messageTopic, partition, message }) => {
-          if (finished) return;
-          const endOffset = endOffsetByPartition.get(partition);
-          if (fromBeginning !== false && endOffsetSnapshot.available && !firstScannedOffsetByPartition.has(partition)) {
-            const actualStartOffset = stringValue(message.offset);
-            const expectedStartOffset = startOffsetByPartition.get(partition);
-            firstScannedOffsetByPartition.set(partition, actualStartOffset);
-            if (!endOffsetByPartition.has(partition) || !expectedStartOffset || actualStartOffset !== expectedStartOffset) {
-              finish("retention-start-moved");
-              return;
+        eachBatchAutoResolve: true,
+        eachBatch: async ({ batch }) => {
+          for (const message of batch.messages) {
+            if (finished) return;
+            const messageTopic = batch.topic;
+            const partition = batch.partition;
+            const endOffset = endOffsetByPartition.get(partition);
+            if (fromBeginning !== false && endOffsetSnapshot.available && !firstScannedOffsetByPartition.has(partition)) {
+              const actualStartOffset = stringValue(message.offset);
+              const expectedStartOffset = startOffsetByPartition.get(partition);
+              firstScannedOffsetByPartition.set(partition, actualStartOffset);
+              if (!endOffsetByPartition.has(partition) || !expectedStartOffset || actualStartOffset !== expectedStartOffset) {
+                finish("retention-start-moved");
+                return;
+              }
+              verifiedStartPartitions.add(partition);
             }
-            verifiedStartPartitions.add(partition);
-          }
-          if (fromBeginning !== false && endOffset && kafkaOffsetAtOrBeyond(message.offset, endOffset)) {
-            excludedPostBarrierCount += 1;
-            reachedEndPartitions.add(partition);
+            if (fromBeginning !== false && endOffset && kafkaOffsetAtOrBeyond(message.offset, endOffset)) {
+              excludedPostBarrierCount += 1;
+              reachedEndPartitions.add(partition);
+              if (endOffsetSnapshot.available && reachedEndPartitions.size === endOffsetByPartition.size) {
+                finish(verifiedStartPartitions.size === endOffsetByPartition.size ? "end-offset" : "retention-start-unverified");
+              }
+              continue;
+            }
+            scannedCount += 1;
+            lastScannedOffsetByPartition.set(partition, stringValue(message.offset));
+            const valueText = message.value ? Buffer.from(message.value).toString("utf8") : "";
+            const value = parseJson(valueText);
+            if (!value || typeof value !== "object" || Array.isArray(value)) invalidJsonCount += 1;
+            else {
+              parsedCount += 1;
+              if (!eventMatchesFilters(value, { traceId, sessionId, runId, commandId })) filterRejectedCount += 1;
+              else {
+                events.push({
+                  topic: messageTopic,
+                  partition,
+                  offset: message.offset,
+                  key: message.key ? Buffer.from(message.key).toString("utf8") : null,
+                  timestamp: message.timestamp ?? null,
+                  valueSha256: sha256(valueText),
+                  value
+                });
+              }
+            }
+            if (fromBeginning !== false && endOffset && kafkaOffsetReached(message.offset, endOffset)) reachedEndPartitions.add(partition);
             if (endOffsetSnapshot.available && reachedEndPartitions.size === endOffsetByPartition.size) {
               finish(verifiedStartPartitions.size === endOffsetByPartition.size ? "end-offset" : "retention-start-unverified");
+              return;
             }
-            return;
-          }
-          scannedCount += 1;
-          lastScannedOffsetByPartition.set(partition, stringValue(message.offset));
-          const valueText = message.value ? Buffer.from(message.value).toString("utf8") : "";
-          const value = parseJson(valueText);
-          if (!value || typeof value !== "object" || Array.isArray(value)) invalidJsonCount += 1;
-          else {
-            parsedCount += 1;
-            if (!eventMatchesFilters(value, { traceId, sessionId, runId, commandId })) filterRejectedCount += 1;
-            else {
-              events.push({
-                topic: messageTopic,
-                partition,
-                offset: message.offset,
-                key: message.key ? Buffer.from(message.key).toString("utf8") : null,
-                timestamp: message.timestamp ?? null,
-                valueSha256: sha256(valueText),
-                value
-              });
+            if (events.length >= maxEvents) {
+              finish("limit");
+              return;
+            }
+            if (maxScannedRecords !== null && scannedCount >= maxScannedRecords) {
+              finish("scan-limit");
+              return;
             }
           }
-          if (fromBeginning !== false && endOffset && kafkaOffsetReached(message.offset, endOffset)) reachedEndPartitions.add(partition);
-          if (endOffsetSnapshot.available && reachedEndPartitions.size === endOffsetByPartition.size) {
-            return finish(verifiedStartPartitions.size === endOffsetByPartition.size ? "end-offset" : "retention-start-unverified");
-          }
-          if (events.length >= maxEvents) finish("limit");
-          else if (maxScannedRecords !== null && scannedCount >= maxScannedRecords) finish("scan-limit");
         }
       }).catch(reject);
     });
+    timing.scanMs = Date.now() - stageStartedAtMs;
   } catch (error) {
     if (error?.code === "kafka_query_aborted") completionReason = "aborted";
     else throw error;
   } finally {
+    stageStartedAtMs = Date.now();
     if (timer) clearTimeout(timer);
     removeAbortListener?.();
     if (running) await boundedKafkaPromise(consumer.stop(), remainingKafkaQueryCleanupBudget(deadlineMs)).catch(() => undefined);
     await boundedDisconnect(consumer, remainingKafkaQueryCleanupBudget(deadlineMs));
+    timing.cleanupMs = Date.now() - stageStartedAtMs;
   }
   return kafkaQueryResult();
 
@@ -1126,12 +1177,51 @@ export async function queryKafkaEventStream({ env = process.env, stream = "hwlab
       scanLimit: maxScannedRecords,
       timeoutMs: budgetMs,
       filters: compactObject({ traceId, sessionId, runId, commandId }),
+      partitionKeyScoped: targetPartition !== null,
+      targetPartition,
+      timing: {
+        ...timing,
+        totalMs: Date.now() - queryStartedAtMs,
+        valuesPrinted: false
+      },
       events,
       valuesPrinted: false
     };
   }
 }
 
+export function kafkaPartitionForKey(partitionKey, topic, offsets = []) {
+  const key = stringValue(partitionKey);
+  const partitionMetadata = offsets
+    .map((entry) => integerValue(entry?.partition))
+    .filter((partition) => partition !== null)
+    .map((partitionId) => ({ partitionId }))
+    .sort((left, right) => left.partitionId - right.partitionId);
+  if (!key || partitionMetadata.length === 0) return null;
+  return Partitioners.DefaultPartitioner()({
+    topic,
+    partitionMetadata,
+    message: { key: Buffer.from(key), value: null, headers: {} }
+  });
+}
+
+async function prepareKafkaQueryGroupOffsets(kafka, { groupId, topic, targetPartition, offsets, deadlineMs, signal }) {
+  const admin = kafka.admin();
+  try {
+    await withinKafkaQueryBudget(admin.connect(), deadlineMs, "query-offset-admin-connect", signal);
+    await withinKafkaQueryBudget(admin.setOffsets({
+      groupId,
+      topic,
+      partitions: offsets.map((entry) => ({
+        partition: entry.partition,
+        offset: entry.partition === targetPartition ? entry.startOffset : entry.endOffset
+      }))
+    }), deadlineMs, "query-offset-admin-set", signal);
+  } finally {
+    await boundedDisconnect(admin, remainingKafkaQueryCleanupBudget(deadlineMs));
+  }
+}
+
 async function kafkaTopicEndOffsetSnapshot(kafka, topic, deadlineMs, signal = null) {
   if (typeof kafka?.admin !== "function") return { available: false, offsets: [], error: { code: "kafka_admin_unavailable", valuesRedacted: true } };
   const admin = kafka.admin();

internal/cloud/server-workbench-realtime-http.test.ts

@@ -146,6 +146,61 @@ test("product realtime uses live Kafka SSE without projection reads", async () =
   }
 });
 
+test("refresh replay scopes the retained Kafka query to the session producer key", async () => {
+  const sessionId = "ses_refresh_partition_key";
+  const queryCalls = [];
+  const subscribers = new Set();
+  const capabilities = { directPublish: true, liveKafkaSse: true, kafkaRefreshReplay: true, transactionalProjector: false, projectionOutboxRelay: false, projectionRealtime: false };
+  const server = createCloudApiServer({
+    accessController: realtimeAccessController({ sessions: [{ id: sessionId, ownerUserId: ACTOR.id }] }),
+    workbenchRuntime: {},
+    kafkaEventBridge: {
+      capabilities,
+      refreshReplay: { groupIdPrefix: "hwlab-test-refresh", timeoutMs: 5000, scanLimit: 100, matchedEventLimit: 20, liveBufferLimit: 20 },
+      ready: Promise.resolve(),
+      subscribeLiveHwlabEvents(listener) { subscribers.add(listener); return () => subscribers.delete(listener); },
+      async queryHwlabEventRetention(options) {
+        queryCalls.push(options);
+        return {
+          topic: "hwlab.event.v1",
+          events: [],
+          completionReason: "end-offset",
+          completion: { reason: "end-offset", complete: true, barrierReached: true, retentionStartVerified: true },
+          reachedEndOffsets: true,
+          endOffsetsAvailable: true,
+          endOffsets: [{ partition: 2, startOffset: "0", endOffset: "10" }],
+          scannedCount: 10,
+          matchedCount: 0
+        };
+      },
+      async stop() {}
+    },
+    env: {
+      HWLAB_KAFKA_DIRECT_PUBLISH_ENABLED: "true",
+      HWLAB_WORKBENCH_LIVE_KAFKA_SSE_ENABLED: "true",
+      HWLAB_WORKBENCH_KAFKA_REFRESH_REPLAY_ENABLED: "true",
+      HWLAB_KAFKA_TRANSACTIONAL_PROJECTOR_ENABLED: "false",
+      HWLAB_KAFKA_PROJECTION_OUTBOX_RELAY_ENABLED: "false",
+      HWLAB_WORKBENCH_PROJECTION_REALTIME_ENABLED: "false",
+      HWLAB_WORKBENCH_KAFKA_REFRESH_REPLAY_GROUP_PREFIX: "hwlab-test-refresh",
+      HWLAB_WORKBENCH_KAFKA_REFRESH_REPLAY_TIMEOUT_MS: "5000",
+      HWLAB_WORKBENCH_KAFKA_REFRESH_REPLAY_SCAN_LIMIT: "100",
+      HWLAB_WORKBENCH_KAFKA_REFRESH_REPLAY_MATCHED_EVENT_LIMIT: "20",
+      HWLAB_WORKBENCH_KAFKA_REFRESH_REPLAY_LIVE_BUFFER_LIMIT: "20"
+    }
+  });
+  await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve));
+  try {
+    const events = await getSseEvents(server.address().port, `/v1/workbench/events?sessionId=${sessionId}`, 1);
+    assert.equal(events[0].event, "workbench.connected");
+    assert.equal(queryCalls.length, 1);
+    assert.equal(queryCalls[0].sessionId, sessionId);
+    assert.equal(queryCalls[0].partitionKey, sessionId);
+  } finally {
+    await new Promise((resolve, reject) => server.close((error) => error ? reject(error) : resolve()));
+  }
+});
+
 test("synchronous projection notification is buffered until the unique initial snapshot completes", async () => {
   const sessionId = "ses_realtime_initial_barrier";
   const traceId = "trc_realtime_initial_barrier";

internal/cloud/server-workbench-realtime-http.ts

@@ -765,6 +765,7 @@ async function handleKafkaRefreshReplayWorkbenchRealtimeHttp(request, response,
       scanLimit: refreshReplay.scanLimit,
       timeoutMs: refreshReplay.timeoutMs,
       groupIdPrefix: refreshReplay.groupIdPrefix,
+      partitionKey: requestedSessionId,
       fromBeginning: true,
       signal
     }),