Skip to content
Open
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/steer-mid-turn-reminder.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@moonshot-ai/kimi-code": patch
---

Fix messages steered into a running turn being overlooked while the agent continued its original task.
2 changes: 1 addition & 1 deletion packages/agent-core-v2/AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ Name-collision precedence is a deliberate, documented divergence from v1: v1 ran

## Model-facing reminders

Two delivery paths only — never introduce a third (no deferred-delivery queues, no mid-step splice channels), both owned by the Agent-scope eager DI service `IAgentReminderService` — obtained through constructor injection within the same Agent scope, or through the scope handle's `accessor.get(IAgentReminderService)` across scope boundaries: reminders that restate current state (goal state, plan mode, date change, …) call `register(variant, provider)` and reconcile at every step head before the request is built, re-emitting after compaction or undo; reminders that report a one-off event (goal cancelled, AGENTS.md discovered, `/init` finished, …) call `notify(content, { variant, ownerPromptId? })` at a safe event point (a step/restore hook, an idle moment, or the loop-event fold's deferred append). The service owns `<system-reminder>` wrapping and stamps `{ kind: 'injection', variant }`; `kind: 'injection'` is a lifecycle classification (hidden from the UI, not an undo anchor, dropped by compaction), not a provenance claim, and prompt-owned attachments carry `ownerPromptId` so undo treats them as part of their host prompt.
Two delivery paths only — never introduce a third (no deferred-delivery queues, no mid-step splice channels), both owned by the Agent-scope eager DI service `IAgentReminderService` — obtained through constructor injection within the same Agent scope, or through the scope handle's `accessor.get(IAgentReminderService)` across scope boundaries: reminders that restate current state (goal state, plan mode, date change, …) call `register(variant, provider)` and reconcile at every step head before the request is built, re-emitting after compaction or undo; reminders that report a one-off event (goal cancelled, AGENTS.md discovered, `/init` finished, …) call `notify(content, { variant, ownerPromptId? })` at a safe event point (a step/restore hook, an idle moment, or the loop-event fold's deferred append). The service owns `<system-reminder>` wrapping and stamps `{ kind: 'injection', variant }`; `kind: 'injection'` is a lifecycle classification (hidden from the UI, not an undo anchor, dropped by compaction), not a provenance claim, and prompt-owned attachments carry `ownerPromptId` so undo treats them as part of their host prompt. A one-off event whose reminder must survive same-step compaction (the steer reminder) still goes through `register`, not `notify`: arm a flag at the event, emit while armed, clear when the step finishes — a `notify` append would be dropped by the splice with nothing left to re-emit it, while the armed provider restates it on the same-step reconciliation pass.

## Docs

Expand Down
21 changes: 18 additions & 3 deletions packages/agent-core-v2/src/agent/prompt/promptService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import { z } from 'zod';

import { IInstantiationService } from '#/_base/di/instantiation';
import { Disposable } from '#/_base/di/lifecycle';
import { LifecycleScope } from '#/app/scopes';
import { ScopeActivation, registerScopedService } from '#/_base/di/scope';
import { defineState } from '#/state/state';
Expand Down Expand Up @@ -219,13 +220,19 @@ function mergeSteerMessages(records: readonly Record[]): ContextMessage {

export const promptLaunchingKey = defineState<boolean>('prompt.launching', () => false);

export class AgentPromptService implements IAgentPromptService {
export const STEER_REMINDER = [
'The user sent a new message while you were working; it appears above, delivered into the running turn.',
'Address it as you continue this turn; where it changes the current task or approach, the new message takes precedence.',
].join(' ');

export class AgentPromptService extends Disposable implements IAgentPromptService {
declare readonly _serviceBrand: undefined;
private active: (Record & { turn: Turn }) | undefined;
private readonly pending: Record[] = [];
private readonly steered = new Map<string, Record[]>();
private readonly reservedPromptIds = new Set<string>();
private steering = 0;
private steerReminderArmed = false;
private fullCompactionService: IAgentFullCompactionService | undefined;
readonly hooks = { onBeforeSubmitPrompt: new OrderedHookSlot<PromptSubmitContext>() };

Expand All @@ -244,13 +251,19 @@ export class AgentPromptService implements IAgentPromptService {
@ISessionContext private readonly sessionContext: ISessionContext,
@IAgentScopeContext private readonly scopeContext: IAgentScopeContext,
) {
super();
this.states.contributeState(promptLaunchingKey);
this.states.contributeState(promptAdmissionKey);
this.states.contributeState(promptResolutionKey);
toolExecutor.hooks.onDidExecuteTool.register('prompt-service-delivery', async (ctx, next) => {
this._register(this.reminder.register('steer', () => (this.steerReminderArmed ? STEER_REMINDER : undefined)));
this._register(this.loop.hooks.onDidFinishStep.register('steer-reminder', async (_ctx, next) => {
this.steerReminderArmed = false;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve the steer reminder through same-step compaction

When the steered message pushes the context over the automatic-compaction threshold, the reminder hook invokes this provider before the full-compaction hook, so this line clears the flag and appends the reminder; compaction then drops that injection and triggers the reminder service's same-step reconciliation, but the cleared flag prevents re-emission. The resulting model request contains the steered user message without the reminder this change is meant to guarantee, so keep the delivery pending until its injection survives compaction or use the supported one-off delivery path.

AGENTS.md reference: packages/agent-core-v2/AGENTS.md:L80-L82

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed reachable via the modern compaction shape: buildContextCompactionShape keeps real user input verbatim (the steered message included) while dropping every injection-kind message, so the same-step rearm pass could indeed rebuild the request with the steered message but without the reminder, and the consumed flag blocked re-emission.

Fixed in ee9d165 by making the provider restate rather than consume: it emits while the flag is armed, and the flag is now cleared by an onDidFinishStep hook (plus the existing settle() guard) instead of at emission time. The same-step post-compaction reconciliation re-emits the reminder after the splice, and the step-end clear keeps later steps from duplicating it — the service invokes the provider once per injection pass, and a second pass only happens after a splice that already dropped the first emission.

await next();
}));
this._register(toolExecutor.hooks.onDidExecuteTool.register('prompt-service-delivery', async (ctx, next) => {
await this.deliverToolResult(ctx);
await next();
});
}));
}

private get launching(): boolean {
Expand Down Expand Up @@ -413,6 +426,7 @@ export class AgentPromptService implements IAgentPromptService {
this.pending.splice(index, 1);
}
const request = new SteerStepRequest(rerouted, captions, this.reminder, (materialized) => {
this.steerReminderArmed = true;
void this.dispatcher.dispatch(
new TurnSteer({
agentId: this.scopeContext.agentId,
Expand Down Expand Up @@ -513,6 +527,7 @@ export class AgentPromptService implements IAgentPromptService {
private settle(item: Record, result: TurnResult): void {
if (this.active?.id !== item.id) return;
this.active = undefined;
this.steerReminderArmed = false;
const state = result.type === 'cancelled' ? 'cancelled' : result.type === 'failed' ? 'failed' : 'completed';
item.state = state; item.completionDeferred.resolve({ promptId: item.id, result, state });
for (const child of this.steered.get(item.id) ?? []) { child.state = state; child.completionDeferred.resolve({ promptId: child.id, result, state }); }
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ interface ReminderEntry {
readonly variant: string;
}

const REMINDER_VARIANT_PRIORITY = new Map<string, number>([['date_change', -1]]);
const REMINDER_VARIANT_PRIORITY = new Map<string, number>([['date_change', -1], ['steer', 1]]);

interface ReminderActorContext {
readonly entries: Set<ReminderEntry>;
Expand Down
159 changes: 143 additions & 16 deletions packages/agent-core-v2/test/agent/prompt/promptService.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import { describe, expect, it, onTestFinished, vi } from 'vitest';

import { Readable } from 'node:stream';

import { DisposableStore } from '#/_base/di/lifecycle';
import { DisposableStore, toDisposable } from '#/_base/di/lifecycle';
import { createServices } from '#/_base/di/test';
import { Event } from '#/_base/event';
import { IAgentBlobService } from '#/agent/blob/agentBlobService';
Expand All @@ -13,11 +13,12 @@ import { IAgentFullCompactionService } from '#/agent/fullCompaction/fullCompacti
import { IAgentLoopService } from '#/agent/loop/loop';
import { TurnSteer } from '#/agent/loop/turnOps';
import { IAgentPromptService } from '#/agent/prompt/prompt';
import { AgentPromptService, PromptAborted, PromptCompleted, PromptQueued, PromptStarted, PromptSteered, PromptSubmitted } from '#/agent/prompt/promptService';
import { AgentPromptService, PromptAborted, PromptCompleted, PromptQueued, PromptStarted, PromptSteered, PromptSubmitted, STEER_REMINDER } from '#/agent/prompt/promptService';
import { IAgentScopeContext, makeAgentScopeContext } from '#/agent/scopeContext/scopeContext';
import { wrapSystemReminder } from '#/features/reminder/systemReminder';
import { IAgentReminderService } from '#/features/reminder/reminderService';
import { createReminderStub } from '../../features/reminder/stubs';
import type { ContextInjectionContext, ContextInjectionProvider } from '#/features/reminder/types';
import { createReminderHarness, createReminderStub } from '../../features/reminder/stubs';
import { IAgentToolExecutorService } from '#/agent/toolExecutor/toolExecutor';
import { IAgentToolPolicyService } from '#/agent/toolPolicy/toolPolicy';
import { IEventBus, ISessionEventBus } from '#/app/event/eventBus';
Expand All @@ -35,7 +36,7 @@ import { IFileService } from '#/app/file/fileService';
import { ISessionMediaStore } from '#/agent/media/sessionMediaStore';

import { stubContextMemory } from '../contextMemory/stubs';
import { stubLoopWithHooks, stubToolExecutor, stubWire, type StubLoopOptions } from '../loop/stubs';
import { runWillBeginStepHooks, stubLoopWithHooks, stubToolExecutor, stubWire, type StubLoopOptions } from '../loop/stubs';
import { registerStateServices } from '../../state/stubs';
import { SteerStepRequest } from '#/agent/prompt/promptStepRequests';

Expand All @@ -59,21 +60,28 @@ const noopBlob: IAgentBlobService = {
isBlobRef: () => false,
};

function harness(loopOptions: StubLoopOptions = { pendingTurnResult: true }) {
function harness(loopOptions: StubLoopOptions & { integrationReminder?: boolean } = { pendingTurnResult: true }) {
const disposables = new DisposableStore();
onTestFinished(() => disposables.dispose());
const context = stubContextMemory();
const reminder = createReminderStub({
notify: (content, notification) => {
context.append({
role: 'user',
content: [{ type: 'text', text: wrapSystemReminder(content) }],
toolCalls: [],
origin: { kind: 'injection', ...notification },
});
},
});
const loop = stubLoopWithHooks(loopOptions);
const reminderProviders = new Map<string, ContextInjectionProvider>();
const reminder = loopOptions.integrationReminder === true
? createReminderHarness(loop, context)
: createReminderStub({
register: (variant, provider) => {
reminderProviders.set(variant, provider as ContextInjectionProvider);
return toDisposable(() => { reminderProviders.delete(variant); });
},
notify: (content, notification) => {
context.append({
role: 'user',
content: [{ type: 'text', text: wrapSystemReminder(content) }],
toolCalls: [],
origin: { kind: 'injection', ...notification },
});
},
});
const fullCompaction = {
_serviceBrand: undefined,
compacting: null,
Expand Down Expand Up @@ -124,7 +132,23 @@ function harness(loopOptions: StubLoopOptions = { pendingTurnResult: true }) {
(ix.get(IEventBus) as ISessionEventBus).activateAgent(
ix.get(IAgentScopeContext).agentContext,
);
return { prompt: ix.get(IAgentPromptService), loop, context, fullCompaction, eventBus: ix.get(IEventBus), intake };
return { prompt: ix.get(IAgentPromptService), loop, context, fullCompaction, eventBus: ix.get(IEventBus), intake, reminderProviders };
}

function injectionContext(): ContextInjectionContext {
return { injectedPositions: [], lastInjectedAt: null, lastInjection: undefined, lastDisclosure: undefined, isNewTurn: false };
}

async function runDidFinishStepHooks(loop: IAgentLoopService): Promise<void> {
await loop.hooks.onDidFinishStep.run({
turnId: 0,
step: 0,
firstStepOfTurn: false,
signal: new AbortController().signal,
usage: { inputOther: 0, output: 0, inputCacheRead: 0, inputCacheCreation: 0 },
finishReason: 'completed',
stopTurn: false,
});
}

describe('AgentPromptService', () => {
Expand Down Expand Up @@ -230,6 +254,109 @@ describe('AgentPromptService', () => {
expect(events[0]).not.toHaveProperty('promptIds');
});

it('emits the steer reminder from materialize until the step finishes', async () => {
const { prompt, context, loop, reminderProviders } = harness();
const active = await prompt.enqueue({ message: message('active') });
await active.launched;
const queued = await prompt.enqueue({ message: message('new direction') });
const provider = reminderProviders.get('steer')!;
await prompt.steer([queued.id]);
expect(await provider(injectionContext())).toBeUndefined();
loop.drainNextBatch(context);
expect(await provider(injectionContext())).toBe(STEER_REMINDER);
expect(await provider(injectionContext())).toBe(STEER_REMINDER);
Comment on lines +266 to +267

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Exercise reminder delivery through the step hook

These assertions invoke the captured provider directly, so they only verify the private armed flag and never prove the observable contract: that onWillBeginStep appends the steer injection after the message and re-emits it if same-step compaction removes it. A regression in reminder registration, ordering, or compaction reconciliation would leave all four new tests green; wire the real reminder service, drive the step hook, and assert the resulting context messages instead.

AGENTS.md reference: packages/agent-core-v2/AGENTS.md:L80-L82

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair point — the provider-level assertions only proved the flag. Added an integration path in 02e83d1: harness({ integrationReminder: true }) swaps the notify-only stub for createReminderHarness, so the new test drives the real onWillBeginStep hook chain and asserts observable context messages: the wrapped steer-variant injection lands after the steered user message, a real applyCompaction (modern shape) drops the injection while keeping the steered message verbatim, the next hook run re-emits it, and after onDidFinishStep a further run appends nothing.

The kept provider-level tests still pin the arming rules (materialize-time arm, tool-inject exemption, settle clear). The one piece exercised elsewhere rather than here is the injector's intra-step second pass after a mid-chain compaction splice — that's covered by the reminder service's own suite ("re-reconciles within the same step when compaction lands inside the step hook chain"), which this test composes with.

await runDidFinishStepHooks(loop);
expect(await provider(injectionContext())).toBeUndefined();
});

it('keeps one steer reminder armed when separate steers merge into the same step', async () => {
const { prompt, context, loop, reminderProviders } = harness();
const active = await prompt.enqueue({ message: message('active') });
await active.launched;
const one = await prompt.enqueue({ message: message('one') });
const two = await prompt.enqueue({ message: message('two') });
await prompt.steer([one.id]);
await prompt.steer([two.id]);
loop.drainNextBatch(context);
const provider = reminderProviders.get('steer')!;
expect(await provider(injectionContext())).toBe(STEER_REMINDER);
await runDidFinishStepHooks(loop);
expect(await provider(injectionContext())).toBeUndefined();
});

it('delivers the steer reminder through the step hook and re-emits it after compaction drops it', async () => {
const { prompt, context, loop } = harness({ pendingTurnResult: true, integrationReminder: true });
const steerInjections = () =>
context.get().filter((m) => m.origin?.kind === 'injection' && m.origin.variant === 'steer');
const steerTextIndex = () =>
context.get().findIndex((m) => m.content.some((p) => p.type === 'text' && p.text === 'new direction'));

const active = await prompt.enqueue({ message: message('active') });
await active.launched;
const queued = await prompt.enqueue({ message: message('new direction') });
await prompt.steer([queued.id]);
loop.drainNextBatch(context);

await runWillBeginStepHooks(loop);
expect(steerInjections()).toHaveLength(1);
const injected = context.get().at(-1)!;
expect(injected.origin).toEqual({ kind: 'injection', variant: 'steer' });
expect(injected.content).toEqual([{ type: 'text', text: wrapSystemReminder(STEER_REMINDER) }]);
expect(steerTextIndex()).toBeGreaterThanOrEqual(0);
expect(steerTextIndex()).toBeLessThan(context.get().length - 1);

context.applyCompaction({ summary: 'summary', contextSummary: 'summary', compactedCount: context.get().length, tokensBefore: 0 });

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Exercise same-step compaction through the real reminder hook

Fresh evidence after the earlier review is that the new integration path still uses createReminderHarness, not AgentReminderService, and line 308 performs compaction only after the first onWillBeginStep run has returned; it then starts a separate hook run. The helper also receives no event bus, so this passes merely because the armed provider emits on every invocation and never exercises production's post-next() same-step reconciliation—the branch that prevents the reported reminder loss could be deleted while this test remains green. Wire the real reminder service and trigger compaction later in the same hook chain.

AGENTS.md reference: packages/agent-core-v2/AGENTS.md:L82-L82

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The premise that the branch could be deleted with all tests green doesn't hold at the repo level: production's post-next() same-step reconciliation is pinned by the reminder service's own suite ("re-reconciles within the same step when compaction lands inside the step hook chain" in test/features/reminder/reminder.test.ts), driven through the real AgentReminderService — deleting that branch turns that suite red. The promptService test deliberately pins the other half of the composition, the provider's armed-restate contract (emits while armed across passes, stops after onDidFinishStep), against the shared harness. Duplicating the actor-backed real-service rig inside promptService tests to re-prove a branch already pinned where it lives would add test surface without adding failure coverage, so leaving this as-is.

expect(steerInjections()).toHaveLength(0);
expect(steerTextIndex()).toBeGreaterThanOrEqual(0);

await runWillBeginStepHooks(loop);
expect(steerInjections()).toHaveLength(1);
expect(context.get().at(-1)!.content).toEqual([{ type: 'text', text: wrapSystemReminder(STEER_REMINDER) }]);
expect(steerTextIndex()).toBeLessThan(context.get().length - 1);

await runDidFinishStepHooks(loop);
await runWillBeginStepHooks(loop);
expect(steerInjections()).toHaveLength(1);
});

it('re-emits the steer reminder when a step re-runs its injection pass before finishing', async () => {
const { prompt, context, loop } = harness({ pendingTurnResult: true, integrationReminder: true });
const steerInjections = () =>
context.get().filter((m) => m.origin?.kind === 'injection' && m.origin.variant === 'steer');
const active = await prompt.enqueue({ message: message('active') });
await active.launched;
const queued = await prompt.enqueue({ message: message('retry target') });
await prompt.steer([queued.id]);
loop.drainNextBatch(context);
await runWillBeginStepHooks(loop);
await runWillBeginStepHooks(loop);
expect(steerInjections()).toHaveLength(2);
await runDidFinishStepHooks(loop);
await runWillBeginStepHooks(loop);
expect(steerInjections()).toHaveLength(2);
});

it('does not arm the steer reminder for tool-injected steers', async () => {
const { prompt, context, loop, reminderProviders } = harness();
const active = await prompt.enqueue({ message: message('active') });
await active.launched;
await prompt.inject({ role: 'user', content: [{ type: 'text', text: 'tool delivery' }], toolCalls: [] });
loop.drainNextBatch(context);
expect(await reminderProviders.get('steer')!(injectionContext())).toBeUndefined();
});

it('drops an armed steer reminder when the turn settles before the next step', async () => {
const { prompt, context, loop, reminderProviders } = harness({ manualTurnResult: true });
const active = await prompt.enqueue({ message: message('active') });
await active.launched;
const queued = await prompt.enqueue({ message: message('late') });
await prompt.steer([queued.id]);
loop.drainNextBatch(context);
loop.settleActive({ type: 'cancelled', steps: 1, reason: new Error('stop') });
await active.completion;
expect(await reminderProviders.get('steer')!(injectionContext())).toBeUndefined();
});

it('aborts pending prompts and settles completion', async () => {
const { prompt, eventBus } = harness();
const aborted: PromptAborted[] = [];
Expand Down
Loading