luna/oh-my-opencode
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(prompt-gate): harden sync and team prompt dispatch

Browse Source
This commit is contained in:
YeonGyu-Kim
2026-05-19 17:43:37 +09:00
parent 1492bffd20
commit bcea4a9d28
44 changed files with 688 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/shared/model-suggestion-retry.test.ts
+3 -2
View File
@@ -561,7 +561,7 @@ describe("promptSyncWithModelSuggestionRetry", () => {
expect(promptAsyncMock).toHaveBeenCalledTimes(0)
})
it("#given sync prompt retry just dispatched #when the same session is prompted again immediately #then the second caller is coalesced by the queue", async () => {
it("#given sync prompt retry just dispatched #when the same session is prompted again immediately #then the second caller is deferred instead of queued", async () => {
// given
const promptMock = mock(async () => undefined)
const client = {
@@ -579,9 +579,10 @@ describe("promptSyncWithModelSuggestionRetry", () => {
// when
await promptSyncWithModelSuggestionRetry(unsafeTestValue(client), args)
await promptSyncWithModelSuggestionRetry(unsafeTestValue(client), args)
const second = promptSyncWithModelSuggestionRetry(unsafeTestValue(client), args)
// then
await expect(second).rejects.toThrow("prompt skipped by gate: reserved")
expect(promptMock).toHaveBeenCalledTimes(1)
})
src/shared/model-suggestion-retry.ts
+2
View File
@@ -164,6 +164,7 @@ export async function promptSyncWithModelSuggestionRetry(
source: "model-suggestion-retry:sync",
settleMs: 0,
checkStatus: false,
checkToolState: false,
...(options.queueBehavior ? { queueBehavior: options.queueBehavior } : {}),
})
if (promptResult.status === "failed") {
@@ -223,6 +224,7 @@ export async function promptSyncWithModelSuggestionRetry(
source: "model-suggestion-retry:sync-retry",
settleMs: 0,
checkStatus: false,
checkToolState: false,
...(options.queueBehavior ? { queueBehavior: options.queueBehavior } : {}),
})
if (promptResult.status === "failed") {
src/shared/prompt-async-gate.ts
+16 -10
View File
@@ -414,15 +414,19 @@ function partIsWaitingOnTool(part: unknown): boolean {
return state.status === "pending" || state.status === "running"
}
function latestAssistantTurnIsWaitingOnTools(messages: unknown[]): boolean {
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 (!isRecord(message) || !Array.isArray(message.parts)) {
return messageFinish(message) === "tool-calls"
const finish = messageFinish(message)
if (finish === undefined) {
return true
}
return messageFinish(message) === "tool-calls" || message.parts.some(partIsWaitingOnTool)
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)) {
@@ -434,7 +438,7 @@ function latestAssistantTurnIsWaitingOnTools(messages: unknown[]): boolean {
return false
}
async function sessionLatestAssistantIsWaitingOnTools<TInput>(args: {
async function sessionLatestAssistantBlocksInternalPrompt<TInput>(args: {
client: { session?: { messages?: (input: { path: { id: string }; query: PromptMessagesQuery }) => Promise<unknown> } }
sessionID: string
input: TInput
@@ -457,9 +461,9 @@ async function sessionLatestAssistantIsWaitingOnTools<TInput>(args: {
args.timeoutMs,
`[prompt-async-gate] ${args.sessionName} session.messages`,
)
return latestAssistantTurnIsWaitingOnTools(getMessagesData(response))
return latestAssistantTurnBlocksInternalPrompt(getMessagesData(response))
} catch (error) {
log("[prompt-async-gate] latest assistant tool-state check failed", {
log("[prompt-async-gate] latest assistant prompt-block check failed", {
sessionID: args.sessionID,
source: args.source,
error: String(error),
@@ -548,7 +552,7 @@ async function dispatchAfterSessionIdle<TInput>(args: {
if (
checkToolState
&& typeof client.session?.messages === "function"
&& await sessionLatestAssistantIsWaitingOnTools({
&& await sessionLatestAssistantBlocksInternalPrompt({
client,
sessionID,
input,
@@ -557,7 +561,7 @@ async function dispatchAfterSessionIdle<TInput>(args: {
timeoutMs: Math.min(dispatchTimeoutMs, getPromptGateMessagesFetchTimeoutMs()),
})
) {
log(`[prompt-async-gate] ${sessionName} skipped because latest assistant is waiting on tools`, {
log(`[prompt-async-gate] ${sessionName} skipped because latest assistant is still active`, {
sessionID,
source,
})
@@ -737,7 +741,9 @@ export async function dispatchInternalPrompt<TInput = PromptAsyncInput>(
return { status: "unavailable" }
}
if (args.queueBehavior === "defer") {
const queueBehavior = args.queueBehavior ?? (args.mode === "sync" ? "defer" : "enqueue")
if (queueBehavior === "defer") {
const activeReservation = getActiveReservation(sessionID)
if (activeReservation) {
return { status: "reserved", reservedBy: activeReservation.source }
src/shared/prompt-async-route-audit.test.ts
+118
View File
@@ -205,6 +205,74 @@ function detectRawPromptInSnippet(contents: string): boolean {
return detected
}
function objectLiteralHasQueueBehavior(node: ts.ObjectLiteralExpression, sourceFile: ts.SourceFile): boolean {
return node.properties.some((property) => {
if (ts.isPropertyAssignment(property) || ts.isShorthandPropertyAssignment(property)) {
return getPropertyName(property.name) === "queueBehavior"
}
if (ts.isSpreadAssignment(property)) {
return property.expression.getText(sourceFile).includes("queueBehavior")
}
return false
})
}
function callExpressionName(node: ts.Expression): string | null {
const callee = unwrapExpression(node)
if (ts.isIdentifier(callee)) {
return callee.text
}
if (ts.isPropertyAccessExpression(callee) || ts.isPropertyAccessChain(callee)) {
return getPropertyName(callee.name)
}
return null
}
function findPromptGateCallsWithoutQueueBehavior(filePath: string, contents: string): number[] {
const sourceFile = ts.createSourceFile(filePath, contents, ts.ScriptTarget.Latest, true, ts.ScriptKind.TS)
const offenders: number[] = []
const visit = (node: ts.Node): void => {
if (ts.isCallExpression(node)) {
if (callExpressionName(node.expression) === "dispatchInternalPrompt") {
const firstArgument = node.arguments[0]
if (
!firstArgument
|| !ts.isObjectLiteralExpression(firstArgument)
|| !objectLiteralHasQueueBehavior(firstArgument, sourceFile)
) {
offenders.push(sourceFile.getLineAndCharacterOfPosition(node.getStart(sourceFile)).line + 1)
}
}
}
ts.forEachChild(node, visit)
}
visit(sourceFile)
return offenders
}
function findPromptRetryCallsWithoutQueueBehavior(filePath: string, contents: string): number[] {
const sourceFile = ts.createSourceFile(filePath, contents, ts.ScriptTarget.Latest, true, ts.ScriptKind.TS)
const offenders: number[] = []
const guardedNames = new Set(["promptWithModelSuggestionRetry", "promptSyncWithModelSuggestionRetry"])
const visit = (node: ts.Node): void => {
if (ts.isCallExpression(node) && guardedNames.has(callExpressionName(node.expression) ?? "")) {
const optionsArgument = node.arguments[2]
if (!optionsArgument || !ts.isObjectLiteralExpression(optionsArgument) || !objectLiteralHasQueueBehavior(optionsArgument, sourceFile)) {
offenders.push(sourceFile.getLineAndCharacterOfPosition(node.getStart(sourceFile)).line + 1)
}
}
ts.forEachChild(node, visit)
}
visit(sourceFile)
return offenders
}
describe("production prompt injection routes", () => {
test("#given a destructuring promptAsync reference #when audit scans snippet #then it is flagged", () => {
// given
@@ -250,6 +318,20 @@ describe("production prompt injection routes", () => {
expect(detected).toBe(true)
})
test("#given indirect dispatchInternalPrompt options #when audit scans snippet #then it is flagged", () => {
// given
const snippet = `
const options = { mode: "async", queueBehavior: "defer" }
await dispatchInternalPrompt(options)
`
// when
const offenders = findPromptGateCallsWithoutQueueBehavior("audit-snippet.ts", snippet)
// then
expect(offenders).toEqual([3])
})
test("#given production TypeScript sources #when prompt routes are audited #then only the shared gate may call raw OpenCode prompt APIs", async () => {
// given
const files = await listSourceFiles(SOURCE_ROOT)
@@ -304,4 +386,40 @@ describe("production prompt injection routes", () => {
// then
expect(offenders).toEqual([])
})
test("#given production TypeScript sources #when prompt gate callers are audited #then every route declares queue behavior explicitly", async () => {
// given
const files = await listSourceFiles(SOURCE_ROOT)
const offenders: string[] = []
// when
for (const filePath of files) {
const contents = await readFile(filePath, "utf8")
const missingLines = findPromptGateCallsWithoutQueueBehavior(filePath, contents)
for (const line of missingLines) {
offenders.push(`${relativeSourcePath(filePath)}:${line}`)
}
}
// then
expect(offenders).toEqual([])
})
test("#given production TypeScript sources #when model-suggestion prompt wrappers are audited #then every retry caller declares queue behavior explicitly", async () => {
// given
const files = await listSourceFiles(SOURCE_ROOT)
const offenders: string[] = []
// when
for (const filePath of files) {
const contents = await readFile(filePath, "utf8")
const missingLines = findPromptRetryCallsWithoutQueueBehavior(filePath, contents)
for (const line of missingLines) {
offenders.push(`${relativeSourcePath(filePath)}:${line}`)
}
}
// then
expect(offenders).toEqual([])
})
})
src/shared/session-route.test.ts
+47 -5
View File
@@ -1,9 +1,14 @@
import { describe, expect, mock, test } from "bun:test"
import { afterEach, describe, expect, mock, test } from "bun:test"
import { unsafeTestValue } from "../../test-support/unsafe-test-value"
import { promptAsyncInDirectory } from "./session-route"
import { releaseAllPromptAsyncReservationsForTesting } from "./prompt-async-gate"
import { promptAsyncInDirectory, promptWithRetryInDirectory } from "./session-route"
describe("promptAsyncInDirectory", () => {
afterEach(() => {
releaseAllPromptAsyncReservationsForTesting()
})
test("#given no session id is present #when routing a promptAsync request #then the helper rejects instead of using an ungated raw prompt", async () => {
// given
const promptAsync = mock(async () => ({ data: "sent" }))
@@ -27,7 +32,7 @@ describe("promptAsyncInDirectory", () => {
expect(promptAsync).toHaveBeenCalledTimes(0)
})
test("#given a routed prompt just dispatched #when the same session is prompted again immediately #then the route coalesces the duplicate", async () => {
test("#given a routed prompt just dispatched #when the same session is prompted again immediately #then the route defers the duplicate", async () => {
// given
const promptAsync = mock(async () => ({ data: "sent" }))
const client = {
@@ -46,7 +51,7 @@ describe("promptAsyncInDirectory", () => {
unsafeTestValue(args),
"/workspace/project",
)
const second = await promptAsyncInDirectory(
const second = promptAsyncInDirectory(
unsafeTestValue(client),
unsafeTestValue(args),
"/workspace/project",
@@ -54,7 +59,44 @@ describe("promptAsyncInDirectory", () => {
// then
expect(first).toEqual({ data: "sent" })
expect(second).toBeUndefined()
await expect(second).rejects.toThrow("promptAsync skipped by gate: reserved")
expect(promptAsync).toHaveBeenCalledTimes(1)
expect(promptAsync.mock.calls[0]?.[0].query).toEqual({ directory: "/workspace/project" })
})
})
describe("promptWithRetryInDirectory", () => {
afterEach(() => {
releaseAllPromptAsyncReservationsForTesting()
})
test("#given a routed retry prompt just dispatched #when the same session is prompted again immediately #then the wrapper defers instead of enqueueing", async () => {
// given
const promptAsync = mock(async () => undefined)
const client = {
session: {
promptAsync,
},
}
const args = {
path: { id: "ses_retry_route_hold" },
body: { parts: [{ type: "text", text: "continue" }] },
}
// when
await promptWithRetryInDirectory(
unsafeTestValue(client),
unsafeTestValue(args),
"/workspace/project",
)
const second = promptWithRetryInDirectory(
unsafeTestValue(client),
unsafeTestValue(args),
"/workspace/project",
)
// then
await expect(second).rejects.toThrow("promptAsync skipped by gate: reserved")
expect(promptAsync).toHaveBeenCalledTimes(1)
expect(promptAsync.mock.calls[0]?.[0].query).toEqual({ directory: "/workspace/project" })
})
src/shared/session-route.ts
+3 -2
View File
@@ -66,6 +66,7 @@ export function promptAsyncInDirectory(
input: routedArgs,
source: "session-route",
settleMs: 0,
queueBehavior: "defer",
}).then((result) => {
if (result.status === "failed") {
throw result.error
@@ -82,7 +83,7 @@ export function promptWithRetryInDirectory(
args: PromptRetryArgs,
directory: string,
): Promise<void> {
return promptWithModelSuggestionRetry(client, routePromptRetry(args, directory))
return promptWithModelSuggestionRetry(client, routePromptRetry(args, directory), { queueBehavior: "defer" })
}
export function promptSyncWithRetryInDirectory(
@@ -90,7 +91,7 @@ export function promptSyncWithRetryInDirectory(
args: PromptSyncRetryArgs,
directory: string,
): Promise<void> {
return promptSyncWithModelSuggestionRetry(client, routePromptSyncRetry(args, directory))
return promptSyncWithModelSuggestionRetry(client, routePromptSyncRetry(args, directory), { queueBehavior: "defer" })
}
export function messagesInDirectory(