From 9fb19154accf6b6f7abfbf7a9820ccda517bc87e Mon Sep 17 00:00:00 2001 From: Luyu Cheng <2239547+chengluyu@users.noreply.github.com> Date: Thu, 9 Jul 2026 15:54:50 +0800 Subject: [PATCH] fix: keep prompt goals running until terminal (#1516) * fix: keep prompt goals running until terminal * fix: reject invalid prompt goal commands * fix: ignore stale prompt goal status checks --- .changeset/fix-prompt-goal-mode.md | 5 + apps/kimi-code/src/cli/goal-prompt.ts | 7 +- apps/kimi-code/src/cli/run-prompt.ts | 70 +++++++++--- apps/kimi-code/test/cli/goal-prompt.test.ts | 116 ++++++++++++++++++++ 4 files changed, 182 insertions(+), 16 deletions(-) create mode 100644 .changeset/fix-prompt-goal-mode.md diff --git a/.changeset/fix-prompt-goal-mode.md b/.changeset/fix-prompt-goal-mode.md new file mode 100644 index 000000000..0aa7146f1 --- /dev/null +++ b/.changeset/fix-prompt-goal-mode.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fix prompt-mode goals so they run until completion and report invalid goal commands before sending prompts. diff --git a/apps/kimi-code/src/cli/goal-prompt.ts b/apps/kimi-code/src/cli/goal-prompt.ts index 5ab0cfe85..57f223ad4 100644 --- a/apps/kimi-code/src/cli/goal-prompt.ts +++ b/apps/kimi-code/src/cli/goal-prompt.ts @@ -46,13 +46,18 @@ const GOAL_PREFIX = /^\/goal(\s|$)/; * Parses a headless prompt into a goal-create request, or `undefined` when the * prompt is not a `/goal` create command (so the caller runs it as a normal * prompt). Non-create goal subcommands are not supported headless and fall - * through to normal prompt handling. + * through to normal prompt handling. Malformed create commands throw instead of + * falling through, so validation errors are reported before anything is sent to + * the model. */ export function parseHeadlessGoalCreate(prompt: string): HeadlessGoalCreate | undefined { const trimmed = prompt.trim(); if (!GOAL_PREFIX.test(trimmed)) return undefined; const args = trimmed.replace(/^\/goal/, '').trim(); const parsed = parseGoalCommand(args); + if (parsed.kind === 'error') { + throw new Error(parsed.message); + } if (parsed.kind !== 'create') return undefined; return { objective: parsed.objective, replace: parsed.replace }; } diff --git a/apps/kimi-code/src/cli/run-prompt.ts b/apps/kimi-code/src/cli/run-prompt.ts index 100368203..3bb3a705f 100644 --- a/apps/kimi-code/src/cli/run-prompt.ts +++ b/apps/kimi-code/src/cli/run-prompt.ts @@ -234,7 +234,7 @@ async function runHeadlessGoal( try { // The objective is sent as the normal prompt; goal continuation keeps the // turn alive until a terminal state is reached. - await runPromptTurn(session, goal.objective, outputFormat, stdout, stderr); + await runPromptTurn(session, goal.objective, outputFormat, stdout, stderr, true); } finally { unsubscribeGoalEvents(); const snapshot = completedSnapshot ?? (await session.getGoal()).goal; @@ -432,9 +432,11 @@ function runPromptTurn( outputFormat: PromptOutputFormat, stdout: PromptOutput, stderr: PromptOutput, + waitForGoalTerminal = false, ): Promise { let activeTurnId: number | undefined; let activeAgentId: string | undefined; + let latestStartedTurnId: number | undefined; const outputWriter = outputFormat === 'stream-json' ? new PromptJsonWriter(stdout) @@ -469,6 +471,18 @@ function runPromptTurn( } activeTurnId = event.turnId; activeAgentId = event.agentId; + latestStartedTurnId = event.turnId; + return; + } + if ( + waitForGoalTerminal && + event.type === 'goal.updated' && + event.agentId === PROMPT_MAIN_AGENT_ID && + activeTurnId === undefined && + event.snapshot !== null && + event.snapshot.status !== 'active' + ) { + void finishCompletedTurn(); return; } if ( @@ -515,19 +529,29 @@ function runPromptTurn( return; case 'turn.ended': if (event.reason === 'completed') { - void (async () => { - // Flush the buffered assistant message before draining background - // tasks: in stream-json mode the final message is only emitted by - // finish(), so a long background wait would otherwise withhold the - // main turn's result until the drain settles. - outputWriter.flushAssistant(); - try { - await session.waitForBackgroundTasksOnPrint(); - } catch (error) { - log.warn('waitForBackgroundTasksOnPrint failed', { error }); - } - finish(); - })(); + outputWriter.flushAssistant(); + if (waitForGoalTerminal) { + const completedTurnId = event.turnId; + activeTurnId = undefined; + activeAgentId = undefined; + void (async () => { + try { + const { goal } = await session.getGoal(); + if ( + activeTurnId !== undefined || + latestStartedTurnId !== completedTurnId + ) { + return; + } + if (goal?.status === 'active') return; + await finishCompletedTurn(); + } catch (error) { + finish(error instanceof Error ? error : new Error(String(error))); + } + })(); + return; + } + void finishCompletedTurn(); return; } finish(new Error(formatTurnEndedFailure(event))); @@ -560,6 +584,20 @@ function runPromptTurn( session.prompt(prompt).catch((error: unknown) => { finish(error instanceof Error ? error : new Error(String(error))); }); + + async function finishCompletedTurn(): Promise { + // Flush the buffered assistant message before draining background tasks: + // in stream-json mode the final message is only emitted by finish(), so a + // long background wait would otherwise withhold the main turn's result + // until the drain settles. + outputWriter.flushAssistant(); + try { + await session.waitForBackgroundTasksOnPrint(); + } catch (error) { + log.warn('waitForBackgroundTasksOnPrint failed', { error }); + } + finish(); + } }); } @@ -610,7 +648,9 @@ class PromptTranscriptWriter implements PromptTurnWriter { writeToolResult(): void {} - flushAssistant(): void {} + flushAssistant(): void { + this.assistantWriter.finish(); + } discardAssistant(): void {} diff --git a/apps/kimi-code/test/cli/goal-prompt.test.ts b/apps/kimi-code/test/cli/goal-prompt.test.ts index 04780bd26..8d75c4721 100644 --- a/apps/kimi-code/test/cli/goal-prompt.test.ts +++ b/apps/kimi-code/test/cli/goal-prompt.test.ts @@ -46,6 +46,12 @@ describe('parseHeadlessGoalCreate', () => { expect(parseHeadlessGoalCreate('/goal status')).toBeUndefined(); expect(parseHeadlessGoalCreate('/goal pause')).toBeUndefined(); }); + + it('rejects malformed goal create prompts instead of falling through', () => { + expect(() => parseHeadlessGoalCreate(`/goal ${'x'.repeat(4001)}`)).toThrow( + 'Goal objective is too long', + ); + }); }); describe('goal summary', () => { @@ -97,6 +103,7 @@ const mocks = vi.hoisted(() => { handler(mainEvent({ type: 'turn.ended', turnId: 1, reason: 'completed' })); } }), + waitForBackgroundTasksOnPrint: vi.fn(async () => {}), }; return { session, @@ -164,6 +171,8 @@ describe('runPrompt headless goal mode', () => { mocks.experimentalFeatures = [{ id: 'micro_compaction', enabled: true }]; mocks.sessions = []; mocks.session.createGoal.mockClear(); + mocks.session.prompt.mockClear(); + mocks.session.waitForBackgroundTasksOnPrint.mockClear(); mocks.session.getStatus.mockResolvedValue({ permission: 'auto', model: 'k2' } as never); mocks.session.getGoal.mockResolvedValue({ goal: snapshot({ status: 'complete' }) } as never); }); @@ -243,6 +252,113 @@ describe('runPrompt headless goal mode', () => { expect(mocks.session.prompt).toHaveBeenCalledWith('Ship feature X'); }); + it('keeps listening across continuation turns until the goal is terminal', async () => { + const active = snapshot({ status: 'active', turnsUsed: 1, tokensUsed: 80 }); + const completed = snapshot({ status: 'complete', turnsUsed: 2, tokensUsed: 160 }); + mocks.session.getGoal.mockResolvedValueOnce({ goal: active } as never); + mocks.session.prompt.mockImplementationOnce(async () => { + for (const handler of mocks.eventHandlers) { + handler(mocks.mainEvent({ type: 'turn.started', turnId: 1, origin: { kind: 'user' } })); + handler(mocks.mainEvent({ type: 'assistant.delta', turnId: 1, delta: '1' })); + handler(mocks.mainEvent({ type: 'turn.ended', turnId: 1, reason: 'completed' })); + } + await Promise.resolve(); + for (const handler of mocks.eventHandlers) { + handler( + mocks.mainEvent({ + type: 'turn.started', + turnId: 2, + origin: { kind: 'system_trigger', name: 'goal_continuation' }, + }), + ); + handler(mocks.mainEvent({ type: 'assistant.delta', turnId: 2, delta: '2' })); + handler( + mocks.mainEvent({ + type: 'goal.updated', + snapshot: completed, + change: { kind: 'completion', status: 'complete' }, + }), + ); + handler(mocks.mainEvent({ type: 'turn.ended', turnId: 2, reason: 'completed' })); + } + }); + const stdout = writer(); + const stderr = writer(); + + await runPrompt(opts(), 'test', { + stdout, + stderr, + process: { once: () => {}, off: () => {}, exit: () => undefined as never }, + }); + + expect(stdout.text()).toBe('• 1\n\n• 2\n\n'); + expect(stderr.text()).toContain('Goal [complete]'); + expect(stderr.text()).toContain('turns: 2'); + }); + + it('ignores stale goal checks once a continuation turn has started', async () => { + const completed = snapshot({ status: 'complete', turnsUsed: 2, tokensUsed: 160 }); + let resolveFirstGoal: ((value: { goal: null }) => void) | undefined; + const firstGoal = new Promise<{ goal: null }>((resolve) => { + resolveFirstGoal = resolve; + }); + mocks.session.getGoal + .mockImplementationOnce(() => firstGoal as never) + .mockResolvedValue({ goal: null } as never); + mocks.session.prompt.mockImplementationOnce(async () => { + const emit = (event: Record) => { + for (const handler of [...mocks.eventHandlers]) { + handler(mocks.mainEvent(event)); + } + }; + emit({ type: 'turn.started', turnId: 1, origin: { kind: 'user' } }); + emit({ type: 'assistant.delta', turnId: 1, delta: '1' }); + emit({ type: 'turn.ended', turnId: 1, reason: 'completed' }); + emit({ + type: 'turn.started', + turnId: 2, + origin: { kind: 'system_trigger', name: 'goal_continuation' }, + }); + emit({ type: 'assistant.delta', turnId: 2, delta: '2' }); + emit({ + type: 'goal.updated', + snapshot: completed, + change: { kind: 'completion', status: 'complete' }, + }); + resolveFirstGoal?.({ goal: null }); + await Promise.resolve(); + emit({ type: 'assistant.delta', turnId: 2, delta: ' tail' }); + emit({ type: 'turn.ended', turnId: 2, reason: 'completed' }); + }); + const stdout = writer(); + const stderr = writer(); + + await runPrompt(opts(), 'test', { + stdout, + stderr, + process: { once: () => {}, off: () => {}, exit: () => undefined as never }, + }); + + expect(stdout.text()).toBe('• 1\n\n• 2 tail\n\n'); + expect(stderr.text()).toContain('Goal [complete]'); + }); + + it('does not send an invalid goal create prompt as a normal prompt', async () => { + const stdout = writer(); + const stderr = writer(); + + await expect( + runPrompt(opts({ prompt: `/goal ${'x'.repeat(4001)}` }), 'test', { + stdout, + stderr, + process: { once: () => {}, off: () => {}, exit: () => undefined as never }, + }), + ).rejects.toThrow('Goal objective is too long'); + + expect(mocks.session.createGoal).not.toHaveBeenCalled(); + expect(mocks.session.prompt).not.toHaveBeenCalled(); + }); + it('validates the resumed session model before creating a headless goal', async () => { mocks.sessions = [{ id: 'ses_goal', workDir: process.cwd() }]; mocks.session.getStatus.mockResolvedValueOnce({ permission: 'auto', model: '' } as never);