diff --git a/packages/junior/src/chat/task-execution/state.ts b/packages/junior/src/chat/task-execution/state.ts index 03ec653cf0..2fda06e081 100644 --- a/packages/junior/src/chat/task-execution/state.ts +++ b/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 { + return normalizeStop(await state.get(conversationStopKey(conversationId))); +} + +async function fenceConversationMutation( + state: StateAdapter, + lock: Lock, + conversationId: string, +): Promise { + 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 { + 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 { 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, diff --git a/packages/junior/src/chat/task-execution/store.ts b/packages/junior/src/chat/task-execution/store.ts index b2dc648e11..ae4f8de389 100644 --- a/packages/junior/src/chat/task-execution/store.ts +++ b/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, diff --git a/packages/junior/src/chat/task-execution/worker.ts b/packages/junior/src/chat/task-execution/worker.ts index 03937b019f..95683c8473 100644 --- a/packages/junior/src/chat/task-execution/worker.ts +++ b/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 => { 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, }); diff --git a/packages/junior/tests/component/task-execution/conversation-work.test.ts b/packages/junior/tests/component/task-execution/conversation-work.test.ts index 534c7494a3..ca1518db6e 100644 --- a/packages/junior/tests/component/task-execution/conversation-work.test.ts +++ b/packages/junior/tests/component/task-execution/conversation-work.test.ts @@ -17,6 +17,7 @@ import { countPendingConversationMessages, drainConversationMailbox, getConversationWorkState, + hasConversationStop, listActiveConversationIds, listConversationsByActivity, ackMessages, @@ -1695,7 +1696,25 @@ describe("conversation work execution", () => { vi.useFakeTimers({ now: 1_000 }); let currentNowMs = 1_000; const queue = createConversationWorkQueueTestAdapter(); - await appendInboundMessage({ message: inboundMessage("m1"), nowMs: 1_000 }); + const state = getStateAdapter(); + const workerGetKeys: string[] = []; + const workerState = new Proxy(state, { + get(target, prop) { + if (prop === "get") { + return async (key: string) => { + workerGetKeys.push(key); + return target.get(key); + }; + } + const value = readProxyProperty(target, prop); + return typeof value === "function" ? value.bind(target) : value; + }, + }) as StateAdapter; + await appendInboundMessage({ + message: inboundMessage("m1"), + nowMs: 1_000, + state, + }); const entered = deferred(); const stopObserved = deferred(); const finishRun = deferred(); @@ -1704,6 +1723,7 @@ describe("conversation work execution", () => { nowMs: () => currentNowMs, queue, softYieldAfterMs: 1_000, + state: workerState, run: async (context) => { await context.attempt.ack(); const signal = context.stopSignal?.(); @@ -1723,22 +1743,33 @@ describe("conversation work execution", () => { }); await entered.promise; + const conversationReadsBeforeStop = workerGetKeys.filter( + (key) => key === CONVERSATION_WORK_STATE_KEY, + ).length; await appendInboundMessage({ message: inboundMessage("m2", { createdAtMs: 1_500, receivedAtMs: 1_500, }), nowMs: 1_500, + state, }); await expect( stopConversationWork({ conversationId: CONVERSATION_ID, nowMs: 2_000, + state, }), ).resolves.toMatchObject({ status: "requested" }); await vi.advanceTimersByTimeAsync(500); await stopObserved.promise; + expect( + workerGetKeys.filter((key) => key === CONVERSATION_WORK_STATE_KEY), + ).toHaveLength(conversationReadsBeforeStop); + expect( + workerGetKeys.includes(`junior:conversation:v2:stop:${CONVERSATION_ID}`), + ).toBe(true); await appendInboundMessage({ message: inboundMessage("m3", { @@ -1748,19 +1779,65 @@ describe("conversation work execution", () => { receivedAtMs: 2_000, }), nowMs: 2_000, + state, }); currentNowMs = 2_000; finishRun.resolve(); await expect(running).resolves.toEqual({ status: "pending_requeued" }); await expect( - getConversationWorkState({ conversationId: CONVERSATION_ID }), + getConversationWorkState({ conversationId: CONVERSATION_ID, state }), ).resolves.toMatchObject({ - execution: { stop: undefined }, messages: [expect.objectContaining({ inboundMessageId: "m3" })], }); }); + it("does not apply a stale stop marker to a new run", async () => { + const queue = createConversationWorkQueueTestAdapter(); + const entered = deferred(); + const finish = deferred(); + await appendInboundMessage({ message: inboundMessage("m1"), nowMs: 1_000 }); + + const first = processConversationWork(conversationQueueMessage(), { + queue, + run: async (context) => { + await context.attempt.ack(); + entered.resolve(); + await finish.promise; + return { status: "completed" }; + }, + }); + await entered.promise; + await expect( + stopConversationWork({ + conversationId: CONVERSATION_ID, + nowMs: 2_000, + }), + ).resolves.toMatchObject({ status: "requested" }); + finish.resolve(); + await expect(first).resolves.toEqual({ status: "completed" }); + + await appendInboundMessage({ + message: inboundMessage("m2", { + createdAtMs: 3_000, + receivedAtMs: 3_000, + }), + nowMs: 3_000, + }); + await expect( + processConversationWork(conversationQueueMessage(), { + queue, + run: async (context) => { + const signal = context.stopSignal?.(); + if (!signal) throw new Error("Expected a Conversation stop signal"); + expect(signal.aborted).toBe(false); + await context.attempt.ack(); + return { status: "completed" }; + }, + }), + ).resolves.toEqual({ status: "completed" }); + }); + it("resumes a paused Turn when its stop missed the live poll", async () => { vi.useFakeTimers({ now: 1_000 }); let currentNowMs = 1_000; @@ -1794,14 +1871,16 @@ describe("conversation work execution", () => { finish.resolve(); await expect(running).resolves.toEqual({ status: "pending_requeued" }); - await expect( - getConversationWorkState({ conversationId: CONVERSATION_ID }), - ).resolves.toMatchObject({ - execution: { - status: "paused", - stop: expect.objectContaining({ inboundMessageIds: [] }), - }, + const paused = await getConversationWorkState({ + conversationId: CONVERSATION_ID, }); + expect(paused?.execution.status).toBe("paused"); + await expect( + hasConversationStop({ + conversationId: CONVERSATION_ID, + runId: paused!.execution.runId!, + }), + ).resolves.toBe(true); await expect( processConversationWork(queue.takeMessage(), { @@ -1818,7 +1897,7 @@ describe("conversation work execution", () => { await expect( getConversationWorkState({ conversationId: CONVERSATION_ID }), ).resolves.toMatchObject({ - execution: { status: "idle", stop: undefined }, + execution: { status: "idle" }, }); }); @@ -1875,14 +1954,18 @@ describe("conversation work execution", () => { releaseCompletion.resolve(); await expect(running).resolves.toEqual({ status: "pending_requeued" }); - await expect( - getConversationWorkState({ conversationId: CONVERSATION_ID, state }), - ).resolves.toMatchObject({ - execution: { - status: "paused", - stop: expect.objectContaining({ runId: expect.any(String) }), - }, + const paused = await getConversationWorkState({ + conversationId: CONVERSATION_ID, + state, }); + expect(paused?.execution.status).toBe("paused"); + await expect( + hasConversationStop({ + conversationId: CONVERSATION_ID, + runId: paused!.execution.runId!, + state, + }), + ).resolves.toBe(true); await expect( processConversationWork(queue.takeMessage(), { @@ -1899,7 +1982,7 @@ describe("conversation work execution", () => { await expect( getConversationWorkState({ conversationId: CONVERSATION_ID, state }), ).resolves.toMatchObject({ - execution: { status: "idle", stop: undefined }, + execution: { status: "idle" }, }); });