perf(chat): use lightweight stop markers

packages/junior/src/chat/task-execution/state.ts

@@ -135,7 +135,7 @@ export interface Lease {
   token: string;
 }
 
-/** Durable request to stop the current Conversation run. */
+/** Durable request to stop one Conversation run. */
 export interface ConversationStop {
   inboundMessageIds: string[];
   runId: string;
@@ -152,7 +152,6 @@ export interface ConversationExecution {
   pendingMessages: InboundMessage[];
   runId?: string;
   status: ExecutionStatus;
-  stop?: ConversationStop;
   updatedAtMs?: number;
 }
 
@@ -268,6 +267,10 @@ function conversationKey(conversationId: string): string {
   return `${CONVERSATION_PREFIX}:${conversationId}`;
 }
 
+function conversationStopKey(conversationId: string): string {
+  return `${CONVERSATION_PREFIX}:stop:${conversationId}`;
+}
+
 function indexLockKey(indexKey: string): string {
   return `${indexKey}:lock`;
 }
@@ -469,8 +472,6 @@ function normalizeExecution(
 
   const lease = normalizeLease(value.lease);
   const runId = toOptionalString(value.runId);
-  const normalizedStop = normalizeStop(value.stop);
-  const stop = normalizedStop?.runId === runId ? normalizedStop : undefined;
   const normalizedStatus =
     status === "idle" && lease
       ? "running"
@@ -492,7 +493,6 @@ function normalizeExecution(
     lastProgressAtMs: toOptionalNumber(value.lastProgressAtMs),
     retryCount: toOptionalNumber(value.retryCount),
     runId,
-    stop,
     updatedAtMs: toOptionalNumber(value.updatedAtMs),
   };
 }
@@ -1011,13 +1011,7 @@ async function writeConversation(
     ...conversation,
     execution,
   };
-  const fenced = await state.extendLock(
-    lock,
-    CONVERSATION_MUTATION_LOCK_TTL_MS,
-  );
-  if (!fenced) {
-    throw new ConversationMutationFencedError(next.conversationId);
-  }
+  await fenceConversationMutation(state, lock, next.conversationId);
   // TODO(dcramer): Remove the stored publishExternally field after no deployed
   // mailbox reader requires it. Destination is used only to write this old
   // Redis shape. Current workers ignore the field.
@@ -1740,6 +1734,36 @@ function isHumanFacingMessage(message: InboundMessage): boolean {
   return message.source === "web" || message.source === "slack";
 }
 
+async function readConversationStop(
+  state: StateAdapter,
+  conversationId: string,
+): Promise<ConversationStop | undefined> {
+  return normalizeStop(await state.get(conversationStopKey(conversationId)));
+}
+
+async function fenceConversationMutation(
+  state: StateAdapter,
+  lock: Lock,
+  conversationId: string,
+): Promise<void> {
+  if (!(await state.extendLock(lock, CONVERSATION_MUTATION_LOCK_TTL_MS))) {
+    throw new ConversationMutationFencedError(conversationId);
+  }
+}
+
+/** Return whether one Conversation run has a durable stop request. */
+export async function hasConversationStop(args: {
+  conversationId: string;
+  runId: string;
+  state?: StateAdapter;
+}): Promise<boolean> {
+  const state = await getConnectedState(args.state);
+  return (
+    (await readConversationStop(state, args.conversationId))?.runId ===
+    args.runId
+  );
+}
+
 /** Persist a stop request for the current run without process affinity. */
 export async function stopConversationWork(args: {
   conversationId: string;
@@ -1757,18 +1781,22 @@ export async function stopConversationWork(args: {
     const inboundMessageIds = current.execution.pendingMessages
       .filter(isHumanFacingMessage)
       .map((message) => message.inboundMessageId);
-    await writeConversation(
-      state,
-      lock,
-      withExecutionUpdate(
-        current,
-        {
-          ...current.execution,
-          runId,
-          stop: { inboundMessageIds, runId },
-        },
-        nowMs,
-      ),
+    if (current.execution.runId === undefined) {
+      await writeConversation(
+        state,
+        lock,
+        withExecutionUpdate(
+          current,
+          { ...current.execution, runId },
+          nowMs,
+        ),
+      );
+    }
+    await fenceConversationMutation(state, lock, args.conversationId);
+    await state.set(
+      conversationStopKey(args.conversationId),
+      { inboundMessageIds, runId } satisfies ConversationStop,
+      JUNIOR_THREAD_STATE_TTL_MS,
     );
     return { runId, status: "requested" };
   });
@@ -1788,7 +1816,7 @@ export async function completeConversationStop(args: {
     if (!current || current.execution.lease?.token !== args.leaseToken) {
       return { status: "lost_lease", removedInboundMessageIds: [] };
     }
-    const stop = current.execution.stop;
+    const stop = await readConversationStop(state, args.conversationId);
     if (!stop || stop.runId !== args.runId) {
       return { status: "none", removedInboundMessageIds: [] };
     }
@@ -1818,11 +1846,12 @@ export async function completeConversationStop(args: {
               ? undefined
               : current.execution.lastEnqueuedAtMs,
           pendingMessages,
-          stop: undefined,
         },
         nowMs,
       ),
     );
+    await fenceConversationMutation(state, lock, args.conversationId);
+    await state.delete(conversationStopKey(args.conversationId));
     return { status: "cleared", removedInboundMessageIds };
   });
 }
@@ -1884,7 +1913,6 @@ export async function cancelHumanFacingPendingMessages(args: {
           retryCount: becomesIdle ? 0 : current.execution.retryCount,
           runId: becomesIdle ? undefined : current.execution.runId,
           status: becomesIdle ? "idle" : current.execution.status,
-          stop: becomesIdle ? undefined : current.execution.stop,
         },
         nowMs,
       ),
@@ -2023,7 +2051,6 @@ export async function recordConversationRetry(args: {
           pendingMessages: stopped ? [] : current.execution.pendingMessages,
           runId: stopped ? undefined : current.execution.runId,
           status: stopped ? "failed" : "paused",
-          stop: stopped ? undefined : current.execution.stop,
         },
         nowMs,
       ),
@@ -2049,9 +2076,11 @@ export async function completeConversationWork(args: {
       return "lost_lease";
     }
     const hasPending = pendingMessages(current).length > 0;
-    const needsRun =
-      current.execution.status === "paused" ||
-      (args.resumeIfStopped === true && current.execution.stop !== undefined);
+    const stopped =
+      args.resumeIfStopped === true &&
+      (await readConversationStop(state, args.conversationId))?.runId ===
+        current.execution.runId;
+    const needsRun = current.execution.status === "paused" || stopped;
     const runnable = needsRun || hasPending;
     await writeConversation(
       state,
@@ -2071,7 +2100,6 @@ export async function completeConversationWork(args: {
               ? current.execution.retryCount
               : 0,
           runId: runnable ? current.execution.runId : undefined,
-          stop: runnable ? current.execution.stop : undefined,
         },
         nowMs,
       ),
@@ -2206,7 +2234,6 @@ export async function deadLetterAttempt(args: {
           lease: undefined,
           status: runnable ? "pending" : "failed",
           runId: runnable ? current.execution.runId : undefined,
-          stop: runnable ? current.execution.stop : undefined,
         },
         nowMs,
       ),
@@ -2245,7 +2272,6 @@ export async function clearExpiredConversationLease(args: {
           pendingMessages: stopped ? [] : current.execution.pendingMessages,
           runId: stopped ? undefined : current.execution.runId,
           status: stopped ? "failed" : "paused",
-          stop: stopped ? undefined : current.execution.stop,
         },
         nowMs,
       ),
@@ -2261,6 +2287,7 @@ export async function deleteConversationState(args: {
 }): Promise<void> {
   await withConversationMutation(args, async (state) => {
     await state.delete(conversationKey(args.conversationId));
+    await state.delete(conversationStopKey(args.conversationId));
     await removeIndexEntry({
       state,
       indexKey: CONVERSATION_ACTIVE_INDEX_KEY,

packages/junior/src/chat/task-execution/store.ts

@@ -14,6 +14,7 @@ export {
   CONVERSATION_WORK_MAX_DELIVERY_ATTEMPTS,
   CONVERSATION_WORK_MAX_RETRIES,
   CONVERSATION_WORK_STALE_ENQUEUE_MS,
+  hasConversationStop,
   isFinalAttempt,
   isInvalidConversationRecordError,
   type AgentInput,

packages/junior/src/chat/task-execution/worker.ts

@@ -24,6 +24,7 @@ import {
   drainConversationMailbox,
   ensureConversationWake,
   getConversationWorkState,
+  hasConversationStop,
   isFinalAttempt,
   isInvalidConversationRecordError,
   recordAttemptFailure,
@@ -285,10 +286,10 @@ function startLeaseCheckIn(args: {
   return timer;
 }
 
-/** Poll shared state only when a worker adapter asks to observe remote stops. */
+/** Poll the run-scoped stop marker only when an adapter observes remote stops. */
 function createConversationStopSignal(args: {
   conversationId: string;
-  initialStopRunId?: string;
+  initiallyStopped: boolean;
   options: ProcessConversationWorkOptions;
   runId: string;
 }) {
@@ -303,21 +304,18 @@ function createConversationStopSignal(args: {
       controller.abort(new Error("Conversation work stopped"));
     }
   };
-  if (args.initialStopRunId === args.runId) {
-    requestStop();
-  }
+  if (args.initiallyStopped) requestStop();
 
   const check = async (): Promise<void> => {
     if (checking || controller.signal.aborted) return;
     checking = true;
     try {
-      const current = await getConversationWorkState({
+      const stopped = await hasConversationStop({
         conversationId: args.conversationId,
+        runId: args.runId,
         state: args.options.state,
       });
-      if (current?.execution.stop?.runId === args.runId) {
-        requestStop();
-      }
+      if (stopped) requestStop();
     } catch (error) {
       if (!failureCaptured) {
         failureCaptured = true;
@@ -633,7 +631,11 @@ async function processConversationWorkInContext(
       };
       const stop = createConversationStopSignal({
         conversationId,
-        initialStopRunId: leasedWork.execution.stop?.runId,
+        initiallyStopped: await hasConversationStop({
+          conversationId,
+          runId,
+          state: options.state,
+        }),
         options,
         runId,
       });