fix(acp): restore IteratorResult return type to nextSdkMessageOrAbort

The simplified SDKMessage | undefined return type collapsed two distinct
states: generator truly done vs generator yielding undefined. This broke
forwardSessionUpdates which needs to distinguish the two — when the
generator yields null/undefined it should continue (calling next() again),
not break out of the loop.

Restored the original IteratorResult<SDKMessage, void> return type so
done and yielded-null are distinct again.
This commit is contained in:
James Feng 2026-06-08 20:07:53 +08:00
parent 3009a58991
commit 3a2f358d27
2 changed files with 22 additions and 26 deletions

View File

@ -698,7 +698,7 @@ describe('toDisplayPath', () => {
// ── forwardSessionUpdates ─────────────────────────────────────────
describe('nextSdkMessageOrAbort', () => {
test('returns undefined when aborted while waiting for next message', async () => {
test('returns done:true when aborted while waiting for next message', async () => {
const ac = new AbortController()
const pending = nextSdkMessageOrAbort(makeWaitingStream(), ac.signal)
ac.abort()
@ -708,19 +708,19 @@ describe('nextSdkMessageOrAbort', () => {
new Promise<'timeout'>(resolve => setTimeout(resolve, 100, 'timeout')),
])
expect(result).toBeUndefined()
expect(result).toEqual({ done: true, value: undefined })
})
test('returns undefined when stream is done', async () => {
test('returns done:true when stream is done', async () => {
const result = await nextSdkMessageOrAbort(
makeStream([]),
new AbortController().signal,
)
expect(result).toBeUndefined()
expect(result).toEqual({ done: true, value: undefined })
})
test('returns a valid SDKMessage', async () => {
test('returns a valid SDKMessage via IteratorResult', async () => {
const msg = {
type: 'assistant',
message: {
@ -734,7 +734,7 @@ describe('nextSdkMessageOrAbort', () => {
new AbortController().signal,
)
expect(result).toBe(msg)
expect(result).toEqual({ done: false, value: msg })
})
})

View File

@ -719,28 +719,22 @@ export function toolUpdateFromEditToolResponse(toolResponse: unknown): {
export function nextSdkMessageOrAbort(
sdkMessages: AsyncGenerator<SDKMessage, void, unknown>,
abortSignal: AbortSignal,
): Promise<SDKMessage | undefined> {
): Promise<IteratorResult<SDKMessage, void>> {
if (abortSignal.aborted) {
return Promise.resolve(undefined)
return Promise.resolve({ done: true, value: undefined })
}
let abortHandler: (() => void) | undefined
const abortPromise = new Promise<undefined>(resolve => {
abortHandler = () => resolve(undefined)
abortSignal.addEventListener('abort', abortHandler, { once: true })
const abortPromise = new Promise<IteratorResult<SDKMessage, void>>(
resolve => {
abortHandler = () => resolve({ done: true, value: undefined })
abortSignal.addEventListener('abort', abortHandler, { once: true })
},
)
return Promise.race([sdkMessages.next(), abortPromise]).finally(() => {
if (abortHandler) {
abortSignal.removeEventListener('abort', abortHandler)
}
})
const nextPromise = sdkMessages.next()
return Promise.race([nextPromise, abortPromise])
.then(result => {
if (result === undefined) {
void nextPromise.catch(() => {})
return undefined
}
if (result.done === true) return undefined
return result.value
})
.finally(() => {
if (abortHandler) abortSignal.removeEventListener('abort', abortHandler)
})
}
// ── Main forwarding function ──────────────────────────────────────
@ -779,8 +773,10 @@ export async function forwardSessionUpdates(
// Race the next message against the abort signal so we unblock
// immediately when cancelled, even if the generator is waiting for
// a slow API response.
const rawMsg = await nextSdkMessageOrAbort(sdkMessages, abortSignal)
if (rawMsg === undefined || abortSignal.aborted) break
const nextResult = await nextSdkMessageOrAbort(sdkMessages, abortSignal)
if (nextResult.done || abortSignal.aborted) break
const rawMsg = nextResult.value
if (rawMsg == null) continue
const msg = rawMsg as BridgeSDKMessage
switch (msg.type) {