diff --git a/src/shared/prompt-async-gate.ts b/src/shared/prompt-async-gate.ts index a54f98713..234f1e65d 100644 --- a/src/shared/prompt-async-gate.ts +++ b/src/shared/prompt-async-gate.ts @@ -1,179 +1,50 @@ import { log } from "./logger" +import { DEFAULT_SESSION_IDLE_SETTLE_MS } from "./session-idle-settle" import { - isSyntheticOrInternalUserMessage, - type InternalInitiatorMessageLike, - type InternalInitiatorTextPartLike, -} from "./internal-initiator-marker" + clearPromptQueueStateForTesting, + enqueueInternalPrompt, + getQueuedPromptBlocker, + isPromptQueueDraining, + nextPromptQueueID, + releaseInFlightPromptMatchingDedupe, + schedulePromptQueueDrain, +} from "./prompt-async-gate/queue" import { - DEFAULT_SESSION_IDLE_SETTLE_MS, - isSessionActive, - settleAfterSessionIdle, -} from "./session-idle-settle" + clearPromptReservationsForTesting, + deletePromptReservation, + getActiveReservation, + getPromptReservation, + reservationSourceMatches, +} from "./prompt-async-gate/reservations" +import { dispatchAfterSessionIdle } from "./prompt-async-gate/session-idle-dispatch" +import { + DEFAULT_PROMPT_ASYNC_POST_DISPATCH_HOLD_MS, + DEFAULT_PROMPT_DISPATCH_TIMEOUT_MS, + DEFAULT_PROMPT_QUEUE_RETRY_MS, + resetPromptGateTimingForTesting, +} from "./prompt-async-gate/timing" +import type { + InternalPromptDispatchArgs, + InternalPromptDispatchResult, + PromptAsyncInput, + PromptAsyncReservationReleaseOptions, +} from "./prompt-async-gate/types" -export const DEFAULT_PROMPT_ASYNC_POST_DISPATCH_HOLD_MS = 2_000 -export const DEFAULT_PROMPT_DISPATCH_TIMEOUT_MS = 30_000 -export const DEFAULT_PROMPT_GATE_MESSAGES_FETCH_TIMEOUT_MS = 5_000 -export const DEFAULT_PROMPT_QUEUE_RETRY_MS = 250 +export { + DEFAULT_PROMPT_ASYNC_POST_DISPATCH_HOLD_MS, + DEFAULT_PROMPT_DISPATCH_TIMEOUT_MS, + DEFAULT_PROMPT_GATE_MESSAGES_FETCH_TIMEOUT_MS, + DEFAULT_PROMPT_QUEUE_RETRY_MS, + _setPromptGateMessagesFetchTimeoutMsForTesting, +} from "./prompt-async-gate/timing" -type PromptAsyncInput = { - path?: { id?: string } - body?: unknown - query?: unknown - signal?: unknown - [key: string]: unknown -} - -type PromptMessagesQuery = { - directory: string - limit?: number -} - -type PromptAsyncClient = { - session?: { - status?: () => Promise - messages?: (input: { path: { id: string }; query: PromptMessagesQuery }) => Promise - promptAsync?: (input: TInput) => Promise - } -} - -type PromptClient = { - session?: { - status?: () => Promise - messages?: (input: { path: { id: string }; query: PromptMessagesQuery }) => Promise - prompt?: (input: TInput) => Promise - } -} - -export type InternalPromptDispatchMode = "async" | "sync" -export type InternalPromptQueueBehavior = "enqueue" | "defer" - -type InternalPromptDispatchCommonArgs = { - sessionID: string - input: TInput - source: string - dedupeKey?: string - queueBehavior?: InternalPromptQueueBehavior - queue?: boolean - queueRetryMs?: number - settleMs?: number - postDispatchHoldMs?: number - dispatchTimeoutMs?: number - checkStatus?: boolean - checkToolState?: boolean -} - -export type InternalPromptDispatchArgs = InternalPromptDispatchCommonArgs & ( - | { mode: "async"; client: PromptAsyncClient } - | { mode: "sync"; client: PromptClient } -) - -type PromptAsyncReservation = { - source: string - dedupeKey: string - reservedAt: number - token: symbol - expiresAt?: number -} - -declare function setTimeout(callback: () => void, delay?: number): unknown -declare function clearTimeout(timeout: unknown): void - -let promptGateMessagesFetchTimeoutMsForTesting: number | undefined - -export type InternalPromptDispatchResult = - | { status: "dispatched"; response: unknown } - | { status: "queued"; queuedBy: string; position: number } - | { status: "active" } - | { status: "reserved"; reservedBy: string } - | { status: "unavailable" } - | { status: "failed"; error: unknown; dispatchAttempted: boolean } - -export type PromptAsyncGateResult = InternalPromptDispatchResult - -type PromptAsyncReservationReleaseOptions = { - reservedBy?: string | readonly string[] - reservedByPrefix?: string | readonly string[] -} - -const promptAsyncReservations = new Map() -const promptQueues = new Map() -const promptQueueDraining = new Set() -const promptQueueInFlight = new Map() -const promptQueueTimers = new Map() -let promptQueueSequence = 0 - -type PromptDispatchClient = { - session?: { - status?: () => Promise - messages?: (input: { path: { id: string }; query: PromptMessagesQuery }) => Promise - } -} - -type QueuedInternalPrompt = { - id: number - sessionID: string - sessionName: "promptAsync" | "prompt" - client: PromptDispatchClient - input: unknown - source: string - dedupeKey: string - settleMs: number - postDispatchHoldMs: number - dispatchTimeoutMs: number - queueRetryMs: number - checkStatus: boolean - checkToolState: boolean - dispatch: (input: unknown) => Promise -} - -export function _setPromptGateMessagesFetchTimeoutMsForTesting(value: number | undefined): void { - promptGateMessagesFetchTimeoutMsForTesting = value -} - -function getPromptGateMessagesFetchTimeoutMs(): number { - return promptGateMessagesFetchTimeoutMsForTesting ?? DEFAULT_PROMPT_GATE_MESSAGES_FETCH_TIMEOUT_MS -} - -function pruneExpiredReservations(now = Date.now()): void { - const expiredSessionIDs: string[] = [] - for (const [sessionID, reservation] of promptAsyncReservations) { - if (typeof reservation.expiresAt === "number" && reservation.expiresAt <= now) { - promptAsyncReservations.delete(sessionID) - expiredSessionIDs.push(sessionID) - log("[prompt-async-gate] expired reservation released", { - sessionID, - source: reservation.source, - }) - } - } - for (const sessionID of expiredSessionIDs) { - schedulePromptQueueDrain(sessionID, 0) - } -} - -function getActiveReservation(sessionID: string): PromptAsyncReservation | undefined { - pruneExpiredReservations() - return promptAsyncReservations.get(sessionID) -} - -function getPromptQueue(sessionID: string): QueuedInternalPrompt[] { - const existing = promptQueues.get(sessionID) - if (existing) { - return existing - } - - const queue: QueuedInternalPrompt[] = [] - promptQueues.set(sessionID, queue) - return queue -} - -function setPromptQueue(sessionID: string, queue: QueuedInternalPrompt[]): void { - if (queue.length === 0) { - promptQueues.delete(sessionID) - return - } - promptQueues.set(sessionID, queue) -} +export type { + InternalPromptDispatchArgs, + InternalPromptDispatchMode, + InternalPromptDispatchResult, + InternalPromptQueueBehavior, + PromptAsyncGateResult, +} from "./prompt-async-gate/types" function stringifyPromptInputForDedupe(input: unknown): string { try { @@ -197,539 +68,6 @@ function createDefaultDedupeKey(source: string, input: unknown): string { return `${source}:${fingerprint.length}:${fingerprint.slice(0, 8192)}` } -function queuedResult(entry: QueuedInternalPrompt, position: number, queuedBy = entry.source): InternalPromptDispatchResult { - return { - status: "queued", - queuedBy, - position, - } -} - -function clearPromptQueueTimer(sessionID: string): void { - const timer = promptQueueTimers.get(sessionID) - if (timer !== undefined) { - clearTimeout(timer) - promptQueueTimers.delete(sessionID) - } -} - -function schedulePromptQueueDrain(sessionID: string, delayMs: number): void { - const queue = promptQueues.get(sessionID) - if (!queue || queue.length === 0) { - clearPromptQueueTimer(sessionID) - return - } - - clearPromptQueueTimer(sessionID) - const timer = setTimeout(() => { - promptQueueTimers.delete(sessionID) - void drainPromptQueue(sessionID).catch((error: unknown) => { - log("[prompt-async-gate] queued prompt drain failed", { - sessionID, - error: String(error), - }) - }) - }, Math.max(0, delayMs)) - promptQueueTimers.set(sessionID, timer) -} - -function removePromptQueueEntry(sessionID: string, entry: QueuedInternalPrompt): void { - const queue = promptQueues.get(sessionID) - if (!queue) { - return - } - const nextQueue = queue.filter((queued) => queued.id !== entry.id) - setPromptQueue(sessionID, nextQueue) -} - -function getQueuedPromptBlocker(sessionID: string): string | undefined { - const inFlight = promptQueueInFlight.get(sessionID) - if (inFlight) { - return inFlight.source - } - - const queue = promptQueues.get(sessionID) - return queue?.[0]?.source -} - -function reservationSourceMatches( - reservationSource: string, - expectedSource: string | readonly string[], - expectedPrefix?: string | readonly string[], -): boolean { - if (typeof expectedSource === "string") { - if (reservationSource === expectedSource) { - return true - } - } else if (expectedSource.includes(reservationSource)) { - return true - } - - if (expectedPrefix === undefined) { - return false - } - - const prefixes = typeof expectedPrefix === "string" ? [expectedPrefix] : expectedPrefix - return prefixes - .filter((prefix) => prefix.length > 0 && prefix.endsWith(":")) - .some((prefix) => reservationSource.startsWith(prefix)) -} - -async function withDispatchTimeout( - operation: Promise, - dispatchTimeoutMs: number, - operationName: string, -): Promise { - if (dispatchTimeoutMs <= 0) { - return operation - } - - let timeoutID: unknown - const timeoutPromise = new Promise((_, reject) => { - timeoutID = setTimeout(() => { - reject(new Error(`${operationName} timed out after ${dispatchTimeoutMs}ms`)) - }, dispatchTimeoutMs) - }) - - try { - return await Promise.race([operation, timeoutPromise]) - } finally { - if (timeoutID !== undefined) { - clearTimeout(timeoutID) - } - } -} - -function isRecord(value: unknown): value is Record { - return typeof value === "object" && value !== null -} - -function getPromptQuery(input: unknown): PromptMessagesQuery { - if (!isRecord(input)) { - return { directory: "" } - } - const query = input.query - if (!isRecord(query)) { - return { directory: "" } - } - - const promptQuery: PromptMessagesQuery = { directory: "" } - if (typeof query.directory === "string") { - promptQuery.directory = query.directory - } - if (typeof query.limit === "number") { - promptQuery.limit = query.limit - } - return promptQuery -} - -function getMessagesData(response: unknown): unknown[] { - if (isRecord(response) && Array.isArray(response.data)) { - return response.data - } - return Array.isArray(response) ? response : [] -} - -function messageRole(message: unknown): string | undefined { - if (!isRecord(message)) { - return undefined - } - const info = message.info - if (isRecord(info) && typeof info.role === "string") { - return info.role - } - return typeof message.role === "string" ? message.role : undefined -} - -function messageFinish(message: unknown): string | true | undefined { - if (!isRecord(message)) { - return undefined - } - const info = message.info - if (isRecord(info)) { - if (info.finish === true) { - return true - } - if (typeof info.finish === "string" && info.finish.length > 0) { - return info.finish - } - } - if (message.finish === true) { - return true - } - return typeof message.finish === "string" && message.finish.length > 0 ? message.finish : undefined -} - -function messageCompleted(message: unknown): boolean { - if (!isRecord(message)) { - return false - } - const info = message.info - const time = isRecord(info) && isRecord(info.time) ? info.time : undefined - const completed = time?.completed - if (typeof completed === "number" && Number.isFinite(completed)) { - return true - } - return typeof completed === "string" && completed.length > 0 -} - -function toInternalInitiatorTextPartLike(part: unknown): InternalInitiatorTextPartLike { - const result: InternalInitiatorTextPartLike = {} - if (!isRecord(part)) { - return result - } - - if (typeof part.type === "string") { - result.type = part.type - } - if (typeof part.text === "string") { - result.text = part.text - } - if (typeof part.synthetic === "boolean") { - result.synthetic = part.synthetic - } - return result -} - -function toInternalInitiatorMessageLike(message: unknown): InternalInitiatorMessageLike | undefined { - if (!isRecord(message)) { - return undefined - } - - const result: InternalInitiatorMessageLike = {} - const info = message.info - if (isRecord(info) && typeof info.role === "string") { - result.info = { role: info.role } - } - if (typeof message.role === "string") { - result.role = message.role - } - if (Array.isArray(message.parts)) { - result.parts = message.parts.map(toInternalInitiatorTextPartLike) - } - return result -} - -function messageIsSyntheticOrInternalUser(message: unknown): boolean { - const initiatorMessage = toInternalInitiatorMessageLike(message) - return initiatorMessage !== undefined && isSyntheticOrInternalUserMessage(initiatorMessage) -} - -function partIsWaitingOnTool(part: unknown): boolean { - if (!isRecord(part)) { - return false - } - if ( - part.type !== "tool" - && part.type !== "tool_use" - && part.type !== "tool-call" - && part.type !== "tool-invocation" - ) { - return false - } - - const state = part.state - if (!isRecord(state)) { - return false - } - return state.status === "pending" || state.status === "running" -} - -function latestAssistantTurnBlocksInternalPrompt(messages: unknown[]): boolean { - for (let index = messages.length - 1; index >= 0; index--) { - const message = messages[index] - const role = messageRole(message) - if (role === "assistant") { - if (messageCompleted(message)) { - return false - } - const finish = messageFinish(message) - if (finish === true) { - return false - } - if (finish === undefined || finish === "unknown") { - return true - } - if (!isRecord(message) || !Array.isArray(message.parts)) { - return finish === "tool-calls" - } - return finish === "tool-calls" || message.parts.some(partIsWaitingOnTool) - } - if (role === "user") { - if (messageIsSyntheticOrInternalUser(message)) { - continue - } - return false - } - } - return false -} - -async function sessionLatestAssistantBlocksInternalPrompt(args: { - client: { session?: { messages?: (input: { path: { id: string }; query: PromptMessagesQuery }) => Promise } } - sessionID: string - input: TInput - sessionName: "promptAsync" | "prompt" - source: string - timeoutMs: number -}): Promise { - const session = args.client.session - if (typeof session?.messages !== "function") { - return false - } - const messages = session.messages.bind(session) - - try { - const response = await withDispatchTimeout( - messages({ - path: { id: args.sessionID }, - query: getPromptQuery(args.input), - }), - args.timeoutMs, - `[prompt-async-gate] ${args.sessionName} session.messages`, - ) - return latestAssistantTurnBlocksInternalPrompt(getMessagesData(response)) - } catch (error) { - log("[prompt-async-gate] latest assistant prompt-block check failed", { - sessionID: args.sessionID, - source: args.source, - error: String(error), - }) - return false - } -} - -async function dispatchAfterSessionIdle(args: { - sessionName: "promptAsync" | "prompt" - client: { - session?: { - status?: () => Promise - messages?: (input: { path: { id: string }; query: PromptMessagesQuery }) => Promise - } - } - sessionID: string - input: TInput - source: string - dedupeKey: string - settleMs: number - postDispatchHoldMs: number - dispatchTimeoutMs: number - checkStatus: boolean - checkToolState: boolean - dispatch: (input: TInput) => Promise -}): Promise { - const { - sessionName, - client, - sessionID, - input, - source, - dedupeKey, - settleMs, - postDispatchHoldMs, - dispatchTimeoutMs, - checkStatus, - checkToolState, - dispatch, - } = args - - const existing = getActiveReservation(sessionID) - if (existing) { - log(`[prompt-async-gate] ${sessionName} skipped because session is reserved`, { - sessionID, - source, - reservedBy: existing.source, - reservedAgeMs: Date.now() - existing.reservedAt, - }) - return { status: "reserved", reservedBy: existing.source } - } - - const reservation: PromptAsyncReservation = { - source, - dedupeKey, - reservedAt: Date.now(), - token: Symbol(source), - } - promptAsyncReservations.set(sessionID, reservation) - let dispatchAttempted = false - - try { - const canReadStatus = checkStatus && typeof client.session?.status === "function" - if (settleMs > 0) { - await settleAfterSessionIdle(settleMs) - } - - let sessionActive = false - if (canReadStatus) { - try { - sessionActive = await withDispatchTimeout( - isSessionActive(client, sessionID), - Math.min(dispatchTimeoutMs, 5000), - `[prompt-async-gate] ${sessionName} isSessionActive`, - ) - } catch { - sessionActive = false - } - } - if (sessionActive) { - log(`[prompt-async-gate] ${sessionName} skipped because session is active`, { sessionID, source }) - return { status: "active" } - } - - if ( - checkToolState - && typeof client.session?.messages === "function" - && await sessionLatestAssistantBlocksInternalPrompt({ - client, - sessionID, - input, - sessionName, - source, - timeoutMs: Math.min(dispatchTimeoutMs, getPromptGateMessagesFetchTimeoutMs()), - }) - ) { - log(`[prompt-async-gate] ${sessionName} skipped because latest assistant is still active`, { - sessionID, - source, - }) - return { status: "active" } - } - - log(`[prompt-async-gate] ${sessionName} dispatching`, { sessionID, source }) - dispatchAttempted = true - const response = await withDispatchTimeout( - dispatch(input), - dispatchTimeoutMs, - `[prompt-async-gate] ${sessionName} dispatch`, - ) - log(`[prompt-async-gate] ${sessionName} dispatched`, { sessionID, source }) - return { status: "dispatched", response } - } catch (error) { - log(`[prompt-async-gate] ${sessionName} failed`, { sessionID, source, error: String(error) }) - return { status: "failed", error, dispatchAttempted } - } finally { - const current = promptAsyncReservations.get(sessionID) - if (current?.token === reservation.token) { - if (dispatchAttempted && postDispatchHoldMs > 0) { - reservation.expiresAt = Date.now() + postDispatchHoldMs - } else { - promptAsyncReservations.delete(sessionID) - } - } - } -} - -async function drainPromptQueue(sessionID: string, awaitedEntry?: QueuedInternalPrompt): Promise { - if (promptQueueDraining.has(sessionID)) { - return awaitedEntry ? queuedResult(awaitedEntry, 1) : undefined - } - - promptQueueDraining.add(sessionID) - clearPromptQueueTimer(sessionID) - - let awaitedResult: InternalPromptDispatchResult | undefined - try { - while (true) { - const queue = promptQueues.get(sessionID) - const entry = queue?.[0] - if (!entry) { - break - } - - promptQueueInFlight.set(sessionID, entry) - const result = await dispatchAfterSessionIdle({ - sessionName: entry.sessionName, - client: entry.client, - sessionID: entry.sessionID, - input: entry.input, - source: entry.source, - dedupeKey: entry.dedupeKey, - settleMs: entry.settleMs, - postDispatchHoldMs: entry.postDispatchHoldMs, - dispatchTimeoutMs: entry.dispatchTimeoutMs, - checkStatus: entry.checkStatus, - checkToolState: entry.checkToolState, - dispatch: entry.dispatch, - }) - if (promptQueueInFlight.get(sessionID)?.id === entry.id) { - promptQueueInFlight.delete(sessionID) - } - - if (result.status === "active" || result.status === "reserved") { - const queued = queuedResult( - entry, - 1, - result.status === "reserved" ? result.reservedBy : entry.source, - ) - if (awaitedEntry?.id === entry.id) { - awaitedResult = queued - } - schedulePromptQueueDrain(sessionID, entry.queueRetryMs) - break - } - - removePromptQueueEntry(sessionID, entry) - if (awaitedEntry?.id === entry.id) { - awaitedResult = result - } - - const remainingQueue = promptQueues.get(sessionID) - if (!remainingQueue || remainingQueue.length === 0) { - break - } - - schedulePromptQueueDrain(sessionID, entry.postDispatchHoldMs) - break - } - } finally { - promptQueueDraining.delete(sessionID) - } - - return awaitedResult -} - -async function enqueueInternalPrompt(entry: QueuedInternalPrompt): Promise { - const activeReservation = getActiveReservation(entry.sessionID) - if (activeReservation?.dedupeKey === entry.dedupeKey) { - log("[prompt-async-gate] queued prompt coalesced with recent dispatch", { - sessionID: entry.sessionID, - source: entry.source, - queuedBy: activeReservation.source, - }) - return queuedResult(entry, 0, activeReservation.source) - } - - const queue = getPromptQueue(entry.sessionID) - const existingIndex = queue.findIndex((queued) => queued.dedupeKey === entry.dedupeKey) - if (existingIndex >= 0) { - const existing = queue[existingIndex] - if (existing) { - log("[prompt-async-gate] queued prompt coalesced with pending dispatch", { - sessionID: entry.sessionID, - source: entry.source, - queuedBy: existing.source, - position: existingIndex + 1, - }) - return queuedResult(existing, existingIndex + 1) - } - } - - queue.push(entry) - log("[prompt-async-gate] queued prompt accepted", { - sessionID: entry.sessionID, - source: entry.source, - position: queue.length, - }) - - if (queue.length > 1 || promptQueueDraining.has(entry.sessionID)) { - schedulePromptQueueDrain(entry.sessionID, 0) - return queuedResult(entry, queue.length) - } - - const result = await drainPromptQueue(entry.sessionID, entry) - return result ?? queuedResult(entry, 1) -} - export async function dispatchInternalPrompt( args: InternalPromptDispatchArgs, ): Promise { @@ -777,7 +115,7 @@ export async function dispatchInternalPrompt( } const queuedBy = getQueuedPromptBlocker(sessionID) - if (queuedBy !== undefined || promptQueueDraining.has(sessionID)) { + if (queuedBy !== undefined || isPromptQueueDraining(sessionID)) { return { status: "reserved", reservedBy: queuedBy ?? source } } @@ -799,7 +137,7 @@ export async function dispatchInternalPrompt( if (args.queue !== false) { return enqueueInternalPrompt({ - id: promptQueueSequence += 1, + id: nextPromptQueueID(), sessionID, sessionName, client, @@ -812,7 +150,7 @@ export async function dispatchInternalPrompt( queueRetryMs, checkStatus: args.checkStatus !== false, checkToolState: args.checkToolState !== false, - dispatch: dispatch as (dispatchInput: unknown) => Promise, + dispatch: async (_dispatchInput: unknown) => dispatch(input), }) } @@ -833,15 +171,9 @@ export async function dispatchInternalPrompt( } export function releaseAllPromptAsyncReservationsForTesting(): void { - promptAsyncReservations.clear() - promptQueues.clear() - promptQueueDraining.clear() - promptQueueInFlight.clear() - for (const timer of promptQueueTimers.values()) { - clearTimeout(timer) - } - promptQueueTimers.clear() - promptGateMessagesFetchTimeoutMsForTesting = undefined + clearPromptReservationsForTesting() + clearPromptQueueStateForTesting() + resetPromptGateTimingForTesting() } export function isInternalPromptDispatchAccepted(result: InternalPromptDispatchResult): boolean { @@ -853,7 +185,7 @@ export function releasePromptAsyncReservation( source: string, options?: PromptAsyncReservationReleaseOptions, ): boolean { - const existing = promptAsyncReservations.get(sessionID) + const existing = getPromptReservation(sessionID) if (!existing) { return false } @@ -868,13 +200,8 @@ export function releasePromptAsyncReservation( return false } - promptAsyncReservations.delete(sessionID) - const inFlight = promptQueueInFlight.get(sessionID) - if (inFlight?.dedupeKey === existing.dedupeKey) { - removePromptQueueEntry(sessionID, inFlight) - promptQueueInFlight.delete(sessionID) - promptQueueDraining.delete(sessionID) - } + deletePromptReservation(sessionID) + releaseInFlightPromptMatchingDedupe(sessionID, existing.dedupeKey) schedulePromptQueueDrain(sessionID, 0) log("[prompt-async-gate] promptAsync reservation released", { sessionID, diff --git a/src/shared/prompt-async-gate/pending-tool-turn.ts b/src/shared/prompt-async-gate/pending-tool-turn.ts new file mode 100644 index 000000000..ddc36853c --- /dev/null +++ b/src/shared/prompt-async-gate/pending-tool-turn.ts @@ -0,0 +1,206 @@ +import { log } from "../logger" +import { + isSyntheticOrInternalUserMessage, + type InternalInitiatorMessageLike, + type InternalInitiatorTextPartLike, +} from "../internal-initiator-marker" +import { isRecord } from "../record-type-guard" +import { withDispatchTimeout } from "./timing" +import type { PromptDispatchClient, PromptMessagesQuery, PromptSessionName } from "./types" + +function getPromptQuery(input: unknown): PromptMessagesQuery { + if (!isRecord(input)) { + return { directory: "" } + } + const query = input.query + if (!isRecord(query)) { + return { directory: "" } + } + + const promptQuery: PromptMessagesQuery = { directory: "" } + if (typeof query.directory === "string") { + return typeof query.limit === "number" + ? { directory: query.directory, limit: query.limit } + : { directory: query.directory } + } + if (typeof query.limit === "number") { + return { ...promptQuery, limit: query.limit } + } + return promptQuery +} + +function getMessagesData(response: unknown): unknown[] { + if (isRecord(response) && Array.isArray(response.data)) { + return response.data + } + return Array.isArray(response) ? response : [] +} + +function messageRole(message: unknown): string | undefined { + if (!isRecord(message)) { + return undefined + } + const info = message.info + if (isRecord(info) && typeof info.role === "string") { + return info.role + } + return typeof message.role === "string" ? message.role : undefined +} + +function messageFinish(message: unknown): string | true | undefined { + if (!isRecord(message)) { + return undefined + } + const info = message.info + if (isRecord(info)) { + if (info.finish === true) { + return true + } + if (typeof info.finish === "string" && info.finish.length > 0) { + return info.finish + } + } + if (message.finish === true) { + return true + } + return typeof message.finish === "string" && message.finish.length > 0 ? message.finish : undefined +} + +function messageCompleted(message: unknown): boolean { + if (!isRecord(message)) { + return false + } + const info = message.info + const time = isRecord(info) && isRecord(info.time) ? info.time : undefined + const completed = time?.completed + if (typeof completed === "number" && Number.isFinite(completed)) { + return true + } + return typeof completed === "string" && completed.length > 0 +} + +function toInternalInitiatorTextPartLike(part: unknown): InternalInitiatorTextPartLike { + const result: InternalInitiatorTextPartLike = {} + if (!isRecord(part)) { + return result + } + + if (typeof part.type === "string") { + result.type = part.type + } + if (typeof part.text === "string") { + result.text = part.text + } + if (typeof part.synthetic === "boolean") { + result.synthetic = part.synthetic + } + return result +} + +function toInternalInitiatorMessageLike(message: unknown): InternalInitiatorMessageLike | undefined { + if (!isRecord(message)) { + return undefined + } + + const result: InternalInitiatorMessageLike = {} + const info = message.info + if (isRecord(info) && typeof info.role === "string") { + result.info = { role: info.role } + } + if (typeof message.role === "string") { + result.role = message.role + } + if (Array.isArray(message.parts)) { + result.parts = message.parts.map(toInternalInitiatorTextPartLike) + } + return result +} + +function messageIsSyntheticOrInternalUser(message: unknown): boolean { + const initiatorMessage = toInternalInitiatorMessageLike(message) + return initiatorMessage !== undefined && isSyntheticOrInternalUserMessage(initiatorMessage) +} + +function partIsWaitingOnTool(part: unknown): boolean { + if (!isRecord(part)) { + return false + } + if ( + part.type !== "tool" + && part.type !== "tool_use" + && part.type !== "tool-call" + && part.type !== "tool-invocation" + ) { + return false + } + + const state = part.state + if (!isRecord(state)) { + return false + } + return state.status === "pending" || state.status === "running" +} + +function latestAssistantTurnBlocksInternalPrompt(messages: unknown[]): boolean { + for (let index = messages.length - 1; index >= 0; index--) { + const message = messages[index] + const role = messageRole(message) + if (role === "assistant") { + if (messageCompleted(message)) { + return false + } + const finish = messageFinish(message) + if (finish === true) { + return false + } + if (finish === undefined || finish === "unknown") { + return true + } + if (!isRecord(message) || !Array.isArray(message.parts)) { + return finish === "tool-calls" + } + return finish === "tool-calls" || message.parts.some(partIsWaitingOnTool) + } + if (role === "user") { + if (messageIsSyntheticOrInternalUser(message)) { + continue + } + return false + } + } + return false +} + +export async function sessionLatestAssistantBlocksInternalPrompt(args: { + readonly client: PromptDispatchClient + readonly sessionID: string + readonly input: TInput + readonly sessionName: PromptSessionName + readonly source: string + readonly timeoutMs: number +}): Promise { + const session = args.client.session + if (typeof session?.messages !== "function") { + return false + } + const messages = session.messages.bind(session) + + try { + const response = await withDispatchTimeout( + messages({ + path: { id: args.sessionID }, + query: getPromptQuery(args.input), + }), + args.timeoutMs, + `[prompt-async-gate] ${args.sessionName} session.messages`, + ) + return latestAssistantTurnBlocksInternalPrompt(getMessagesData(response)) + } catch (error) { + log("[prompt-async-gate] latest assistant prompt-block check failed", { + sessionID: args.sessionID, + source: args.source, + error: String(error), + }) + return false + } +} diff --git a/src/shared/prompt-async-gate/queue.ts b/src/shared/prompt-async-gate/queue.ts new file mode 100644 index 000000000..2aa649c8c --- /dev/null +++ b/src/shared/prompt-async-gate/queue.ts @@ -0,0 +1,233 @@ +import { log } from "../logger" +import { dispatchAfterSessionIdle } from "./session-idle-dispatch" +import { + getActiveReservation, + setExpiredReservationHandler, +} from "./reservations" +import type { InternalPromptDispatchResult, QueuedInternalPrompt } from "./types" + +declare function setTimeout(callback: () => void, delay?: number): unknown +declare function clearTimeout(timeout: unknown): void + +const promptQueues = new Map() +const promptQueueDraining = new Set() +const promptQueueInFlight = new Map() +const promptQueueTimers = new Map() +let promptQueueSequence = 0 + +setExpiredReservationHandler((sessionID) => { + schedulePromptQueueDrain(sessionID, 0) +}) + +function getPromptQueue(sessionID: string): QueuedInternalPrompt[] { + const existing = promptQueues.get(sessionID) + if (existing) { + return existing + } + + const queue: QueuedInternalPrompt[] = [] + promptQueues.set(sessionID, queue) + return queue +} + +function setPromptQueue(sessionID: string, queue: QueuedInternalPrompt[]): void { + if (queue.length === 0) { + promptQueues.delete(sessionID) + return + } + promptQueues.set(sessionID, queue) +} + +function queuedResult(entry: QueuedInternalPrompt, position: number, queuedBy = entry.source): InternalPromptDispatchResult { + return { + status: "queued", + queuedBy, + position, + } +} + +function clearPromptQueueTimer(sessionID: string): void { + const timer = promptQueueTimers.get(sessionID) + if (timer !== undefined) { + clearTimeout(timer) + promptQueueTimers.delete(sessionID) + } +} + +export function schedulePromptQueueDrain(sessionID: string, delayMs: number): void { + const queue = promptQueues.get(sessionID) + if (!queue || queue.length === 0) { + clearPromptQueueTimer(sessionID) + return + } + + clearPromptQueueTimer(sessionID) + const timer = setTimeout(() => { + promptQueueTimers.delete(sessionID) + void drainPromptQueue(sessionID).catch((error: unknown) => { + log("[prompt-async-gate] queued prompt drain failed", { + sessionID, + error: String(error), + }) + }) + }, Math.max(0, delayMs)) + promptQueueTimers.set(sessionID, timer) +} + +function removePromptQueueEntry(sessionID: string, entry: QueuedInternalPrompt): void { + const queue = promptQueues.get(sessionID) + if (!queue) { + return + } + const nextQueue = queue.filter((queued) => queued.id !== entry.id) + setPromptQueue(sessionID, nextQueue) +} + +export function getQueuedPromptBlocker(sessionID: string): string | undefined { + const inFlight = promptQueueInFlight.get(sessionID) + if (inFlight) { + return inFlight.source + } + + const queue = promptQueues.get(sessionID) + return queue?.[0]?.source +} + +export function isPromptQueueDraining(sessionID: string): boolean { + return promptQueueDraining.has(sessionID) +} + +export function nextPromptQueueID(): number { + promptQueueSequence += 1 + return promptQueueSequence +} + +export function releaseInFlightPromptMatchingDedupe(sessionID: string, dedupeKey: string): void { + const inFlight = promptQueueInFlight.get(sessionID) + if (inFlight?.dedupeKey === dedupeKey) { + removePromptQueueEntry(sessionID, inFlight) + promptQueueInFlight.delete(sessionID) + promptQueueDraining.delete(sessionID) + } +} + +export function clearPromptQueueStateForTesting(): void { + promptQueues.clear() + promptQueueDraining.clear() + promptQueueInFlight.clear() + for (const timer of promptQueueTimers.values()) { + clearTimeout(timer) + } + promptQueueTimers.clear() +} + +async function drainPromptQueue(sessionID: string, awaitedEntry?: QueuedInternalPrompt): Promise { + if (promptQueueDraining.has(sessionID)) { + return awaitedEntry ? queuedResult(awaitedEntry, 1) : undefined + } + + promptQueueDraining.add(sessionID) + clearPromptQueueTimer(sessionID) + + let awaitedResult: InternalPromptDispatchResult | undefined + try { + while (true) { + const queue = promptQueues.get(sessionID) + const entry = queue?.[0] + if (!entry) { + break + } + + promptQueueInFlight.set(sessionID, entry) + const result = await dispatchAfterSessionIdle({ + sessionName: entry.sessionName, + client: entry.client, + sessionID: entry.sessionID, + input: entry.input, + source: entry.source, + dedupeKey: entry.dedupeKey, + settleMs: entry.settleMs, + postDispatchHoldMs: entry.postDispatchHoldMs, + dispatchTimeoutMs: entry.dispatchTimeoutMs, + checkStatus: entry.checkStatus, + checkToolState: entry.checkToolState, + dispatch: entry.dispatch, + }) + if (promptQueueInFlight.get(sessionID)?.id === entry.id) { + promptQueueInFlight.delete(sessionID) + } + + if (result.status === "active" || result.status === "reserved") { + const queued = queuedResult( + entry, + 1, + result.status === "reserved" ? result.reservedBy : entry.source, + ) + if (awaitedEntry?.id === entry.id) { + awaitedResult = queued + } + schedulePromptQueueDrain(sessionID, entry.queueRetryMs) + break + } + + removePromptQueueEntry(sessionID, entry) + if (awaitedEntry?.id === entry.id) { + awaitedResult = result + } + + const remainingQueue = promptQueues.get(sessionID) + if (!remainingQueue || remainingQueue.length === 0) { + break + } + + schedulePromptQueueDrain(sessionID, entry.postDispatchHoldMs) + break + } + } finally { + promptQueueDraining.delete(sessionID) + } + + return awaitedResult +} + +export async function enqueueInternalPrompt(entry: QueuedInternalPrompt): Promise { + const activeReservation = getActiveReservation(entry.sessionID) + if (activeReservation?.dedupeKey === entry.dedupeKey) { + log("[prompt-async-gate] queued prompt coalesced with recent dispatch", { + sessionID: entry.sessionID, + source: entry.source, + queuedBy: activeReservation.source, + }) + return queuedResult(entry, 0, activeReservation.source) + } + + const queue = getPromptQueue(entry.sessionID) + const existingIndex = queue.findIndex((queued) => queued.dedupeKey === entry.dedupeKey) + if (existingIndex >= 0) { + const existing = queue[existingIndex] + if (existing) { + log("[prompt-async-gate] queued prompt coalesced with pending dispatch", { + sessionID: entry.sessionID, + source: entry.source, + queuedBy: existing.source, + position: existingIndex + 1, + }) + return queuedResult(existing, existingIndex + 1) + } + } + + queue.push(entry) + log("[prompt-async-gate] queued prompt accepted", { + sessionID: entry.sessionID, + source: entry.source, + position: queue.length, + }) + + if (queue.length > 1 || promptQueueDraining.has(entry.sessionID)) { + schedulePromptQueueDrain(entry.sessionID, 0) + return queuedResult(entry, queue.length) + } + + const result = await drainPromptQueue(entry.sessionID, entry) + return result ?? queuedResult(entry, 1) +} diff --git a/src/shared/prompt-async-gate/reservations.ts b/src/shared/prompt-async-gate/reservations.ts new file mode 100644 index 000000000..135f5ff93 --- /dev/null +++ b/src/shared/prompt-async-gate/reservations.ts @@ -0,0 +1,95 @@ +import { log } from "../logger" +import type { PromptAsyncReservation, PromptAsyncReservationReleaseOptions } from "./types" + +const promptAsyncReservations = new Map() +let expiredReservationHandler: ((sessionID: string) => void) | undefined + +export function setExpiredReservationHandler(handler: (sessionID: string) => void): void { + expiredReservationHandler = handler +} + +function notifyExpiredReservation(sessionID: string): void { + expiredReservationHandler?.(sessionID) +} + +function pruneExpiredReservations(now = Date.now()): void { + const expiredSessionIDs: string[] = [] + for (const [sessionID, reservation] of promptAsyncReservations) { + if (typeof reservation.expiresAt === "number" && reservation.expiresAt <= now) { + promptAsyncReservations.delete(sessionID) + expiredSessionIDs.push(sessionID) + log("[prompt-async-gate] expired reservation released", { + sessionID, + source: reservation.source, + }) + } + } + for (const sessionID of expiredSessionIDs) { + notifyExpiredReservation(sessionID) + } +} + +export function getActiveReservation(sessionID: string): PromptAsyncReservation | undefined { + pruneExpiredReservations() + return promptAsyncReservations.get(sessionID) +} + +export function getPromptReservation(sessionID: string): PromptAsyncReservation | undefined { + return promptAsyncReservations.get(sessionID) +} + +export function setPromptReservation(sessionID: string, reservation: PromptAsyncReservation): void { + promptAsyncReservations.set(sessionID, reservation) +} + +export function finishPromptReservation( + sessionID: string, + reservation: PromptAsyncReservation, + dispatchAttempted: boolean, + postDispatchHoldMs: number, +): void { + const current = promptAsyncReservations.get(sessionID) + if (current?.token !== reservation.token) { + return + } + + if (dispatchAttempted && postDispatchHoldMs > 0) { + promptAsyncReservations.set(sessionID, { + ...reservation, + expiresAt: Date.now() + postDispatchHoldMs, + }) + return + } + promptAsyncReservations.delete(sessionID) +} + +export function deletePromptReservation(sessionID: string): void { + promptAsyncReservations.delete(sessionID) +} + +export function clearPromptReservationsForTesting(): void { + promptAsyncReservations.clear() +} + +export function reservationSourceMatches( + reservationSource: string, + expectedSource: string | readonly string[], + expectedPrefix?: PromptAsyncReservationReleaseOptions["reservedByPrefix"], +): boolean { + if (typeof expectedSource === "string") { + if (reservationSource === expectedSource) { + return true + } + } else if (expectedSource.includes(reservationSource)) { + return true + } + + if (expectedPrefix === undefined) { + return false + } + + const prefixes = typeof expectedPrefix === "string" ? [expectedPrefix] : expectedPrefix + return prefixes + .filter((prefix) => prefix.length > 0 && prefix.endsWith(":")) + .some((prefix) => reservationSource.startsWith(prefix)) +} diff --git a/src/shared/prompt-async-gate/session-idle-dispatch.ts b/src/shared/prompt-async-gate/session-idle-dispatch.ts new file mode 100644 index 000000000..2e4fc62ef --- /dev/null +++ b/src/shared/prompt-async-gate/session-idle-dispatch.ts @@ -0,0 +1,118 @@ +import { log } from "../logger" +import { isSessionActive, settleAfterSessionIdle } from "../session-idle-settle" +import { sessionLatestAssistantBlocksInternalPrompt } from "./pending-tool-turn" +import { + finishPromptReservation, + getActiveReservation, + setPromptReservation, +} from "./reservations" +import { getPromptGateMessagesFetchTimeoutMs, withDispatchTimeout } from "./timing" +import type { InternalPromptDispatchResult, PromptAsyncReservation, PromptDispatchClient, PromptSessionName } from "./types" + +export async function dispatchAfterSessionIdle(args: { + readonly sessionName: PromptSessionName + readonly client: PromptDispatchClient + readonly sessionID: string + readonly input: TInput + readonly source: string + readonly dedupeKey: string + readonly settleMs: number + readonly postDispatchHoldMs: number + readonly dispatchTimeoutMs: number + readonly checkStatus: boolean + readonly checkToolState: boolean + readonly dispatch: (input: TInput) => Promise +}): Promise { + const { + sessionName, + client, + sessionID, + input, + source, + dedupeKey, + settleMs, + postDispatchHoldMs, + dispatchTimeoutMs, + checkStatus, + checkToolState, + dispatch, + } = args + + const existing = getActiveReservation(sessionID) + if (existing) { + log(`[prompt-async-gate] ${sessionName} skipped because session is reserved`, { + sessionID, + source, + reservedBy: existing.source, + reservedAgeMs: Date.now() - existing.reservedAt, + }) + return { status: "reserved", reservedBy: existing.source } + } + + const reservation: PromptAsyncReservation = { + source, + dedupeKey, + reservedAt: Date.now(), + token: Symbol(source), + } + setPromptReservation(sessionID, reservation) + let dispatchAttempted = false + + try { + const canReadStatus = checkStatus && typeof client.session?.status === "function" + if (settleMs > 0) { + await settleAfterSessionIdle(settleMs) + } + + let sessionActive = false + if (canReadStatus) { + try { + sessionActive = await withDispatchTimeout( + isSessionActive(client, sessionID), + Math.min(dispatchTimeoutMs, 5000), + `[prompt-async-gate] ${sessionName} isSessionActive`, + ) + } catch { + sessionActive = false + } + } + if (sessionActive) { + log(`[prompt-async-gate] ${sessionName} skipped because session is active`, { sessionID, source }) + return { status: "active" } + } + + if ( + checkToolState + && typeof client.session?.messages === "function" + && await sessionLatestAssistantBlocksInternalPrompt({ + client, + sessionID, + input, + sessionName, + source, + timeoutMs: Math.min(dispatchTimeoutMs, getPromptGateMessagesFetchTimeoutMs()), + }) + ) { + log(`[prompt-async-gate] ${sessionName} skipped because latest assistant is still active`, { + sessionID, + source, + }) + return { status: "active" } + } + + log(`[prompt-async-gate] ${sessionName} dispatching`, { sessionID, source }) + dispatchAttempted = true + const response = await withDispatchTimeout( + dispatch(input), + dispatchTimeoutMs, + `[prompt-async-gate] ${sessionName} dispatch`, + ) + log(`[prompt-async-gate] ${sessionName} dispatched`, { sessionID, source }) + return { status: "dispatched", response } + } catch (error) { + log(`[prompt-async-gate] ${sessionName} failed`, { sessionID, source, error: String(error) }) + return { status: "failed", error, dispatchAttempted } + } finally { + finishPromptReservation(sessionID, reservation, dispatchAttempted, postDispatchHoldMs) + } +} diff --git a/src/shared/prompt-async-gate/timing.ts b/src/shared/prompt-async-gate/timing.ts new file mode 100644 index 000000000..fc859f422 --- /dev/null +++ b/src/shared/prompt-async-gate/timing.ts @@ -0,0 +1,46 @@ +export const DEFAULT_PROMPT_ASYNC_POST_DISPATCH_HOLD_MS = 2_000 +export const DEFAULT_PROMPT_DISPATCH_TIMEOUT_MS = 30_000 +export const DEFAULT_PROMPT_GATE_MESSAGES_FETCH_TIMEOUT_MS = 5_000 +export const DEFAULT_PROMPT_QUEUE_RETRY_MS = 250 + +declare function setTimeout(callback: () => void, delay?: number): unknown +declare function clearTimeout(timeout: unknown): void + +let promptGateMessagesFetchTimeoutMsForTesting: number | undefined + +export function _setPromptGateMessagesFetchTimeoutMsForTesting(value: number | undefined): void { + promptGateMessagesFetchTimeoutMsForTesting = value +} + +export function getPromptGateMessagesFetchTimeoutMs(): number { + return promptGateMessagesFetchTimeoutMsForTesting ?? DEFAULT_PROMPT_GATE_MESSAGES_FETCH_TIMEOUT_MS +} + +export function resetPromptGateTimingForTesting(): void { + promptGateMessagesFetchTimeoutMsForTesting = undefined +} + +export async function withDispatchTimeout( + operation: Promise, + dispatchTimeoutMs: number, + operationName: string, +): Promise { + if (dispatchTimeoutMs <= 0) { + return operation + } + + let timeoutID: unknown + const timeoutPromise = new Promise((_, reject) => { + timeoutID = setTimeout(() => { + reject(new Error(`${operationName} timed out after ${dispatchTimeoutMs}ms`)) + }, dispatchTimeoutMs) + }) + + try { + return await Promise.race([operation, timeoutPromise]) + } finally { + if (timeoutID !== undefined) { + clearTimeout(timeoutID) + } + } +} diff --git a/src/shared/prompt-async-gate/types.ts b/src/shared/prompt-async-gate/types.ts new file mode 100644 index 000000000..082693692 --- /dev/null +++ b/src/shared/prompt-async-gate/types.ts @@ -0,0 +1,99 @@ +export type PromptAsyncInput = { + readonly path?: { readonly id?: string } + readonly body?: unknown + readonly query?: unknown + readonly signal?: unknown + readonly [key: string]: unknown +} + +export type PromptMessagesQuery = { + readonly directory: string + readonly limit?: number +} + +export type PromptAsyncClient = { + readonly session?: { + readonly status?: () => Promise + readonly messages?: (input: { readonly path: { readonly id: string }; readonly query: PromptMessagesQuery }) => Promise + readonly promptAsync?: (input: TInput) => Promise + } +} + +export type PromptClient = { + readonly session?: { + readonly status?: () => Promise + readonly messages?: (input: { readonly path: { readonly id: string }; readonly query: PromptMessagesQuery }) => Promise + readonly prompt?: (input: TInput) => Promise + } +} + +export type InternalPromptDispatchMode = "async" | "sync" +export type InternalPromptQueueBehavior = "enqueue" | "defer" +export type PromptSessionName = "promptAsync" | "prompt" + +type InternalPromptDispatchCommonArgs = { + readonly sessionID: string + readonly input: TInput + readonly source: string + readonly dedupeKey?: string + readonly queueBehavior?: InternalPromptQueueBehavior + readonly queue?: boolean + readonly queueRetryMs?: number + readonly settleMs?: number + readonly postDispatchHoldMs?: number + readonly dispatchTimeoutMs?: number + readonly checkStatus?: boolean + readonly checkToolState?: boolean +} + +export type InternalPromptDispatchArgs = InternalPromptDispatchCommonArgs & ( + | { readonly mode: "async"; readonly client: PromptAsyncClient } + | { readonly mode: "sync"; readonly client: PromptClient } +) + +export type PromptAsyncReservation = { + readonly source: string + readonly dedupeKey: string + readonly reservedAt: number + readonly token: symbol + readonly expiresAt?: number +} + +export type InternalPromptDispatchResult = + | { readonly status: "dispatched"; readonly response: unknown } + | { readonly status: "queued"; readonly queuedBy: string; readonly position: number } + | { readonly status: "active" } + | { readonly status: "reserved"; readonly reservedBy: string } + | { readonly status: "unavailable" } + | { readonly status: "failed"; readonly error: unknown; readonly dispatchAttempted: boolean } + +export type PromptAsyncGateResult = InternalPromptDispatchResult + +export type PromptAsyncReservationReleaseOptions = { + readonly reservedBy?: string | readonly string[] + readonly reservedByPrefix?: string | readonly string[] +} + +export type PromptDispatchClient = { + readonly session?: { + readonly status?: () => Promise + readonly messages?: (input: { readonly path: { readonly id: string }; readonly query: PromptMessagesQuery }) => Promise + } +} + +export type QueuedInternalPrompt = { + readonly id: number + readonly sessionID: string + readonly sessionName: PromptSessionName + readonly client: PromptDispatchClient + readonly input: unknown + readonly source: string + readonly dedupeKey: string + readonly settleMs: number + readonly postDispatchHoldMs: number + readonly dispatchTimeoutMs: number + readonly queueRetryMs: number + readonly checkStatus: boolean + readonly checkToolState: boolean + readonly dispatch: (input: unknown) => Promise +}