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

@@ -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.

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 };
 }

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<void> {
   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<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();
+    }
   });
 }
 
@@ -610,7 +648,9 @@ class PromptTranscriptWriter implements PromptTurnWriter {
 
   writeToolResult(): void {}
 
-  flushAssistant(): void {}
+  flushAssistant(): void {
+    this.assistantWriter.finish();
+  }
 
   discardAssistant(): void {}
 

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<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);