diff --git a/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts b/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts index 5cabb1c363722..4c2de8698ac9d 100644 --- a/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts +++ b/apps/server/src/orchestration-v2/DelegatedCompletionDelivery.test.ts @@ -1,4 +1,5 @@ import * as SourceControlProviderRegistry from "../sourceControl/SourceControlProviderRegistry.ts"; +import { backgroundWorkHoldsCompletion } from "@t3tools/shared/orchestrationV2PendingBackgroundWork"; import * as McpProviderSessions from "@t3tools/provider-core/server/McpProviderSessions"; import * as NodeCrypto from "@effect/platform-node/NodeCrypto"; import * as NodeServices from "@effect/platform-node/NodeServices"; @@ -2242,6 +2243,118 @@ it.layer(Layer.merge(layerTest, ProviderContinuationRequests.layer))( assert.notEqual(deliveryRun?.status, "queued"); }), ); + it.effect( + "keeps the parent shell unsettled while its delegated completion wake is pending", + () => + Effect.gen(function* () { + const orchestrator = yield* Orchestrator.OrchestratorV2; + const sink = yield* EventSink.EventSinkV2; + const requests = yield* ProviderContinuationRequests.ProviderContinuationRequests; + const now = yield* DateTime.now; + const threadId = ThreadId.make("thread:pending-wake-shell"); + yield* seedParentWithTerminalTask({ + threadId, + projectId: ProjectId.make("project:pending-wake-shell"), + runId: RunId.make("run:pending-wake-shell-seed"), + rootNodeId: NodeId.make("node:pending-wake-shell-seed"), + taskId: NodeId.make("node:pending-wake-shell-seed-task"), + deliveryState: "acknowledged", + parentRunStatus: "completed", + now, + }); + const messageId = MessageId.make("message:pending-wake-shell"); + yield* orchestrator.dispatch({ + type: "message.dispatch", + createdBy: "user", + creationSource: "web", + commandId: CommandId.make("command:pending-wake-shell-send"), + threadId, + messageId, + text: "Delegate this and end your turn", + attachments: [], + modelSelection, + dispatchMode: { type: "start_immediately" }, + }); + const parentRun = (yield* orchestrator.getThreadProjection(threadId)).runs.find( + (run) => run.userMessageId === messageId, + )!; + const runId = parentRun.id; + yield* orchestrator.dispatch({ + type: "delegated_task.request", + createdBy: "agent", + creationSource: "mcp", + commandId: CommandId.make("command:pending-wake-shell-delegate"), + parentThreadId: threadId, + parentRunId: runId, + parentNodeId: parentRun.rootNodeId!, + task: "Run in the background", + modelSelection, + runtimeMode: "full-access", + interactionMode: "default", + completionWake: "always", + }); + const started = yield* orchestrator.getThreadProjection(threadId); + const task = started.subagents.find((row) => row.runId === runId && row.childThreadId)!; + yield* sink.write({ + events: [ + { + ...runEvent({ + threadId, + runId, + ordinal: parentRun.ordinal, + status: "completed", + now, + }), + payload: { + ...started.runs.find((run) => run.id === runId)!, + status: "completed", + completedAt: now, + }, + }, + ], + }); + assert.isTrue( + backgroundWorkHoldsCompletion( + (yield* orchestrator.getThreadShell(threadId))?.pendingBackgroundTasks ?? [], + ), + ); + + const childThreadId = task.childThreadId!; + const childRun = (yield* orchestrator.getThreadProjection(childThreadId)).runs[0]!; + yield* sink.write({ + events: [ + { + ...runEvent({ + threadId: childThreadId, + runId: childRun.id, + ordinal: childRun.ordinal, + status: "completed", + now, + }), + payload: { ...childRun, status: "completed", completedAt: now }, + }, + ], + }); + let wake = yield* requests.take; + while (wake.threadId !== threadId) wake = yield* requests.take; + const projection = yield* orchestrator.getThreadProjection(threadId); + assert.equal( + projection.subagents.find((row) => row.id === task.id)?.completionDelivery?.state, + "claimed", + ); + assert.isFalse( + projection.runs.some( + (run) => run.userMessageId === wake.delegatedCompletion?.messageId, + ), + ); + + const shell = yield* orchestrator.getThreadShell(threadId); + assert.isTrue( + backgroundWorkHoldsCompletion(shell?.pendingBackgroundTasks ?? []), + `parent shell reads settled (status ${shell?.status}, no held background work) while its delegated completion wake is still pending`, + ); + }), + ); it.effect("pauses queued follow-up cards and resumes them", () => Effect.gen(function* () { const orchestrator = yield* Orchestrator.OrchestratorV2; diff --git a/apps/server/src/orchestration-v2/NotificationMailbox.ts b/apps/server/src/orchestration-v2/NotificationMailbox.ts index 79458677c1af0..90993eac93b22 100644 --- a/apps/server/src/orchestration-v2/NotificationMailbox.ts +++ b/apps/server/src/orchestration-v2/NotificationMailbox.ts @@ -2,7 +2,10 @@ import { latestProviderTurnForAttempt, type MessageId, type OrchestrationV2ConversationMessage, + type OrchestrationV2Run, + type OrchestrationV2Subagent, type OrchestrationV2ThreadProjection, + type OrchestrationV2TurnItem, } from "@t3tools/contracts"; export function isBufferedProviderNotification( @@ -35,3 +38,79 @@ export function isUndeliveredMailboxSteer( "completed") ); } + +type DelegatedCompletionRecords = Pick< + OrchestrationV2ThreadProjection, + "messages" | "runs" | "providerTurns" +>; + +/** The run's claimed delegated completion still owes a wake: none dispatched, or its steer went unread. */ +export function pendingDelegatedCompletionWake( + projection: DelegatedCompletionRecords, + run: Pick, +) { + const delivery = run.delegatedCompletion?.delivery; + if (run.delegatedCompletion?.disposition !== "open" || delivery == null) return undefined; + const dispatched = projection.messages.some((message) => message.id === delivery.messageId); + return dispatched && !isUndeliveredMailboxSteer(projection, delivery.messageId) + ? undefined + : delivery; +} + +const TERMINAL_TASK_STATUSES: ReadonlyArray = [ + "completed", + "failed", + "cancelled", + "interrupted", +]; + +/** + * Finished delegated task items whose completion wake has not started its run: claimed into a + * pending wake, or left pending behind a running wake, whose end reserves them a successor. A run + * without its provider thread is never offered a wake, so its tasks hold nothing. + */ +export function delegatedTaskItemsAwaitingWake( + projection: DelegatedCompletionRecords & { + readonly providerThreads: ReadonlyArray<{ readonly id: string }>; + readonly subagents: ReadonlyArray< + Pick< + OrchestrationV2Subagent, + "id" | "runId" | "origin" | "status" | "completionWake" | "completionDelivery" + > + >; + readonly turnItems: ReadonlyArray; + }, +): ReadonlyArray { + const taskIds = new Set(); + for (const run of projection.runs) { + if ( + run.providerThreadId === null || + !projection.providerThreads.some((thread) => thread.id === run.providerThreadId) + ) { + continue; + } + for (const taskId of pendingDelegatedCompletionWake(projection, run)?.taskIds ?? []) { + taskIds.add(taskId); + } + if (run.delegatedCompletion?.disposition !== "open" || run.delegatedCompletion.delivery == null) + continue; + // Mirrors the successor reservation: a live run's blocking wait takes settled_only results. + const runIsLive = ["preparing", "starting", "running"].includes(run.status); + for (const task of projection.subagents) { + if ( + task.runId === run.id && + task.origin === "app_owned" && + task.completionDelivery?.state === "pending" && + TERMINAL_TASK_STATUSES.includes(task.status) && + (!runIsLive || task.completionWake === "always") + ) { + taskIds.add(task.id); + } + } + } + if (taskIds.size === 0) return []; + return projection.turnItems.filter( + (item) => + item.type === "subagent" && item.origin === "app_owned" && taskIds.has(item.subagentId), + ); +} diff --git a/apps/server/src/orchestration-v2/Orchestrator.ts b/apps/server/src/orchestration-v2/Orchestrator.ts index eab5cda62027e..9588701134efc 100644 --- a/apps/server/src/orchestration-v2/Orchestrator.ts +++ b/apps/server/src/orchestration-v2/Orchestrator.ts @@ -130,6 +130,7 @@ import { import { isBufferedProviderNotification, isUndeliveredMailboxSteer, + pendingDelegatedCompletionWake, } from "./NotificationMailbox.ts"; import { EventSinkV2 } from "./EventSink.ts"; import { @@ -1557,18 +1558,26 @@ const makeOrchestrator = Effect.fn("orchestrationV2.Orchestrator.layer")(functio messageIds: [delivery.messageId], }); const message = messages.messages[0]; - if (message !== undefined) { - if (message.delegatedCompletion === undefined || message.runId === null) return; - const runs = yield* projectionStore.getThreadRecords(threadId, ["runs"], { - runIds: [message.runId], - }); - const run = runs.runs[0]; - const turns = yield* projectionStore.getThreadRecords(threadId, ["providerTurns"], { - providerTurnAttemptIds: run?.activeAttemptId == null ? [] : [run.activeAttemptId], - }); - if (!isUndeliveredMailboxSteer({ ...messages, ...runs, ...turns }, delivery.messageId)) - return; - } + const runs = + message?.delegatedCompletion === undefined || message.runId === null + ? [] + : (yield* projectionStore.getThreadRecords(threadId, ["runs"], { + runIds: [message.runId], + })).runs; + const attemptId = runs[0]?.activeAttemptId; + const providerTurns = + attemptId == null + ? [] + : (yield* projectionStore.getThreadRecords(threadId, ["providerTurns"], { + providerTurnAttemptIds: [attemptId], + })).providerTurns; + if ( + pendingDelegatedCompletionWake( + { messages: messages.messages, runs, providerTurns }, + parentRun, + ) === undefined + ) + return; const projection = yield* projectionStore.getThreadRecords(threadId, ["providerThreads"], { providerThreadIds: parentRun.providerThreadId === null ? [] : [parentRun.providerThreadId], }); diff --git a/apps/server/src/orchestration-v2/ProjectionStore.test.ts b/apps/server/src/orchestration-v2/ProjectionStore.test.ts index 40b7709870585..427e2b35e0e8c 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.test.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.test.ts @@ -12,6 +12,8 @@ import { MessageId, type ModelSelection, type OrchestrationV2ProviderThread, + type OrchestrationV2Run, + type OrchestrationV2Subagent, type OrchestrationV2TurnItem, type OrchestrationV2ThreadHistoryPage, type OrchestrationV2ThreadShellSnapshot, @@ -1530,6 +1532,238 @@ const threadWorktreePathsIncludeArchivedOnRequest = Effect.gen(function* () { ); }); +const delegatedCompletionWakeHoldsEveryShell = Effect.gen(function* () { + const store = yield* ProjectionStore.ProjectionStoreV2; + const now = yield* DateTime.now; + const threadId = yield* addRolledBackRecoveryCandidate("delegated-wake-hold"); + const providerThreadId = ProviderThreadId.make("provider-thread:delegated-wake-hold"); + yield* store.apply({ + id: EventId.make("event:delegated-wake-hold:provider-thread"), + type: "provider-thread.updated", + threadId, + occurredAt: now, + payload: { + id: providerThreadId, + appThreadId: threadId, + driver, + providerInstanceId, + providerSessionId: null, + ownerNodeId: null, + nativeThreadRef: null, + nativeConversationHeadRef: null, + lastRunOrdinal: null, + handoffIds: [], + forkedFrom: null, + status: "active", + firstRunOrdinal: 1, + createdAt: now, + updatedAt: now, + }, + }); + const taskId = NodeId.make("node:delegated-wake-hold:task"); + const wakeMessageId = MessageId.make("message:delegated-wake-hold:wake"); + const parentRun = { + ...(yield* store.getThreadProjection(threadId)).runs[0]!, + id: RunId.make("run:delegated-wake-hold:parent"), + ordinal: 2, + providerThreadId, + userMessageId: MessageId.make("message:delegated-wake-hold:parent"), + rootNodeId: NodeId.make("node:delegated-wake-hold:parent"), + status: "completed" as const, + completedAt: now, + delegatedCompletion: { + disposition: "open" as const, + nextGeneration: 2, + delivery: { generation: 1, messageId: wakeMessageId, taskIds: [taskId] }, + }, + }; + const applyRun = (name: string, run: OrchestrationV2Run) => + store.apply({ + id: EventId.make(`event:delegated-wake-hold:${name}`), + type: "run.updated", + threadId, + runId: run.id, + driver, + providerInstanceId, + occurredAt: now, + payload: run, + }); + yield* applyRun("parent", parentRun); + const itemId = TurnItemId.make("item:delegated-wake-hold:task"); + yield* store.apply({ + id: EventId.make("event:delegated-wake-hold:item"), + type: "turn-item.updated", + threadId, + runId: parentRun.id, + nodeId: taskId, + driver, + occurredAt: now, + payload: { + id: itemId, + type: "subagent", + threadId, + runId: parentRun.id, + nodeId: taskId, + providerThreadId, + providerTurnId: null, + nativeItemRef: null, + parentItemId: null, + ordinal: 10, + status: "completed", + title: "Delegated task", + startedAt: now, + completedAt: now, + updatedAt: now, + subagentId: taskId, + origin: "app_owned", + driver, + providerInstanceId, + childThreadId: null, + prompt: "Do the work", + result: "done", + }, + }); + const applyTask = ( + name: string, + policy: Pick, + ) => + store.apply({ + id: EventId.make(`event:delegated-wake-hold:${name}`), + type: "subagent.updated", + threadId, + runId: parentRun.id, + nodeId: taskId, + driver, + providerInstanceId, + occurredAt: now, + payload: { + id: taskId, + threadId, + runId: parentRun.id, + parentNodeId: parentRun.rootNodeId, + origin: "app_owned", + createdBy: "agent", + driver, + providerInstanceId, + providerThreadId: null, + childThreadId: null, + nativeTaskRef: null, + prompt: "Do the work", + title: "Delegated task", + model: null, + ...policy, + status: "completed", + result: "done", + startedAt: now, + completedAt: now, + updatedAt: now, + }, + }); + yield* applyTask("task", { + completionWake: "always", + completionDelivery: { state: "claimed", observedByRunId: null }, + }); + const heldTaskIds = Effect.gen(function* () { + yield* assertShellIndexCoversShells(); + const taskIds = ( + shell: + | { + readonly pendingBackgroundTasks?: + | ReadonlyArray<{ readonly taskId: string }> + | undefined; + } + | null + | undefined, + ) => (shell?.pendingBackgroundTasks ?? []).map((task) => task.taskId); + const producers = [ + taskIds(yield* store.getThreadShell(threadId)), + taskIds((yield* store.getShellSnapshot()).threads.find((row) => row.id === threadId)), + taskIds( + ProjectionStore.threadShellFromProjection(yield* store.getThreadProjection(threadId)), + ), + taskIds((yield* store.getSettlementCandidates(threadId))[0]), + taskIds((yield* store.getAgentAwarenessSnapshot(threadId)).threads[0]), + ]; + for (const held of producers) assert.deepEqual(held, producers[0]); + return producers[0]; + }); + + assert.deepEqual(yield* heldTaskIds, [itemId]); + + // No provider thread means no wake is ever offered. + yield* applyRun("missing-provider-thread", { + ...parentRun, + providerThreadId: ProviderThreadId.make("provider-thread:delegated-wake-hold:missing"), + }); + assert.deepEqual(yield* heldTaskIds, []); + yield* applyRun("parent-restored", parentRun); + + const wakeRun = { + ...parentRun, + id: RunId.make("run:delegated-wake-hold:wake"), + ordinal: 3, + userMessageId: wakeMessageId, + rootNodeId: NodeId.make("node:delegated-wake-hold:wake"), + delegatedCompletion: { disposition: "open" as const, nextGeneration: 1, delivery: null }, + }; + yield* applyRun("wake-run", wakeRun); + yield* store.apply({ + id: EventId.make("event:delegated-wake-hold:wake-message"), + type: "message.updated", + threadId, + runId: wakeRun.id, + occurredAt: now, + payload: { + id: wakeMessageId, + threadId, + runId: wakeRun.id, + nodeId: null, + role: "user", + text: "Delegated task finished", + attachments: [], + streaming: false, + createdBy: "agent", + creationSource: "server", + delegatedCompletion: { parentRunId: parentRun.id, generation: 1, taskIds: [taskId] }, + createdAt: now, + updatedAt: now, + }, + }); + assert.deepEqual(yield* heldTaskIds, []); + + // A wake steered into a run that ended before reading it is offered again. + yield* applyRun("wake-run-steered", { + ...wakeRun, + userMessageId: MessageId.make("message:delegated-wake-hold:steered-into"), + }); + assert.deepEqual(yield* heldTaskIds, [itemId]); + yield* applyRun("wake-run-read", { ...wakeRun, status: "completed" }); + assert.deepEqual(yield* heldTaskIds, []); + + // A sibling that finished behind that wake waits for the successor its end reserves. + yield* applyTask("task-pending", { + completionWake: "always", + completionDelivery: { state: "pending", observedByRunId: null }, + }); + assert.deepEqual(yield* heldTaskIds, [itemId]); + + // A live parent's blocking wait takes a settled_only result instead. + yield* applyTask("task-settled-only", { + completionWake: "settled_only", + completionDelivery: { state: "pending", observedByRunId: null }, + }); + assert.deepEqual(yield* heldTaskIds, [itemId]); + yield* applyRun("parent-live", { ...parentRun, status: "running", completedAt: null }); + assert.deepEqual(yield* heldTaskIds, []); + + // With no wake outstanding, nothing reserves one. + yield* applyRun("parent-settled", { + ...parentRun, + delegatedCompletion: { disposition: "open", nextGeneration: 2, delivery: null }, + }); + assert.deepEqual(yield* heldTaskIds, []); +}); + const commandReadsSkipSettledLeafNodes = Effect.gen(function* () { const store = yield* ProjectionStore.ProjectionStoreV2; const now = yield* DateTime.now; @@ -2711,6 +2945,10 @@ it.layer(layerTest)("ProjectionStoreV2", (it) => { "keeps restart-cancelled work through a stale run.updated", () => restartCancelledWorkSurvivesStaleRunUpdate, ); + it.effect( + "every shell holds a finished delegated task until its completion wake starts", + () => delegatedCompletionWakeHoldsEveryShell, + ); it.effect("records restart-cancelled work without regressing a run that completed since", () => Effect.gen(function* () { const store = yield* ProjectionStore.ProjectionStoreV2; diff --git a/apps/server/src/orchestration-v2/ProjectionStore.ts b/apps/server/src/orchestration-v2/ProjectionStore.ts index 0c4976dc90598..5f05f26b69e40 100644 --- a/apps/server/src/orchestration-v2/ProjectionStore.ts +++ b/apps/server/src/orchestration-v2/ProjectionStore.ts @@ -5,6 +5,7 @@ import { projectTurnItemForWire } from "./WireProjection.ts"; import * as Stream from "effect/Stream"; import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient"; import { makeThreadFind } from "./ThreadFind.ts"; +import { delegatedTaskItemsAwaitingWake } from "./NotificationMailbox.ts"; import type { OrchestrationV2ThreadLinkMetadata, OrchestrationV2SearchThreadInput, @@ -1010,6 +1011,15 @@ type ShellPendingTurnItemRow = PayloadRow & { readonly owning_provider_instance_id: string | null; }; +type ShellAwaitingWakeRow = ShellPendingTurnItemRow & { + readonly provider_thread_id: string; + readonly task_payload_json: string; + readonly run_payload_json: string; + readonly message_payload_json: string | null; + readonly message_run_payload_json: string | null; + readonly message_turn_payload_json: string | null; +}; + export type AgentAwarenessThread = ProjectThreadAwarenessV2Input["thread"] & Pick< OrchestrationV2ThreadShell, @@ -1616,6 +1626,7 @@ export function threadShellFromProjection( activeProviderThreadId: projection.thread.activeProviderThreadId, runs, pullRequests: projection.thread.pullRequests, + awaitingWakeTurnItems: delegatedTaskItemsAwaitingWake(projection), }); return { createdBy: projection.thread.createdBy, @@ -6611,6 +6622,39 @@ export const layer: Layer.Layer = } `; + // Finished tasks owed a wake are rare and indexed; the run, message and turn joins only run for them. + const selectShellAwaitingWakeRows = (threadIds: ReadonlyArray) => + sql` + SELECT s.thread_id, p.provider_thread_id, item.payload_json, + s.payload_json AS task_payload_json, + owner.provider_instance_id AS owning_provider_instance_id, + r.payload_json AS run_payload_json, + m.payload_json AS message_payload_json, + mr.payload_json AS message_run_payload_json, + CASE WHEN mr.run_id IS NULL THEN NULL ELSE ( + SELECT turn.payload_json FROM orchestration_v2_projection_provider_turns turn + WHERE turn.thread_id = s.thread_id + AND turn.run_attempt_id = json_extract(mr.payload_json, '$.activeAttemptId') + ORDER BY turn.ordinal DESC + LIMIT 1 + ) END AS message_turn_payload_json + FROM orchestration_v2_projection_subagents s + INNER JOIN orchestration_v2_projection_runs r ON r.run_id = s.run_id + INNER JOIN orchestration_v2_projection_provider_threads p + ON p.provider_thread_id = r.provider_thread_id + -- CROSS JOIN pins the claimed-task index first; unpinned, SQLite walks every turn item. + CROSS JOIN orchestration_v2_projection_turn_items item + ON item.node_id = s.subagent_id AND item.thread_id = s.thread_id AND item.type = 'subagent' + LEFT JOIN orchestration_v2_projection_provider_threads owner + ON owner.provider_thread_id = item.provider_thread_id + LEFT JOIN orchestration_v2_projection_messages m + ON m.message_id = json_extract(r.payload_json, '$.delegatedCompletion.delivery.messageId') + LEFT JOIN orchestration_v2_projection_runs mr ON mr.run_id = m.run_id + WHERE s.thread_id IN (SELECT value FROM json_each(${encodeIdList(threadIds)})) + AND s.completion_delivery_state IN ('claimed', 'pending') + AND json_type(r.payload_json, '$.delegatedCompletion.delivery') = 'object' + `; + const runMapsByThreadId = (input: { readonly runRows: ReadonlyArray; readonly itemCountRows: ReadonlyArray; @@ -6650,6 +6694,7 @@ export const layer: Layer.Layer = )(function* (input: { readonly providerThreadRows: ReadonlyArray; readonly pendingTurnItemRows: ReadonlyArray; + readonly awaitingWakeRows: ReadonlyArray; }) { const providerThreadsByThreadId = new Map< ThreadId, @@ -6668,22 +6713,64 @@ export const layer: Layer.Layer = } const pendingTurnItemsByThreadId = new Map>(); + const decodeShellTurnItem = (row: ShellPendingTurnItemRow) => + decodeTurnItemPayload(row.payload_json).pipe( + Effect.map((turnItem) => ({ + ...turnItem, + ...(row.owning_provider_instance_id === null + ? {} + : { + owningProviderInstanceId: ProviderInstanceId.make( + row.owning_provider_instance_id, + ), + }), + })), + ); for (const row of input.pendingTurnItemRows) { - const turnItem = { - ...(yield* decodeTurnItemPayload(row.payload_json)), - ...(row.owning_provider_instance_id === null - ? {} - : { - owningProviderInstanceId: ProviderInstanceId.make(row.owning_provider_instance_id), - }), - }; + const turnItem = yield* decodeShellTurnItem(row); const threadId = ThreadId.make(row.thread_id); const existing = pendingTurnItemsByThreadId.get(threadId) ?? []; existing.push(turnItem); pendingTurnItemsByThreadId.set(threadId, existing); } - return { providerThreadsByThreadId, pendingTurnItemsByThreadId }; + const awaitingWakeTurnItemsByThreadId = new Map>(); + for (const row of input.awaitingWakeRows) { + const message = + row.message_payload_json === null + ? null + : yield* decodeMessagePayload(row.message_payload_json); + const messageRun = + row.message_run_payload_json === null + ? null + : yield* decodeRunPayload(row.message_run_payload_json); + const messageTurn = + row.message_turn_payload_json === null + ? null + : yield* decodeProviderTurnPayload(row.message_turn_payload_json); + const items = delegatedTaskItemsAwaitingWake({ + runs: [ + yield* decodeRunPayload(row.run_payload_json), + ...(messageRun === null ? [] : [messageRun]), + ], + messages: message === null ? [] : [message], + providerTurns: messageTurn === null ? [] : [messageTurn], + providerThreads: [{ id: row.provider_thread_id }], + subagents: [yield* decodeSubagentPayload(row.task_payload_json)], + turnItems: [yield* decodeShellTurnItem(row)], + }); + if (items.length === 0) continue; + const threadId = ThreadId.make(row.thread_id); + const existing = awaitingWakeTurnItemsByThreadId.get(threadId) ?? []; + existing.push(...items); + awaitingWakeTurnItemsByThreadId.set(threadId, existing); + } + + return { + providerThreadsByThreadId, + pendingTurnItemsByThreadId, + awaitingWakeTurnItemsByThreadId, + }; }); const getSettlementCandidates: ProjectionStoreV2Shape["getSettlementCandidates"] = (threadId) => @@ -6756,11 +6843,15 @@ export const layer: Layer.Layer = `; if (rows.length === 0) return []; const threadIds = rows.map((row) => ThreadId.make(row.thread_id)); - const { providerThreadsByThreadId, pendingTurnItemsByThreadId } = - yield* pendingBackgroundDataByThreadId({ - providerThreadRows: yield* selectShellProviderThreadRows(threadIds), - pendingTurnItemRows: yield* selectShellPendingTurnItemRows(threadIds), - }); + const { + providerThreadsByThreadId, + pendingTurnItemsByThreadId, + awaitingWakeTurnItemsByThreadId, + } = yield* pendingBackgroundDataByThreadId({ + providerThreadRows: yield* selectShellProviderThreadRows(threadIds), + pendingTurnItemRows: yield* selectShellPendingTurnItemRows(threadIds), + awaitingWakeRows: yield* selectShellAwaitingWakeRows(threadIds), + }); return yield* Effect.forEach(rows, (row) => Effect.gen(function* () { const thread = yield* decodeThreadPayload(row.payload_json); @@ -6810,6 +6901,7 @@ export const layer: Layer.Layer = activeProviderThreadId: thread.activeProviderThreadId, hasActiveRun: false, pullRequests: thread.pullRequests, + awaitingWakeTurnItems: awaitingWakeTurnItemsByThreadId.get(thread.id) ?? [], }), } satisfies ProjectionSettlementCandidate; }), @@ -6906,6 +6998,10 @@ export const layer: Layer.Layer = ThreadId, ReadonlyArray >; + readonly awaitingWakeTurnItemsByThreadId: ReadonlyMap< + ThreadId, + ReadonlyArray + >; }) { const { row, @@ -6914,6 +7010,7 @@ export const layer: Layer.Layer = shellItemCountsByThreadId, providerThreadsByThreadId, pendingTurnItemsByThreadId, + awaitingWakeTurnItemsByThreadId, } = input; const thread = yield* decodeThreadPayload(row.payload_json); const latestUserAuthoredMessageAt = @@ -7022,6 +7119,7 @@ export const layer: Layer.Layer = activeProviderThreadId: thread.activeProviderThreadId, hasActiveRun: row.active_run_id !== null, pullRequests: thread.pullRequests, + awaitingWakeTurnItems: awaitingWakeTurnItemsByThreadId.get(thread.id) ?? [], }), ]; return { @@ -7120,6 +7218,7 @@ export const layer: Layer.Layer = sequenceRows, providerThreadRows, pendingTurnItemRows, + awaitingWakeRows, ] = yield* Effect.all([ readForThreadIds(selectShellRunRows, forkSourceIds), readForThreadIds(selectShellRunItemCounts, forkSourceIds), @@ -7141,6 +7240,7 @@ export const layer: Layer.Layer = ), ), ), + readForThreadIds(selectShellAwaitingWakeRows), ]); return { targetThreadIds, @@ -7151,6 +7251,7 @@ export const layer: Layer.Layer = sequenceRows, providerThreadRows, pendingTurnItemRows, + awaitingWakeRows, }; }); @@ -7163,22 +7264,25 @@ export const layer: Layer.Layer = sequenceRows, providerThreadRows, pendingTurnItemRows, + awaitingWakeRows, }: Effect.Success>) { const { runOrdinalsByThreadId, itemCountsByThreadId } = runMapsByThreadId({ runRows, itemCountRows, }); const shellItemCountsByThreadId = shellItemCountMap(shellItemCountRows); - const { providerThreadsByThreadId, pendingTurnItemsByThreadId } = - yield* pendingBackgroundDataByThreadId({ providerThreadRows, pendingTurnItemRows }); + const background = yield* pendingBackgroundDataByThreadId({ + providerThreadRows, + pendingTurnItemRows, + awaitingWakeRows, + }); const states = yield* Effect.forEach(threadRows, (row) => shellThreadStateFromRow({ row, runOrdinalsByThreadId, itemCountsByThreadId, shellItemCountsByThreadId, - providerThreadsByThreadId, - pendingTurnItemsByThreadId, + ...background, }), ); const statesByThreadId = new Map(states.map((state) => [state.thread.id, state])); @@ -7246,36 +7350,52 @@ export const layer: Layer.Layer = const getAgentAwarenessSnapshot = Effect.fn("ProjectionStore.getAgentAwarenessSnapshot")( function* (threadId?: ThreadId) { - const { threadRows, providerThreadRows, pendingTurnItemRows, sequenceRows } = - yield* sql.withTransaction( - Effect.gen(function* () { - const threadRows = yield* selectShellThreadRows( - threadId, - threadId === undefined ? { location: "active", listedOnly: true } : {}, - ); - const ids = threadRows.map((row) => ThreadId.make(row.thread_id)); - const providerThreadRows = - ids.length === 0 ? [] : yield* selectShellProviderThreadRows(ids); - const pendingTurnItemRows = - ids.length === 0 - ? [] - : yield* selectShellPendingTurnItemRows( - ids, - threadRows.flatMap((row) => - row.active_run_id !== null && row.active_run_status === "running" - ? [row.active_run_id] - : [], - ), - ); - const sequenceRows = yield* sql<{ readonly snapshot_sequence: number | null }>` + const { + threadRows, + providerThreadRows, + pendingTurnItemRows, + awaitingWakeRows, + sequenceRows, + } = yield* sql.withTransaction( + Effect.gen(function* () { + const threadRows = yield* selectShellThreadRows( + threadId, + threadId === undefined ? { location: "active", listedOnly: true } : {}, + ); + const ids = threadRows.map((row) => ThreadId.make(row.thread_id)); + const providerThreadRows = + ids.length === 0 ? [] : yield* selectShellProviderThreadRows(ids); + const pendingTurnItemRows = + ids.length === 0 + ? [] + : yield* selectShellPendingTurnItemRows( + ids, + threadRows.flatMap((row) => + row.active_run_id !== null && row.active_run_status === "running" + ? [row.active_run_id] + : [], + ), + ); + const awaitingWakeRows = + ids.length === 0 ? [] : yield* selectShellAwaitingWakeRows(ids); + const sequenceRows = yield* sql<{ readonly snapshot_sequence: number | null }>` SELECT MAX(sequence) AS snapshot_sequence FROM orchestration_events WHERE application_event_version = 2 AND aggregate_kind = 'thread' `; - return { threadRows, providerThreadRows, pendingTurnItemRows, sequenceRows }; - }), - ); - const { providerThreadsByThreadId, pendingTurnItemsByThreadId } = - yield* pendingBackgroundDataByThreadId({ providerThreadRows, pendingTurnItemRows }); + return { + threadRows, + providerThreadRows, + pendingTurnItemRows, + awaitingWakeRows, + sequenceRows, + }; + }), + ); + const background = yield* pendingBackgroundDataByThreadId({ + providerThreadRows, + pendingTurnItemRows, + awaitingWakeRows, + }); const threads = yield* Effect.forEach(threadRows, (row) => shellThreadStateFromRow({ row, @@ -7283,8 +7403,7 @@ export const layer: Layer.Layer = runOrdinalsByThreadId: new Map(), itemCountsByThreadId: new Map(), shellItemCountsByThreadId: new Map(), - providerThreadsByThreadId, - pendingTurnItemsByThreadId, + ...background, }).pipe( Effect.map((state): AgentAwarenessThread => ({ id: state.thread.id, @@ -7398,6 +7517,13 @@ export const layer: Layer.Layer = SELECT 1 FROM json_each(t.payload_json, '$.pullRequests') link WHERE json_type(link.value, '$.watch') IS NOT NULL ) + OR EXISTS ( + SELECT 1 FROM orchestration_v2_projection_subagents task + INNER JOIN orchestration_v2_projection_runs task_run ON task_run.run_id = task.run_id + WHERE task.thread_id = t.thread_id + AND task.completion_delivery_state IN ('claimed', 'pending') + AND json_type(task_run.payload_json, '$.delegatedCompletion.delivery') = 'object' + ) ) AS may_be_live FROM orchestration_v2_projection_threads t WHERE t.deleted_at IS NULL @@ -7553,6 +7679,7 @@ export const layer: Layer.Layer = shellItemCountRows, providerThreadRows, pendingTurnItemRows, + awaitingWakeRows, ] = yield* Effect.all([ forkSourceIds.length === 0 ? Effect.succeed([] as ReadonlyArray) @@ -7570,14 +7697,18 @@ export const layer: Layer.Layer = : [], ), ), + selectShellAwaitingWakeRows(threadIds), ]); const { runOrdinalsByThreadId, itemCountsByThreadId } = runMapsByThreadId({ runRows, itemCountRows, }); const shellItemCountsByThreadId = shellItemCountMap(shellItemCountRows); - const { providerThreadsByThreadId, pendingTurnItemsByThreadId } = - yield* pendingBackgroundDataByThreadId({ providerThreadRows, pendingTurnItemRows }); + const background = yield* pendingBackgroundDataByThreadId({ + providerThreadRows, + pendingTurnItemRows, + awaitingWakeRows, + }); const states = yield* Effect.forEach([...rowsByThreadId.values()], (row) => shellThreadStateFromRow({ @@ -7585,8 +7716,7 @@ export const layer: Layer.Layer = runOrdinalsByThreadId, itemCountsByThreadId, shellItemCountsByThreadId, - providerThreadsByThreadId, - pendingTurnItemsByThreadId, + ...background, }), ); const statesByThreadId = new Map(states.map((state) => [state.thread.id, state])); diff --git a/apps/server/src/orchestration-v2/ProviderContinuationService.ts b/apps/server/src/orchestration-v2/ProviderContinuationService.ts index 1fe30ee0645c8..3f9b69da863cc 100644 --- a/apps/server/src/orchestration-v2/ProviderContinuationService.ts +++ b/apps/server/src/orchestration-v2/ProviderContinuationService.ts @@ -6,7 +6,7 @@ import * as Ref from "effect/Ref"; import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator"; import * as ProviderContinuationRequests from "@t3tools/provider-core/server/ProviderContinuationRequests"; import * as ThreadManagementService from "./ThreadManagementService.ts"; -import { isUndeliveredMailboxSteer } from "./NotificationMailbox.ts"; +import { pendingDelegatedCompletionWake } from "./NotificationMailbox.ts"; const CONTINUATION_MESSAGE_TEXT = "Background task completed."; @@ -24,21 +24,12 @@ function currentDelegatedCompletionDelivery( >, ) { const sourceRun = projection.runs.find((candidate) => candidate.id === completion.parentRunId); - const delivery = sourceRun?.delegatedCompletion?.delivery; - const alreadyDispatched = projection.messages.some( - (message) => message.id === completion.messageId, - ); - if ( - sourceRun?.delegatedCompletion?.disposition !== "open" || - delivery === null || - delivery === undefined || - delivery.generation !== completion.generation || - delivery.messageId !== completion.messageId || - (alreadyDispatched && !isUndeliveredMailboxSteer(projection, completion.messageId)) - ) { - return undefined; - } - return delivery; + const delivery = + sourceRun === undefined ? undefined : pendingDelegatedCompletionWake(projection, sourceRun); + return delivery?.generation === completion.generation && + delivery.messageId === completion.messageId + ? delivery + : undefined; } function delegatedCompletionRetryKey( diff --git a/packages/shared/src/orchestrationV2PendingBackgroundWork.ts b/packages/shared/src/orchestrationV2PendingBackgroundWork.ts index dfdf9a24ace1e..3f34135b55ceb 100644 --- a/packages/shared/src/orchestrationV2PendingBackgroundWork.ts +++ b/packages/shared/src/orchestrationV2PendingBackgroundWork.ts @@ -245,6 +245,8 @@ export function derivePendingBackgroundWork(input: { /** Excludes rolled-back runs; optional only when the caller already filtered them, as SQL shell reads do. */ readonly runs?: ReadonlyArray; readonly pullRequests?: ReadonlyArray | undefined; + /** Finished delegated task items whose completion wake has not started; they hold like running ones. */ + readonly awaitingWakeTurnItems?: ReadonlyArray; }): ReadonlyArray { const hasActiveRun = input.hasActiveRun ?? @@ -286,7 +288,13 @@ export function derivePendingBackgroundWork(input: { } } - for (const item of pendingBackgroundTurnItems(input)) { + for (const item of [ + ...pendingBackgroundTurnItems(input), + // A finished task has no provider process, so it must not keep the caller's MCP session live. + ...(input.awaitingWakeTurnItems ?? []).map( + ({ owningProviderInstanceId: _owner, ...item }) => item, + ), + ]) { const taskId = nativeTaskIdFromTurnItem(item); if (byTaskId.has(taskId)) { continue;