Ver código fonte

fix: harden message transform guards

Zaradacht Taifour 1 mês atrás
pai
commit
cc93aaf44e

+ 4 - 2
src/hooks/filter-available-skills/index.ts

@@ -117,9 +117,11 @@ export function createFilterAvailableSkillsHook(
   return {
     'experimental.chat.messages.transform': async (
       _input: Record<string, never>,
-      output: { messages: MessageWithParts[] },
+      output: { messages?: unknown },
     ): Promise<void> => {
-      const messages = output.messages.filter(isMessageWithParts);
+      const messages = (Array.isArray(output.messages) ? output.messages : []).filter(
+        isMessageWithParts,
+      );
       if (messages.length === 0) {
         return;
       }

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

@@ -121,6 +121,17 @@ describe('createPhaseReminderHook', () => {
     expect(output.messages).toEqual([]);
   });
 
+  test('handles missing or non-array messages', async () => {
+    const hook = createPhaseReminderHook();
+
+    await expect(
+      hook['experimental.chat.messages.transform']({}, {}),
+    ).resolves.toBeUndefined();
+    await expect(
+      hook['experimental.chat.messages.transform']({}, { messages: {} }),
+    ).resolves.toBeUndefined();
+  });
+
   test('handles no user messages', async () => {
     const hook = createPhaseReminderHook();
     const output = {
@@ -136,4 +147,27 @@ describe('createPhaseReminderHook', () => {
 
     expect(output.messages[0].parts[0].text).toBe('Hi');
   });
+
+  test('skips malformed messages while still appending to latest valid user message', async () => {
+    const hook = createPhaseReminderHook();
+    const output = {
+      messages: [
+        {},
+        { info: { role: 'assistant' } },
+        { parts: [{ type: 'text', text: 'missing info' }] },
+        {
+          info: { role: 'user', agent: 'orchestrator' },
+          parts: [{ type: 'text', text: 'hello' }],
+        },
+      ],
+    };
+
+    await expect(
+      hook['experimental.chat.messages.transform']({}, output as never),
+    ).resolves.toBeUndefined();
+
+    expect(output.messages[3].parts.length).toBe(2);
+    expect(output.messages[3].parts[0].text).toBe('hello');
+    expect(output.messages[3].parts[1].text).toBe(PHASE_REMINDER);
+  });
 });

+ 8 - 4
src/hooks/phase-reminder/index.ts

@@ -7,7 +7,7 @@
  */
 import { PHASE_REMINDER } from '../../config/constants';
 import { SLIM_INTERNAL_INITIATOR_MARKER } from '../../utils';
-import type { MessageWithParts } from '../types';
+import { isUserMessageWithParts, type MessageWithParts } from '../types';
 
 export { PHASE_REMINDER };
 
@@ -20,9 +20,9 @@ export function createPhaseReminderHook() {
   return {
     'experimental.chat.messages.transform': async (
       _input: Record<string, never>,
-      output: { messages: MessageWithParts[] },
+      output: { messages?: unknown },
     ): Promise<void> => {
-      const { messages } = output;
+      const messages = Array.isArray(output.messages) ? output.messages : [];
 
       if (messages.length === 0) {
         return;
@@ -30,7 +30,7 @@ export function createPhaseReminderHook() {
 
       let lastUserMessageIndex = -1;
       for (let i = messages.length - 1; i >= 0; i--) {
-        if (messages[i].info.role === 'user') {
+        if (isUserMessageWithParts(messages[i])) {
           lastUserMessageIndex = i;
           break;
         }
@@ -41,6 +41,10 @@ export function createPhaseReminderHook() {
       }
 
       const lastUserMessage = messages[lastUserMessageIndex];
+      if (!isUserMessageWithParts(lastUserMessage)) {
+        return;
+      }
+
       const agent = lastUserMessage.info.agent;
       if (agent && agent !== 'orchestrator') {
         return;

+ 6 - 4
src/hooks/task-session-manager/index.ts

@@ -643,9 +643,11 @@ export function createTaskSessionManagerHook(
 
     'experimental.chat.messages.transform': async (
       _input: Record<string, never>,
-      output: { messages: MessageWithParts[] },
+      output: { messages?: unknown },
     ): Promise<void> => {
-      for (const [messageIndex, message] of output.messages.entries()) {
+      const messages = Array.isArray(output.messages) ? output.messages : [];
+
+      for (const [messageIndex, message] of messages.entries()) {
         if (!isUserMessageWithParts(message)) continue;
         if (message.info.agent && message.info.agent !== 'orchestrator') {
           continue;
@@ -662,8 +664,8 @@ export function createTaskSessionManagerHook(
         }
       }
 
-      for (let i = output.messages.length - 1; i >= 0; i -= 1) {
-        const message = output.messages[i];
+      for (let i = messages.length - 1; i >= 0; i -= 1) {
+        const message = messages[i];
         if (!isUserMessageWithParts(message)) continue;
         if (message.info.agent && message.info.agent !== 'orchestrator') return;
         if (

+ 4 - 2
src/index.ts

@@ -1035,9 +1035,11 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
     // API (doesn't show in UI)
     'experimental.chat.messages.transform': async (
       input: Record<string, never>,
-      output: { messages: unknown[] },
+      output: { messages?: unknown },
     ): Promise<void> => {
-      const messages = output.messages.filter(isMessageWithParts);
+      const messages = (Array.isArray(output.messages) ? output.messages : []).filter(
+        isMessageWithParts,
+      );
 
       for (const message of messages) {
         if (!isUserMessageWithParts(message)) {