Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/fix-prompt-goal-mode.md
Original file line number Diff line number Diff line change
@@ -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.
7 changes: 6 additions & 1 deletion apps/kimi-code/src/cli/goal-prompt.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 };
}
Expand Down
70 changes: 55 additions & 15 deletions apps/kimi-code/src/cli/run-prompt.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -432,9 +432,11 @@ function runPromptTurn(
outputFormat: PromptOutputFormat,
stdout: PromptOutput,
stderr: PromptOutput,
waitForGoalTerminal = false,
): Promise<void> {
let activeTurnId: number | undefined;
let activeAgentId: string | undefined;
let latestStartedTurnId: number | undefined;
const outputWriter =
outputFormat === 'stream-json'
? new PromptJsonWriter(stdout)
Expand Down Expand Up @@ -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 (
Expand Down Expand Up @@ -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();
Comment thread
chengluyu marked this conversation as resolved.
} catch (error) {
finish(error instanceof Error ? error : new Error(String(error)));
}
})();
return;
}
void finishCompletedTurn();
return;
}
finish(new Error(formatTurnEndedFailure(event)));
Expand Down Expand Up @@ -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<void> {
// 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();
}
});
}

Expand Down Expand Up @@ -610,7 +648,9 @@ class PromptTranscriptWriter implements PromptTurnWriter {

writeToolResult(): void {}

flushAssistant(): void {}
flushAssistant(): void {
this.assistantWriter.finish();
}

discardAssistant(): void {}

Expand Down
116 changes: 116 additions & 0 deletions apps/kimi-code/test/cli/goal-prompt.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down Expand Up @@ -97,6 +103,7 @@ const mocks = vi.hoisted(() => {
handler(mainEvent({ type: 'turn.ended', turnId: 1, reason: 'completed' }));
}
}),
waitForBackgroundTasksOnPrint: vi.fn(async () => {}),
};
return {
session,
Expand Down Expand Up @@ -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);
});
Expand Down Expand Up @@ -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<string, unknown>) => {
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);
Expand Down
Loading