diff --git a/.changeset/flush-stream-json-output.md b/.changeset/flush-stream-json-output.md new file mode 100644 index 0000000000..5f1475a88c --- /dev/null +++ b/.changeset/flush-stream-json-output.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Respect stream JSON backpressure without allowing a stalled consumer to block cleanup or exit indefinitely. diff --git a/apps/kimi-code/src/cli/prompt-render.ts b/apps/kimi-code/src/cli/prompt-render.ts index 0e2f35238d..4e7ff21c81 100644 --- a/apps/kimi-code/src/cli/prompt-render.ts +++ b/apps/kimi-code/src/cli/prompt-render.ts @@ -43,6 +43,41 @@ interface RetryingEventLike { export interface PromptOutput { readonly columns?: number | undefined; write(chunk: string): boolean; + once?(event: 'drain', listener: () => void): unknown; +} + +interface PromptOutputDrainState { + pending: Promise | undefined; +} + +const outputDrainStates = new WeakMap(); + +export function writePromptOutput(output: PromptOutput, chunk: string): void { + if (output.write(chunk) || output.once === undefined) return; + + let state = outputDrainStates.get(output); + if (state === undefined) { + state = { pending: undefined }; + outputDrainStates.set(output, state); + } + if (state.pending !== undefined) return; + + let resolveDrain!: () => void; + state.pending = new Promise((resolve) => { + resolveDrain = resolve; + }); + output.once('drain', () => { + state.pending = undefined; + resolveDrain(); + }); +} + +export async function drainPromptOutput(output: PromptOutput): Promise { + while (true) { + const pending = outputDrainStates.get(output)?.pending; + if (pending === undefined) return; + await pending; + } } const PROMPT_BLOCK_BULLET = '• '; @@ -262,7 +297,7 @@ export class PromptJsonWriter implements PromptTurnWriter { private writeJsonLine( message: PromptJsonAssistantMessage | PromptJsonToolMessage | PromptJsonRetryMetaMessage, ): void { - this.stdout.write(`${JSON.stringify(message)}\n`); + writePromptOutput(this.stdout, `${JSON.stringify(message)}\n`); } } @@ -381,7 +416,7 @@ export function writeExperimentalVersion( type: 'system.version', version, }; - stdout.write(`${JSON.stringify(message)}\n`); + writePromptOutput(stdout, `${JSON.stringify(message)}\n`); return; } stderr.write(`kimi version ${version}\n`); @@ -403,7 +438,7 @@ export function writeResumeHint( command, content, }; - stdout.write(`${JSON.stringify(message)}\n`); + writePromptOutput(stdout, `${JSON.stringify(message)}\n`); return; } stderr.write(`${content}\n`); diff --git a/apps/kimi-code/src/cli/run-prompt.ts b/apps/kimi-code/src/cli/run-prompt.ts index abee299622..9d1f874da8 100644 --- a/apps/kimi-code/src/cli/run-prompt.ts +++ b/apps/kimi-code/src/cli/run-prompt.ts @@ -16,7 +16,11 @@ import { } from '@moonshot-ai/kimi-code-sdk'; import { resolve } from 'pathe'; -import { CLI_SHUTDOWN_TIMEOUT_MS, PROMPT_CLEANUP_TIMEOUT_MS } from '#/constant/app'; +import { + CLI_SHUTDOWN_TIMEOUT_MS, + HEADLESS_STDIO_DRAIN_TIMEOUT_MS, + PROMPT_CLEANUP_TIMEOUT_MS, +} from '#/constant/app'; import { isKimiV2Enabled } from './experimental-v2'; import { resolveOutputFormat } from './options'; @@ -29,7 +33,13 @@ import { type HeadlessGoalCreate, } from './goal-prompt'; import type { PromptHarness, PromptSession } from './prompt-session'; -import { PromptJsonWriter, PromptTranscriptWriter, writeResumeHint } from './prompt-render'; +import { + drainPromptOutput, + PromptJsonWriter, + PromptTranscriptWriter, + writePromptOutput, + writeResumeHint, +} from './prompt-render'; import { createCliTelemetryBootstrap, initializeCliTelemetry } from './telemetry'; import { createKimiCodeHostIdentity } from './version'; @@ -146,6 +156,7 @@ export async function runPrompt( let restorePromptSessionPermission = async (): Promise => {}; let removeTerminationCleanup: (() => void) | undefined; let cleanupPromise: Promise | undefined; + let outputDrainAttempted = false; const cleanupPromptRun = async (): Promise => { const pending = (cleanupPromise ??= (async () => { removeTerminationCleanup?.(); @@ -153,8 +164,12 @@ export async function runPrompt( try { await restorePromptSessionPermission(); } finally { - await shutdownTelemetry({ timeoutMs: CLI_SHUTDOWN_TIMEOUT_MS }); - await harness.close(); + try { + await shutdownTelemetry({ timeoutMs: CLI_SHUTDOWN_TIMEOUT_MS }); + await harness.close(); + } finally { + if (!outputDrainAttempted) await drainPromptOutput(stdout); + } } })()); // Bound cleanup so a wedged shutdown step (e.g. a SessionEnd hook, MCP @@ -213,6 +228,8 @@ export async function runPrompt( ); } writeResumeHint(session.id, outputFormat, stdout, stderr); + outputDrainAttempted = true; + await raceWithTimeout(drainPromptOutput(stdout), HEADLESS_STDIO_DRAIN_TIMEOUT_MS); withTelemetryContext({ sessionId: session.id }).track('exit', { duration_ms: Date.now() - startedAt, @@ -268,7 +285,7 @@ async function runHeadlessGoal( unsubscribeGoalEvents(); const snapshot = completedSnapshot ?? (await session.getGoal()).goal; if (outputFormat === 'stream-json') { - stdout.write(`${JSON.stringify(goalSummaryJson(snapshot))}\n`); + writePromptOutput(stdout, `${JSON.stringify(goalSummaryJson(snapshot))}\n`); } else { stderr.write(`${formatGoalSummaryText(snapshot)}\n`); } diff --git a/apps/kimi-code/src/cli/v2/run-v2-print.ts b/apps/kimi-code/src/cli/v2/run-v2-print.ts index ff0c96c590..a41a1efcce 100644 --- a/apps/kimi-code/src/cli/v2/run-v2-print.ts +++ b/apps/kimi-code/src/cli/v2/run-v2-print.ts @@ -63,6 +63,7 @@ import { resolve } from 'pathe'; import { CLI_SHUTDOWN_TIMEOUT_MS, CLI_USER_AGENT_PRODUCT, + HEADLESS_STDIO_DRAIN_TIMEOUT_MS, PROMPT_CLEANUP_TIMEOUT_MS, } from '#/constant/app'; @@ -85,11 +86,13 @@ import { createKimiCodeHostIdentity } from '../version'; import { resolveOutputFormat } from '../options'; import type { CLIOptions, PromptOutputFormat } from '../options'; import { + drainPromptOutput, type PromptOutput, PromptJsonWriter, type PromptTurnWriter, PromptTranscriptWriter, writeExperimentalVersion, + writePromptOutput, writeResumeHint, } from '../prompt-render'; @@ -165,16 +168,21 @@ export async function runV2Print( let removeTerminationCleanup: (() => void) | undefined; let cleanupPromise: Promise | undefined; let telemetryService: ITelemetryService | undefined; + let outputDrainAttempted = false; const cleanup = async (): Promise => { const pending = (cleanupPromise ??= (async () => { removeTerminationCleanup?.(); try { await restorePermission(); } finally { - if (telemetryService !== undefined) { - await raceWithTimeout(telemetryService.shutdown(), CLI_SHUTDOWN_TIMEOUT_MS); + try { + if (telemetryService !== undefined) { + await raceWithTimeout(telemetryService.shutdown(), CLI_SHUTDOWN_TIMEOUT_MS); + } + app.dispose(); + } finally { + if (!outputDrainAttempted) await drainPromptOutput(stdout); } - app.dispose(); } })()); await raceWithTimeout(pending, PROMPT_CLEANUP_TIMEOUT_MS); @@ -232,6 +240,8 @@ export async function runV2Print( ); } writeResumeHint(resolved.session.id, outputFormat, stdout, stderr); + outputDrainAttempted = true; + await raceWithTimeout(drainPromptOutput(stdout), HEADLESS_STDIO_DRAIN_TIMEOUT_MS); telemetryService.withContext({ sessionId: resolved.session.id }).track2('exit', { duration_ms: Date.now() - startedAt, @@ -535,7 +545,7 @@ async function runNativeGoal( subscription.dispose(); const snapshot = completedSnapshot ?? goalService.getGoal().goal; if (outputFormat === 'stream-json') { - stdout.write(`${JSON.stringify(goalSummaryJson(snapshot))}\n`); + writePromptOutput(stdout, `${JSON.stringify(goalSummaryJson(snapshot))}\n`); } else { stderr.write(`${formatGoalSummaryText(snapshot)}\n`); } diff --git a/apps/kimi-code/test/cli/run-prompt.test.ts b/apps/kimi-code/test/cli/run-prompt.test.ts index 4bad127d89..759f9f2746 100644 --- a/apps/kimi-code/test/cli/run-prompt.test.ts +++ b/apps/kimi-code/test/cli/run-prompt.test.ts @@ -2,7 +2,10 @@ import type { createKimiDeviceId as createKimiDeviceIdFn } from '@moonshot-ai/ki import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { runPrompt } from '#/cli/run-prompt'; -import { PROMPT_CLEANUP_TIMEOUT_MS } from '#/constant/app'; +import { + HEADLESS_STDIO_DRAIN_TIMEOUT_MS, + PROMPT_CLEANUP_TIMEOUT_MS, +} from '#/constant/app'; type CreateKimiDeviceId = typeof createKimiDeviceIdFn; @@ -203,6 +206,27 @@ function writer(columns?: number) { }; } +function backpressuredWriter() { + let text = ''; + let blocked = true; + const drainListeners = new Set<() => void>(); + return { + write: vi.fn((chunk: string) => { + text += chunk; + return !blocked; + }), + once: vi.fn((event: 'drain', listener: () => void) => { + if (event === 'drain') drainListeners.add(listener); + }), + drain: () => { + blocked = false; + for (const listener of drainListeners) listener(); + drainListeners.clear(); + }, + text: () => text, + }; +} + function fakeProcess() { const listeners = new Map Promise | void>(); return { @@ -680,6 +704,56 @@ describe('runPrompt', () => { expect(stderr.text()).toBe(''); }); + it('waits for stream-json stdout backpressure before completing', async () => { + const stdout = backpressuredWriter(); + const stderr = writer(); + let settled = false; + + const run = runPrompt(opts({ outputFormat: 'stream-json' }), '1.2.3-test', { + stdout, + stderr, + }).then(() => { + settled = true; + }); + + await waitForAssertion(() => { + expect(stdout.text()).toContain('"type":"session.resume_hint"'); + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(settled).toBe(false); + expect(stdout.once).toHaveBeenCalledWith('drain', expect.any(Function)); + + stdout.drain(); + await run; + expect(settled).toBe(true); + }); + + it('completes after the stdio timeout if stream-json stdout stays backpressured', async () => { + vi.useFakeTimers(); + try { + const stdout = backpressuredWriter(); + let settled = false; + const run = runPrompt(opts({ outputFormat: 'stream-json' }), '1.2.3-test', { + stdout, + stderr: writer(), + }).then(() => { + settled = true; + }); + + await vi.advanceTimersByTimeAsync(HEADLESS_STDIO_DRAIN_TIMEOUT_MS - 1); + expect(settled).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + await run; + expect(settled).toBe(true); + expect(mocks.shutdownTelemetry).toHaveBeenCalled(); + expect(mocks.harnessClose).toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + it('writes stream-json tool calls and tool results as JSONL messages', async () => { mocks.session.prompt.mockImplementationOnce(async () => { for (const handler of mocks.eventHandlers) { @@ -1033,6 +1107,116 @@ describe('runPrompt', () => { expect(mocks.harnessClose).toHaveBeenCalledTimes(1); }); + it('drains pending stream-json output before exiting on a signal', async () => { + let releasePrompt!: () => void; + mocks.session.prompt.mockImplementationOnce(async () => { + for (const handler of mocks.eventHandlers) { + handler( + mocks.mainEvent({ type: 'turn.started', turnId: 12, origin: { kind: 'user' } }), + ); + handler(mocks.mainEvent({ type: 'assistant.delta', turnId: 12, delta: 'final answer' })); + handler( + mocks.mainEvent({ + type: 'tool.result', + turnId: 12, + toolCallId: 'tc_done', + output: 'done', + }), + ); + } + await new Promise((resolve) => { + releasePrompt = resolve; + }); + }); + const stdout = backpressuredWriter(); + const processMock = fakeProcess(); + const run = runPrompt(opts({ outputFormat: 'stream-json' }), '1.2.3-test', { + stdout, + stderr: writer(), + process: processMock, + } as Parameters[2] & { process: ReturnType }); + + await waitForAssertion(() => { + expect(stdout.text()).toContain('{"role":"assistant","content":"final answer"}'); + expect(processMock.listener('SIGTERM')).toBeDefined(); + }); + + const signalCleanup = processMock.listener('SIGTERM')?.(); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(processMock.exit).not.toHaveBeenCalled(); + + stdout.drain(); + await signalCleanup; + expect(processMock.exit).toHaveBeenCalledWith(143); + + for (const handler of mocks.eventHandlers) { + handler(mocks.mainEvent({ type: 'turn.ended', turnId: 12, reason: 'completed' })); + } + releasePrompt(); + await run; + }); + + it('restores permission before timing out a wedged signal output drain', async () => { + try { + let releasePrompt!: () => void; + mocks.session.prompt.mockImplementationOnce(async () => { + for (const handler of mocks.eventHandlers) { + handler( + mocks.mainEvent({ type: 'turn.started', turnId: 13, origin: { kind: 'user' } }), + ); + handler( + mocks.mainEvent({ type: 'assistant.delta', turnId: 13, delta: 'final answer' }), + ); + handler( + mocks.mainEvent({ + type: 'tool.result', + turnId: 13, + toolCallId: 'tc_done', + output: 'done', + }), + ); + } + await new Promise((resolve) => { + releasePrompt = resolve; + }); + }); + const stdout = backpressuredWriter(); + const processMock = fakeProcess(); + const run = runPrompt( + opts({ session: 'ses_existing', outputFormat: 'stream-json' }), + '1.2.3-test', + { stdout, stderr: writer(), process: processMock } as Parameters< + typeof runPrompt + >[2] & { process: ReturnType }, + ); + + await waitForAssertion(() => { + expect(stdout.text()).toContain('final answer'); + expect(processMock.listener('SIGTERM')).toBeDefined(); + }); + vi.useFakeTimers(); + const signalCleanup = processMock.listener('SIGTERM')?.(); + await vi.advanceTimersByTimeAsync(0); + + expect(mocks.session.setPermission).toHaveBeenNthCalledWith(2, 'manual'); + expect(mocks.harnessClose).toHaveBeenCalled(); + expect(processMock.exit).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(PROMPT_CLEANUP_TIMEOUT_MS); + await signalCleanup; + expect(processMock.exit).toHaveBeenCalledWith(143); + + stdout.drain(); + for (const handler of mocks.eventHandlers) { + handler(mocks.mainEvent({ type: 'turn.ended', turnId: 13, reason: 'completed' })); + } + releasePrompt(); + await run; + } finally { + vi.useRealTimers(); + } + }); + it('waits for the pending auto permission write before signal restore', async () => { let releaseAutoPermission!: () => void; let releasePrompt!: () => void; diff --git a/apps/kimi-code/test/cli/v2-run-print.test.ts b/apps/kimi-code/test/cli/v2-run-print.test.ts index e6a5e5ce3e..c3bf79a098 100644 --- a/apps/kimi-code/test/cli/v2-run-print.test.ts +++ b/apps/kimi-code/test/cli/v2-run-print.test.ts @@ -4,6 +4,11 @@ import { join } from 'node:path'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { + HEADLESS_STDIO_DRAIN_TIMEOUT_MS, + PROMPT_CLEANUP_TIMEOUT_MS, +} from '#/constant/app'; + import { IAgentCatalogRuntimeOptions, IAgentGoalService, @@ -104,6 +109,39 @@ function writer() { }; } +function backpressuredWriter() { + let text = ''; + let blocked = true; + const drainListeners = new Set<() => void>(); + return { + write: vi.fn((chunk: string) => { + text += chunk; + return !blocked; + }), + once: vi.fn((_event: 'drain', listener: () => void) => drainListeners.add(listener)), + drain: () => { + blocked = false; + for (const listener of drainListeners) listener(); + drainListeners.clear(); + }, + text: () => text, + }; +} + +function fakeProcess() { + const listeners = new Map Promise | void>(); + return { + once: vi.fn((signal: NodeJS.Signals, listener: () => Promise | void) => { + listeners.set(signal, listener); + }), + off: vi.fn((signal: NodeJS.Signals, listener: () => Promise | void) => { + if (listeners.get(signal) === listener) listeners.delete(signal); + }), + exit: vi.fn(), + listener: (signal: NodeJS.Signals) => listeners.get(signal), + }; +} + function opts(overrides: Record = {}) { return { session: undefined, @@ -267,6 +305,119 @@ describe('runV2Print', () => { expect(app.dispose).toHaveBeenCalled(); }); + it('waits for stream-json stdout backpressure before completing', async () => { + const stdout = backpressuredWriter(); + const { app, agent } = makeFakeHarness(); + mocks.bootstrap.mockReturnValue({ app }); + mocks.ensureMainAgent.mockResolvedValue(agent); + let settled = false; + + const run = runV2Print(opts({ outputFormat: 'stream-json' }) as never, '1.2.3-test', { + stdout, + stderr: writer(), + }).then(() => { + settled = true; + }); + + for (let attempt = 0; attempt < 20 && !stdout.text().includes('session.resume_hint'); attempt += 1) { + await new Promise((resolve) => setTimeout(resolve, 0)); + } + expect(stdout.text()).toContain('session.resume_hint'); + expect(settled).toBe(false); + + stdout.drain(); + await run; + expect(settled).toBe(true); + }); + + it('completes cleanup after the stdio timeout if stdout never drains', async () => { + vi.useFakeTimers(); + try { + const stdout = backpressuredWriter(); + const { app, agent } = makeFakeHarness(); + mocks.bootstrap.mockReturnValue({ app }); + mocks.ensureMainAgent.mockResolvedValue(agent); + let settled = false; + const run = runV2Print( + opts({ outputFormat: 'stream-json' }) as never, + '1.2.3-test', + { stdout, stderr: writer() }, + ).then(() => { + settled = true; + }); + + await vi.advanceTimersByTimeAsync(HEADLESS_STDIO_DRAIN_TIMEOUT_MS - 1); + expect(settled).toBe(false); + + await vi.advanceTimersByTimeAsync(1); + await run; + expect(settled).toBe(true); + expect(app.dispose).toHaveBeenCalled(); + } finally { + vi.useRealTimers(); + } + }); + + it('restores permission before timing out a wedged signal output drain', async () => { + let releaseTurn!: () => void; + const stdout = backpressuredWriter(); + const processMock = fakeProcess(); + const { app, agent, agentServices, appServices } = makeFakeHarness(); + const permission = agentServices.get(IAgentPermissionModeService) as { + mode: string; + setMode: ReturnType; + }; + permission.mode = 'manual'; + const promptService = agentServices.get(IAgentPromptService) as { + enqueue: ReturnType; + }; + promptService.enqueue.mockResolvedValueOnce({ + launched: Promise.resolve({ + id: 1, + result: new Promise((resolve) => { + releaseTurn = () => { + resolve({ type: 'completed' }); + }; + }), + }), + }); + const index = appServices.get(ISessionIndex) as { list: ReturnType }; + index.list.mockResolvedValueOnce({ items: [{ id: 'ses_v2', cwd: process.cwd() }] }); + mocks.bootstrap.mockReturnValue({ app }); + mocks.ensureMainAgent.mockResolvedValue(agent); + + const run = runV2Print( + opts({ session: 'ses_v2', outputFormat: 'stream-json' }) as never, + '1.2.3-test', + { stdout, stderr: writer(), process: processMock }, + ); + for (let attempt = 0; attempt < 20 && permission.setMode.mock.calls.length === 0; attempt += 1) { + await new Promise((resolve) => setTimeout(resolve, 0)); + } + expect(permission.setMode).toHaveBeenNthCalledWith(1, 'auto'); + expect(stdout.text()).toContain('system.version'); + + vi.useFakeTimers(); + try { + const signalCleanup = processMock.listener('SIGTERM')?.(); + await vi.advanceTimersByTimeAsync(0); + + expect(permission.setMode).toHaveBeenNthCalledWith(2, 'manual'); + expect(app.dispose).toHaveBeenCalled(); + expect(processMock.exit).not.toHaveBeenCalled(); + + await vi.advanceTimersByTimeAsync(PROMPT_CLEANUP_TIMEOUT_MS); + await signalCleanup; + expect(processMock.exit).toHaveBeenCalledWith(143); + + stdout.drain(); + releaseTurn(); + await run; + } finally { + vi.useRealTimers(); + } + }); + it('seeds explicit skill dirs from --skillsDir into bootstrap', async () => { const stdout = writer(); const stderr = writer();