fix(prompt-gate): share message reservations
This commit is contained in:
YeonGyu-Kim
@@ -1,6 +1,7 @@
|
||||
import { afterEach, describe, expect, test } from "bun:test"
|
||||
|
||||
import {
|
||||
promptAfterSessionIdle,
|
||||
promptAsyncAfterSessionIdle,
|
||||
releaseAllPromptAsyncReservationsForTesting,
|
||||
} from "./prompt-async-gate"
|
||||
@@ -81,4 +82,48 @@ describe("promptAsyncAfterSessionIdle", () => {
|
||||
expect(result.status).toBe("active")
|
||||
expect(promptCalls).toBe(0)
|
||||
})
|
||||
|
||||
test("#given two internal prompt calls race for one idle session #when they dispatch concurrently #then only one prompt is accepted", async () => {
|
||||
// given
|
||||
let promptCalls = 0
|
||||
let releasePrompt: (() => void) | undefined
|
||||
const promptGate = new Promise<void>((resolve) => {
|
||||
releasePrompt = resolve
|
||||
})
|
||||
const client = {
|
||||
session: {
|
||||
status: async () => ({ data: { ses_prompt_race: { type: "idle" } } }),
|
||||
prompt: async () => {
|
||||
promptCalls += 1
|
||||
await promptGate
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
// when
|
||||
const first = promptAfterSessionIdle({
|
||||
client,
|
||||
sessionID: "ses_prompt_race",
|
||||
input: { path: { id: "ses_prompt_race" }, body: { parts: [] } },
|
||||
source: "test:prompt:first",
|
||||
settleMs: 0,
|
||||
postDispatchHoldMs: 0,
|
||||
})
|
||||
await Promise.resolve()
|
||||
const second = await promptAfterSessionIdle({
|
||||
client,
|
||||
sessionID: "ses_prompt_race",
|
||||
input: { path: { id: "ses_prompt_race" }, body: { parts: [] } },
|
||||
source: "test:prompt:second",
|
||||
settleMs: 0,
|
||||
postDispatchHoldMs: 0,
|
||||
})
|
||||
releasePrompt?.()
|
||||
const firstResult = await first
|
||||
|
||||
// then
|
||||
expect(firstResult.status).toBe("dispatched")
|
||||
expect(second.status).toBe("reserved")
|
||||
expect(promptCalls).toBe(1)
|
||||
})
|
||||
})
|
||||
|
||||
@@ -1,111 +1 @@
|
||||
import { log } from "../../shared/logger"
|
||||
import {
|
||||
DEFAULT_SESSION_IDLE_SETTLE_MS,
|
||||
isSessionActive,
|
||||
settleAfterSessionIdle,
|
||||
} from "./session-idle-settle"
|
||||
|
||||
export const DEFAULT_PROMPT_ASYNC_POST_DISPATCH_HOLD_MS = 250
|
||||
|
||||
type PromptAsyncInput = {
|
||||
path?: { id?: string }
|
||||
body?: unknown
|
||||
query?: unknown
|
||||
signal?: unknown
|
||||
[key: string]: unknown
|
||||
}
|
||||
|
||||
type PromptAsyncClient<TInput> = {
|
||||
session?: {
|
||||
status?: () => Promise<unknown>
|
||||
promptAsync?: (input: TInput) => Promise<unknown>
|
||||
}
|
||||
}
|
||||
|
||||
type PromptAsyncReservation = {
|
||||
source: string
|
||||
reservedAt: number
|
||||
token: symbol
|
||||
}
|
||||
|
||||
export type PromptAsyncGateResult =
|
||||
| { status: "dispatched"; response: unknown }
|
||||
| { status: "active" }
|
||||
| { status: "reserved"; reservedBy: string }
|
||||
| { status: "unavailable" }
|
||||
| { status: "failed"; error: unknown }
|
||||
|
||||
const promptAsyncReservations = new Map<string, PromptAsyncReservation>()
|
||||
|
||||
export async function promptAsyncAfterSessionIdle<TInput = PromptAsyncInput>(args: {
|
||||
client: PromptAsyncClient<TInput>
|
||||
sessionID: string
|
||||
input: TInput
|
||||
source: string
|
||||
settleMs?: number
|
||||
postDispatchHoldMs?: number
|
||||
}): Promise<PromptAsyncGateResult> {
|
||||
const {
|
||||
client,
|
||||
sessionID,
|
||||
input,
|
||||
source,
|
||||
settleMs = DEFAULT_SESSION_IDLE_SETTLE_MS,
|
||||
} = args
|
||||
const postDispatchHoldMs = args.postDispatchHoldMs ?? (
|
||||
settleMs > 0 ? DEFAULT_PROMPT_ASYNC_POST_DISPATCH_HOLD_MS : 0
|
||||
)
|
||||
|
||||
if (typeof client.session?.promptAsync !== "function") {
|
||||
log("[prompt-async-gate] promptAsync unavailable", { sessionID, source })
|
||||
return { status: "unavailable" }
|
||||
}
|
||||
|
||||
const existing = promptAsyncReservations.get(sessionID)
|
||||
if (existing) {
|
||||
log("[prompt-async-gate] promptAsync 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,
|
||||
reservedAt: Date.now(),
|
||||
token: Symbol(source),
|
||||
}
|
||||
promptAsyncReservations.set(sessionID, reservation)
|
||||
|
||||
try {
|
||||
const canReadStatus = typeof client.session?.status === "function"
|
||||
await settleAfterSessionIdle(settleMs)
|
||||
|
||||
if (canReadStatus && await isSessionActive(client, sessionID)) {
|
||||
log("[prompt-async-gate] promptAsync skipped because session is active", { sessionID, source })
|
||||
return { status: "active" }
|
||||
}
|
||||
|
||||
log("[prompt-async-gate] promptAsync dispatching", { sessionID, source })
|
||||
const response = await client.session.promptAsync(input)
|
||||
if (canReadStatus) {
|
||||
await settleAfterSessionIdle(postDispatchHoldMs)
|
||||
}
|
||||
log("[prompt-async-gate] promptAsync dispatched", { sessionID, source })
|
||||
return { status: "dispatched", response }
|
||||
} catch (error) {
|
||||
log("[prompt-async-gate] promptAsync failed", { sessionID, source, error: String(error) })
|
||||
return { status: "failed", error }
|
||||
} finally {
|
||||
const current = promptAsyncReservations.get(sessionID)
|
||||
if (current?.token === reservation.token) {
|
||||
promptAsyncReservations.delete(sessionID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export function releaseAllPromptAsyncReservationsForTesting(): void {
|
||||
promptAsyncReservations.clear()
|
||||
}
|
||||
export * from "../../shared/prompt-async-gate"
|
||||
|
||||
@@ -1,61 +1 @@
|
||||
export const DEFAULT_SESSION_IDLE_SETTLE_MS = 150
|
||||
|
||||
export function settleAfterSessionIdle(ms = DEFAULT_SESSION_IDLE_SETTLE_MS): Promise<void> {
|
||||
return ms > 0 ? new Promise((resolve) => setTimeout(resolve, ms)) : Promise.resolve()
|
||||
}
|
||||
|
||||
type SessionStatusClient = {
|
||||
session?: {
|
||||
status?: () => Promise<unknown>
|
||||
}
|
||||
}
|
||||
|
||||
const ACTIVE_SESSION_STATUSES = new Set(["busy", "retry", "running"])
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null
|
||||
}
|
||||
|
||||
function getSessionStatusPayload(response: unknown): Record<string, unknown> {
|
||||
if (isRecord(response) && isRecord(response.data)) {
|
||||
return response.data
|
||||
}
|
||||
|
||||
if (isRecord(response)) {
|
||||
return response
|
||||
}
|
||||
|
||||
return {}
|
||||
}
|
||||
|
||||
export function isActiveSessionStatusType(statusType: string): boolean {
|
||||
return ACTIVE_SESSION_STATUSES.has(statusType)
|
||||
}
|
||||
|
||||
export async function isSessionActive(client: SessionStatusClient, sessionID: string): Promise<boolean> {
|
||||
if (typeof client.session?.status !== "function") {
|
||||
return false
|
||||
}
|
||||
|
||||
try {
|
||||
const statusResult = await client.session.status()
|
||||
const status = getSessionStatusPayload(statusResult)[sessionID]
|
||||
if (!isRecord(status)) {
|
||||
return false
|
||||
}
|
||||
|
||||
const statusType = status.type
|
||||
return typeof statusType === "string" && isActiveSessionStatusType(statusType)
|
||||
} catch {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
export async function shouldPromptAfterSessionIdle(
|
||||
client: SessionStatusClient,
|
||||
sessionID: string,
|
||||
settleMs = DEFAULT_SESSION_IDLE_SETTLE_MS,
|
||||
): Promise<boolean> {
|
||||
await settleAfterSessionIdle(settleMs)
|
||||
return !(await isSessionActive(client, sessionID))
|
||||
}
|
||||
export * from "../../shared/session-idle-settle"
|
||||
|
||||
Reference in New Issue
Block a user