diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index e199decf35260..bcec0a13cf20a 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -3454,6 +3454,19 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio } if (command.type === "thread.settle") { const projection = yield* loadProjectionForCommand(command, ["runs", "runtimeRequests"]); + if ( + command.expectedLatestRunId !== undefined && + projection.runs.reduce( + (latest, run) => (latest === undefined || run.ordinal > latest.ordinal ? run : latest), + undefined, + )?.id !== command.expectedLatestRunId + ) { + return yield* new OrchestratorDispatchError({ + commandId: command.commandId, + commandType: command.type, + cause: "The thread received another run before settlement.", + }); + } const queuedMessages = yield* projectionStore .getMessagesById( command.threadId, diff --git a/apps/server/src/orchestration-v2/ThreadManagementService.test.ts b/apps/server/src/orchestration-v2/ThreadManagementService.test.ts index 9285f72d46f56..742616f27bd3e 100644 --- a/apps/server/src/orchestration-v2/ThreadManagementService.test.ts +++ b/apps/server/src/orchestration-v2/ThreadManagementService.test.ts @@ -559,28 +559,108 @@ it.effect.each(["live", "no-session", "unsupported-runtime"] as const)( ); it.effect.each([ - { status: "completed" as const, settles: true }, - { status: "failed" as const, settles: false }, - { status: "interrupted" as const, settles: false }, -])("settleAfterRun settles only when the run $status", ({ status, settles }) => + { status: "completed" as const, followUp: "none", settles: true }, + { status: "failed" as const, followUp: "none", settles: false }, + { status: "interrupted" as const, followUp: "none", settles: false }, + { status: "completed" as const, followUp: "notification", settles: true }, + { status: "completed" as const, followUp: "replayed-notification", settles: true }, + { status: "completed" as const, followUp: "delegated", settles: true }, + { status: "completed" as const, followUp: "failed-notification", settles: false }, + { status: "completed" as const, followUp: "held-notification", settles: true }, + { status: "completed" as const, followUp: "user", settles: false }, + { status: "completed" as const, followUp: "steered-notification", settles: false }, +])("settleAfterRun handles $status with $followUp follow-up", ({ status, followUp, settles }) => Effect.gen(function* () { const projectId = ProjectId.make("project:thread-management:settle-after-run"); const threadId = ThreadId.make("thread:thread-management:settle-after-run"); const runId = RunId.make("run:thread-management:settle-after-run"); + const followUpId = RunId.make(`${runId}:follow-up`); + const subscribed = yield* Deferred.make(); + const events = yield* Queue.unbounded(); + const replayRejected = yield* Deferred.make(); + let sequence = 0; + let sequenceReads = 0; + let rejectedCommandId: CommandId | undefined; + const runs = [{ id: runId, ordinal: 1, status, userMessageId: MessageId.make("original") }]; + const nextRun = { + id: followUpId, + ordinal: 2, + status: "running" as OrchestrationV2Run["status"], + userMessageId: MessageId.make("follow-up"), + }; const dispatched: Array = []; + // The notification starts after the completed run was observed, before settlement commits. + let raced = false; const layerTest = ThreadManagementService.layer.pipe( Layer.provide(Layer.mock(ProviderSessionManager.ProviderSessionManagerV2)({})), Layer.provide( Layer.mock(Orchestrator.OrchestratorV2)({ - getThreadEventSequence: () => Effect.succeed(0), + getThreadEventSequence: () => + Effect.sync(() => { + if (followUp === "replayed-notification" && sequenceReads++ === 2) sequence = 1; + return sequence; + }), restartContinuationPending: () => Effect.succeed(false), - getThreadRecords: () => - Effect.succeed({ - thread: { id: threadId, projectId, deletedAt: null }, - runs: [{ id: runId, status }], - } as unknown as OrchestrationV2ThreadProjection), + streamStoredEventsFrom: () => + Stream.fromEffect(Deferred.succeed(subscribed, undefined)).pipe( + Stream.drain, + Stream.concat(Stream.fromQueue(events)), + ), + getThreadRecords: (_threadId, _fields, filter) => + Effect.sync( + () => + ({ + thread: { id: threadId, projectId, deletedAt: null }, + runs: raced ? [...runs, nextRun] : runs, + messages: [ + { id: MessageId.make("original"), runId, role: "user" as const }, + { + id: nextRun.userMessageId, + runId: followUpId, + role: "user" as const, + ...(followUp.endsWith("notification") ? { notification: {} } : {}), + ...(followUp === "delegated" ? { delegatedCompletion: {} } : {}), + }, + { id: MessageId.make("reply"), runId: followUpId, role: "assistant" as const }, + ...(followUp === "steered-notification" + ? [ + { + id: MessageId.make("user-steer"), + runId: followUpId, + role: "user" as const, + }, + ] + : []), + ].filter( + (message) => + (filter?.messageIds === undefined || + filter.messageIds.includes(message.id)) && + (filter?.messageRunIds === undefined || + filter.messageRunIds.includes(message.runId)) && + (filter?.messageRoles === undefined || + filter.messageRoles.includes(message.role)), + ), + }) as unknown as OrchestrationV2ThreadProjection, + ), dispatch: (command) => - Effect.sync(() => { + Effect.gen(function* () { + if (followUp !== "none" && !raced) { + raced = true; + rejectedCommandId = command.commandId; + return yield* new Orchestrator.OrchestratorDispatchError({ + commandId: command.commandId, + commandType: command.type, + cause: "Thread is still running", + }); + } + if (followUp === "replayed-notification" && command.commandId === rejectedCommandId) { + yield* Deferred.succeed(replayRejected, undefined); + return yield* new Orchestrator.OrchestratorCommandPreviouslyRejectedError({ + commandId: command.commandId, + commandType: command.type, + detail: "Thread is still running", + }); + } dispatched.push(command.type); return { sequence: 1, storedEvents: [] } as never; }), @@ -592,9 +672,33 @@ it.effect.each([ const service = yield* ThreadManagementService.ThreadManagementService.pipe( Effect.provide(layerTest), ); - - yield* service.settleAfterRun({ projectId, threadId, runId }); - + const fiber = yield* service + .settleAfterRun({ projectId, threadId, runId }) + .pipe(Effect.forkChild); + if (followUp !== "none") { + yield* Deferred.await(subscribed); + expect(dispatched).toEqual([]); + nextRun.status = + followUp === "failed-notification" + ? "failed" + : followUp === "held-notification" + ? "queued" + : "completed"; + yield* Queue.offer(events, { + sequence: 1, + event: { type: "run.updated", threadId, payload: nextRun }, + } as unknown as OrchestrationV2StoredEvent); + } + if (followUp === "replayed-notification") { + yield* Deferred.await(replayRejected); + expect(dispatched).toEqual([]); + sequence = 2; + yield* Queue.offer(events, { + sequence, + event: { type: "run.updated", threadId, payload: nextRun }, + } as unknown as OrchestrationV2StoredEvent); + } + yield* Fiber.join(fiber); expect(dispatched).toEqual(settles ? ["thread.settle"] : []); }), ); diff --git a/apps/server/src/orchestration-v2/ThreadManagementService.ts b/apps/server/src/orchestration-v2/ThreadManagementService.ts index b6c1e27120ae8..46ea454310554 100644 --- a/apps/server/src/orchestration-v2/ThreadManagementService.ts +++ b/apps/server/src/orchestration-v2/ThreadManagementService.ts @@ -446,8 +446,8 @@ export interface ThreadManagementServiceShape { readonly waitForThread: ( input: ThreadManagementWaitInput, ) => Effect.Effect; - /** Waits for the agent's run to end before settling its thread and stopping the session. - * Other endings or settle refusals leave the thread active. */ + /** Waits for the agent's run and automatic follow-ups before settling. + * A new user turn or an unsuccessful ending leaves the thread active. */ readonly settleAfterRun: (input: { readonly projectId: ProjectId; readonly threadId: ThreadId; @@ -1044,11 +1044,74 @@ const make = Effect.gen(function* () { )(function* (input) { const { run } = yield* waitForThread({ ...input, timeoutMs: SETTLE_AFTER_RUN_WAIT_MS }); if (run?.status !== "completed") return; - yield* dispatch({ - type: "thread.settle", - commandId: CommandId.make(`server:settle-after-run:${input.runId}`), - threadId: input.threadId, + const afterSequence = yield* orchestrator.getThreadEventSequence(input.threadId); + const trySettle = Effect.gen(function* () { + const sequence = yield* orchestrator.getThreadEventSequence(input.threadId); + const records = yield* getProjectThreadRecords(input, ["runs"]); + const following = records.runs.filter((candidate) => candidate.ordinal > run.ordinal); + if (following.length > 0) { + const { messages } = yield* getProjectThreadRecords(input, ["messages"], { + messageRunIds: following.map((candidate) => candidate.id), + messageRoles: ["user"], + }); + const automaticMessageIds = new Set( + messages + .filter( + (message) => + message.role === "user" && + (message.notification !== undefined || message.delegatedCompletion !== undefined), + ) + .map((message) => message.id), + ); + if ( + messages.some((message) => !automaticMessageIds.has(message.id)) || + following.some((candidate) => !automaticMessageIds.has(candidate.userMessageId)) + ) { + return true; + } + if ( + following.some( + (candidate) => candidate.status !== "queued" && !isTerminalRunStatus(candidate.status), + ) + ) + return false; + if ( + following.some( + (candidate) => candidate.status !== "queued" && candidate.status !== "completed", + ) + ) + return true; + } + const latest = latestRun(records); + if (latest === undefined) return true; + return yield* dispatch({ + type: "thread.settle", + commandId: CommandId.make(`server:settle-after-run:${input.runId}:${sequence}`), + threadId: input.threadId, + expectedLatestRunId: latest.id, + }).pipe( + Effect.as(true), + // A notification can start between the projection read and the guarded dispatch. + Effect.catchTags({ + OrchestratorDispatchError: () => Effect.succeed(false), + OrchestratorCommandPreviouslyRejectedError: () => Effect.succeed(false), + }), + ); }); + yield* Stream.fromEffect(trySettle).pipe( + Stream.concat( + orchestrator + .streamStoredEventsFrom({ + threadId: input.threadId, + afterSequence, + eventType: "run.updated", + }) + .pipe(Stream.mapEffect(() => trySettle)), + ), + Stream.filter((done) => done), + Stream.runHead, + Effect.timeoutOption(Duration.millis(SETTLE_AFTER_RUN_WAIT_MS)), + ); }); const settleThread: ThreadManagementServiceShape["settleThread"] = Effect.fn( diff --git a/apps/server/src/orchestration-v2/runtimeLayer.test.ts b/apps/server/src/orchestration-v2/runtimeLayer.test.ts index af95d71e6121b..b90f491dfa94e 100644 --- a/apps/server/src/orchestration-v2/runtimeLayer.test.ts +++ b/apps/server/src/orchestration-v2/runtimeLayer.test.ts @@ -4859,11 +4859,29 @@ it.layer(layerTest)("OrchestrationV2LayerLive lifecycle", (it) => { }); }); yield* holdQueueAfterRestart("settle-automatic-restart"); + const beforeSettle = yield* orchestrator.getThreadProjection(threadId); + const oldestRun = beforeSettle.runs.reduce((oldest, run) => + run.ordinal < oldest.ordinal ? run : oldest, + ); + const latestRun = beforeSettle.runs.reduce((latest, run) => + run.ordinal > latest.ordinal ? run : latest, + ); + const staleSettle = yield* orchestrator + .dispatch({ + type: "thread.settle", + commandId: CommandId.make("settle-automatic-stale"), + threadId, + expectedLatestRunId: oldestRun.id, + }) + .pipe(Effect.flip); + assert.equal(staleSettle._tag, "OrchestratorDispatchError"); + assert.isNull((yield* orchestrator.getThreadProjection(threadId)).thread.settledAt); yield* orchestrator.dispatch({ type: "thread.settle", commandId: CommandId.make("settle-automatic-settle"), threadId, + expectedLatestRunId: latestRun.id, }); const settled = yield* orchestrator.getThreadProjection(threadId); assert.isNotNull(settled.thread.settledAt); diff --git a/packages/contracts/src/orchestrationV2.ts b/packages/contracts/src/orchestrationV2.ts index 8234864349504..461ec5e4b3f3e 100644 --- a/packages/contracts/src/orchestrationV2.ts +++ b/packages/contracts/src/orchestrationV2.ts @@ -2975,6 +2975,7 @@ export const OrchestrationV2Command = Schema.Union([ type: Schema.Literal("thread.settle"), commandId: CommandId, threadId: ThreadId, + expectedLatestRunId: Schema.optional(RunId), settledAt: Schema.optional(Schema.DateTimeUtc), }), /** Rejects changes since snapshotAt or any explicit settled override to preserve user actions. */