Browse Source

Preserve phase reminders across turns

DanMaly 3 weeks ago
parent
commit
7796cbba1a
2 changed files with 142 additions and 23 deletions
  1. 110 0
      src/hooks/phase-reminder/index.test.ts
  2. 32 23
      src/hooks/phase-reminder/index.ts

+ 110 - 0
src/hooks/phase-reminder/index.test.ts

@@ -35,6 +35,95 @@ describe('createPhaseReminderHook', () => {
     });
   });
 
+  test('appends one reminder to every historical orchestrator user message in the session', async () => {
+    const hook = createPhaseReminderHook();
+    const output = {
+      messages: [
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's1' },
+          parts: [{ type: 'text', text: 'first' }],
+        },
+        {
+          info: { role: 'user', agent: 'explorer', sessionID: 's1' },
+          parts: [{ type: 'text', text: 'specialist' }],
+        },
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's2' },
+          parts: [{ type: 'text', text: 'other session' }],
+        },
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's1' },
+          parts: [{ type: 'text', text: 'latest' }],
+        },
+      ],
+    };
+
+    await hook['experimental.chat.messages.transform']({}, output);
+
+    expect(output.messages[0].parts).toHaveLength(2);
+    expect(output.messages[3].parts).toHaveLength(2);
+    expect(output.messages[1].parts).toHaveLength(1);
+    expect(output.messages[2].parts).toHaveLength(1);
+  });
+
+  test('reconstructs byte-identical historical messages on the next turn', async () => {
+    const hook = createPhaseReminderHook();
+    const turnN = {
+      messages: [
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's1' },
+          parts: [{ type: 'text', text: 'first' }],
+        },
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's1' },
+          parts: [{ type: 'text', text: 'second' }],
+        },
+      ],
+    };
+
+    await hook['experimental.chat.messages.transform']({}, turnN);
+    const transformedHistory = structuredClone(turnN.messages);
+    const turnNPlusOne = {
+      messages: [
+        ...turnN.messages.map((message) => ({
+          ...message,
+          parts: message.parts.filter((part) => !part.synthetic),
+        })),
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's1' },
+          parts: [{ type: 'text', text: 'third' }],
+        },
+      ],
+    };
+
+    await hook['experimental.chat.messages.transform']({}, turnNPlusOne);
+
+    expect(turnNPlusOne.messages.slice(0, -1)).toEqual(transformedHistory);
+  });
+
+  test('is idempotent when run twice on the same messages', async () => {
+    const hook = createPhaseReminderHook();
+    const output = {
+      messages: [
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's1' },
+          parts: [{ type: 'text', text: 'first' }],
+        },
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's1' },
+          parts: [{ type: 'text', text: 'latest' }],
+        },
+      ],
+    };
+
+    await hook['experimental.chat.messages.transform']({}, output);
+    await hook['experimental.chat.messages.transform']({}, output);
+
+    for (const message of output.messages) {
+      expect(message.parts.filter((part) => part.synthetic)).toHaveLength(1);
+    }
+  });
+
   test('skips non-orchestrator sessions', async () => {
     const hook = createPhaseReminderHook();
     const output = {
@@ -126,6 +215,27 @@ describe('createPhaseReminderHook', () => {
     ).toBe(false);
   });
 
+  test('replays historical reminders when the latest user message is internal', async () => {
+    const hook = createPhaseReminderHook();
+    const output = {
+      messages: [
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's1' },
+          parts: [{ type: 'text', text: 'historical' }],
+        },
+        {
+          info: { role: 'user', agent: 'orchestrator', sessionID: 's1' },
+          parts: [createInternalAgentTextPart('internal notification')],
+        },
+      ],
+    };
+
+    await hook['experimental.chat.messages.transform']({}, output);
+
+    expect(output.messages[0].parts).toHaveLength(2);
+    expect(output.messages[1].parts).toHaveLength(1);
+  });
+
   test('does not let user-visible internal marker suppress injection', async () => {
     const hook = createPhaseReminderHook();
     const output = {

+ 32 - 23
src/hooks/phase-reminder/index.ts

@@ -8,7 +8,11 @@
 import { PHASE_REMINDER } from '../../config/constants';
 import { isInternalInitiatorPart } from '../../utils';
 import { isRecord } from '../../utils/guards';
-import { findLatestUserMessage, type MessagePart } from '../types';
+import {
+  findLatestUserMessage,
+  isUserMessageWithParts,
+  type MessagePart,
+} from '../types';
 
 export { PHASE_REMINDER };
 
@@ -53,32 +57,37 @@ export function createPhaseReminderHook(options: PhaseReminderOptions = {}) {
         return;
       }
 
-      const textPartIndex = lastUserMessage.parts.findIndex(
-        (p) => p.type === 'text' && p.text !== undefined,
-      );
-
-      if (textPartIndex === -1) {
-        return;
-      }
-
-      const originalPart = lastUserMessage.parts[textPartIndex];
-      if (isInternalInitiatorPart(originalPart)) {
-        return;
-      }
-      if (lastUserMessage.parts.some(hasPhaseReminder)) {
-        return;
-      }
-
       // post-file-tool-nudge must run first so its tagged part deduplicates.
       // Append reminder as a new, separate message part instead of mutating
       // the user-authored text. This prevents the reminder from leaking into
       // the UI display and chat history (issue #448).
-      lastUserMessage.parts.push({
-        type: 'text',
-        synthetic: true,
-        text: PHASE_REMINDER,
-        metadata: { [PHASE_REMINDER_METADATA_KEY]: true },
-      });
+      for (const message of messages) {
+        if (
+          !isUserMessageWithParts(message) ||
+          message.info.agent !== 'orchestrator' ||
+          message.info.sessionID !== sessionID
+        ) {
+          continue;
+        }
+
+        const textPart = message.parts.find(
+          (part) => part.type === 'text' && part.text !== undefined,
+        );
+        if (
+          !textPart ||
+          isInternalInitiatorPart(textPart) ||
+          message.parts.some(hasPhaseReminder)
+        ) {
+          continue;
+        }
+
+        message.parts.push({
+          type: 'text',
+          synthetic: true,
+          text: PHASE_REMINDER,
+          metadata: { [PHASE_REMINDER_METADATA_KEY]: true },
+        });
+      }
     },
   };
 }