🤖 Inject mode transition context for flaky test #224

src/services/aiService.ts

@@ -22,6 +22,7 @@ import {
   validateAnthropicCompliance,
   addInterruptedSentinel,
   filterEmptyAssistantMessages,
+  injectModeTransition,
 } from "@/utils/messages/modelMessageTransform";
 import { applyCacheControl } from "@/utils/ai/cacheStrategy";
 import type { HistoryService } from "./historyService";
@@ -450,9 +451,12 @@ export class AIService extends EventEmitter {
       // Add [CONTINUE] sentinel to partial messages (for model context)
       const messagesWithSentinel = addInterruptedSentinel(filteredMessages);
 
+      // Inject mode transition context if mode changed from last assistant message
+      const messagesWithModeContext = injectModeTransition(messagesWithSentinel, mode);
+
       // Apply centralized tool-output redaction BEFORE converting to provider ModelMessages
       // This keeps the persisted/UI history intact while trimming heavy fields for the request
-      const redactedForProvider = applyToolOutputRedaction(messagesWithSentinel);
+      const redactedForProvider = applyToolOutputRedaction(messagesWithModeContext);
       log.debug_obj(`${workspaceId}/2a_redacted_messages.json`, redactedForProvider);
 
       // Convert CmuxMessage to ModelMessage format using Vercel AI SDK utility
@@ -525,6 +529,7 @@ export class AIService extends EventEmitter {
         timestamp: Date.now(),
         model: modelString,
         systemMessageTokens,
+        mode, // Track the mode for this assistant response
       });
 
       // Append to history to get historySequence assigned
@@ -676,6 +681,7 @@ export class AIService extends EventEmitter {
         {
           systemMessageTokens,
           timestamp: Date.now(),
+          mode, // Pass mode so it persists in final history entry
         },
         providerOptions,
         maxOutputTokens,

src/types/message.ts

@@ -18,6 +18,7 @@ export interface CmuxMetadata {
   errorType?: StreamErrorType; // Error type/category if stream failed
   compacted?: boolean; // Whether this message is a compacted summary of previous history
   toolPolicy?: ToolPolicy; // Tool policy active when this message was sent (user messages only)
+  mode?: string; // The mode (plan/exec/etc) active when this message was sent (assistant messages only)
 }
 
 // Extended tool part type that supports interrupted tool calls (input-available state)

src/utils/messages/modelMessageTransform.test.ts

@@ -4,6 +4,7 @@ import {
   transformModelMessages,
   validateAnthropicCompliance,
   addInterruptedSentinel,
+  injectModeTransition,
 } from "./modelMessageTransform";
 import type { CmuxMessage } from "@/types/message";
 
@@ -712,3 +713,138 @@ describe("modelMessageTransform", () => {
     });
   });
 });
+
+describe("injectModeTransition", () => {
+  it("should inject transition message when mode changes", () => {
+    const messages: CmuxMessage[] = [
+      {
+        id: "user-1",
+        role: "user",
+        parts: [{ type: "text", text: "Let's plan a feature" }],
+        metadata: { timestamp: 1000 },
+      },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "Here's the plan..." }],
+        metadata: { timestamp: 2000, mode: "plan" },
+      },
+      {
+        id: "user-2",
+        role: "user",
+        parts: [{ type: "text", text: "Now execute it" }],
+        metadata: { timestamp: 3000 },
+      },
+    ];
+
+    const result = injectModeTransition(messages, "exec");
+
+    // Should have 4 messages: user, assistant, mode-transition, user
+    expect(result.length).toBe(4);
+
+    // Third message should be mode transition
+    expect(result[2].role).toBe("user");
+    expect(result[2].metadata?.synthetic).toBe(true);
+    expect(result[2].parts[0]).toMatchObject({
+      type: "text",
+      text: "[Mode switched from plan to exec. Follow exec mode instructions.]",
+    });
+
+    // Original messages should be preserved
+    expect(result[0]).toEqual(messages[0]);
+    expect(result[1]).toEqual(messages[1]);
+    expect(result[3]).toEqual(messages[2]); // Last user message shifted
+  });
+
+  it("should not inject transition when mode is the same", () => {
+    const messages: CmuxMessage[] = [
+      {
+        id: "user-1",
+        role: "user",
+        parts: [{ type: "text", text: "Let's plan" }],
+        metadata: { timestamp: 1000 },
+      },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "Planning..." }],
+        metadata: { timestamp: 2000, mode: "plan" },
+      },
+      {
+        id: "user-2",
+        role: "user",
+        parts: [{ type: "text", text: "Continue planning" }],
+        metadata: { timestamp: 3000 },
+      },
+    ];
+
+    const result = injectModeTransition(messages, "plan");
+
+    // Should be unchanged
+    expect(result.length).toBe(3);
+    expect(result).toEqual(messages);
+  });
+
+  it("should not inject transition when no previous mode exists", () => {
+    const messages: CmuxMessage[] = [
+      {
+        id: "user-1",
+        role: "user",
+        parts: [{ type: "text", text: "Hello" }],
+        metadata: { timestamp: 1000 },
+      },
+    ];
+
+    const result = injectModeTransition(messages, "exec");
+
+    // Should be unchanged (no assistant message to compare)
+    expect(result.length).toBe(1);
+    expect(result).toEqual(messages);
+  });
+
+  it("should not inject transition when no mode specified", () => {
+    const messages: CmuxMessage[] = [
+      {
+        id: "user-1",
+        role: "user",
+        parts: [{ type: "text", text: "Hello" }],
+        metadata: { timestamp: 1000 },
+      },
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "Hi" }],
+        metadata: { timestamp: 2000, mode: "plan" },
+      },
+      {
+        id: "user-2",
+        role: "user",
+        parts: [{ type: "text", text: "Continue" }],
+        metadata: { timestamp: 3000 },
+      },
+    ];
+
+    const result = injectModeTransition(messages, undefined);
+
+    // Should be unchanged
+    expect(result.length).toBe(3);
+    expect(result).toEqual(messages);
+  });
+
+  it("should handle conversation with no user messages", () => {
+    const messages: CmuxMessage[] = [
+      {
+        id: "assistant-1",
+        role: "assistant",
+        parts: [{ type: "text", text: "Hi" }],
+        metadata: { timestamp: 2000, mode: "plan" },
+      },
+    ];
+
+    const result = injectModeTransition(messages, "exec");
+
+    // Should be unchanged (no user message to inject before)
+    expect(result.length).toBe(1);
+    expect(result).toEqual(messages);
+  });
+});

src/utils/messages/modelMessageTransform.ts

@@ -106,6 +106,84 @@ export function addInterruptedSentinel(messages: CmuxMessage[]): CmuxMessage[] {
   return result;
 }
 
+/**
+ * Inject mode transition context when mode changes mid-conversation.
+ * Inserts a synthetic user message before the final user message to signal the mode switch.
+ * This provides temporal context that helps models understand they should follow new mode instructions.
+ *
+ * @param messages The conversation history
+ * @param currentMode The mode for the upcoming assistant response (e.g., "plan", "exec")
+ * @returns Messages with mode transition context injected if needed
+ */
+export function injectModeTransition(messages: CmuxMessage[], currentMode?: string): CmuxMessage[] {
+  // No mode specified, nothing to do
+  if (!currentMode) {
+    return messages;
+  }
+
+  // Need at least one message to have a conversation
+  if (messages.length === 0) {
+    return messages;
+  }
+
+  // Find the last assistant message to check its mode
+  const lastAssistantMessage = [...messages].reverse().find((m) => m.role === "assistant");
+  const lastMode = lastAssistantMessage?.metadata?.mode;
+
+  // No mode transition if no previous mode or same mode
+  if (!lastMode || lastMode === currentMode) {
+    return messages;
+  }
+
+  // Mode transition detected! Inject a synthetic user message before the last user message
+  // This provides temporal context: user says "switch modes" before their actual request
+
+  // Find the index of the last user message
+  let lastUserIndex = -1;
+  for (let i = messages.length - 1; i >= 0; i--) {
+    if (messages[i].role === "user") {
+      lastUserIndex = i;
+      break;
+    }
+  }
+
+  // If there's no user message, can't inject transition (nothing to inject before)
+  if (lastUserIndex === -1) {
+    return messages;
+  }
+
+  const result: CmuxMessage[] = [];
+
+  // Add all messages up to (but not including) the last user message
+  for (let i = 0; i < lastUserIndex; i++) {
+    result.push(messages[i]);
+  }
+
+  // Inject mode transition message right before the last user message
+  const transitionMessage: CmuxMessage = {
+    id: `mode-transition-${Date.now()}`,
+    role: "user",
+    parts: [
+      {
+        type: "text",
+        text: `[Mode switched from ${lastMode} to ${currentMode}. Follow ${currentMode} mode instructions.]`,
+      },
+    ],
+    metadata: {
+      timestamp: Date.now(),
+      synthetic: true,
+    },
+  };
+  result.push(transitionMessage);
+
+  // Add the last user message and any remaining messages
+  for (let i = lastUserIndex; i < messages.length; i++) {
+    result.push(messages[i]);
+  }
+
+  return result;
+}
+
 /**
  * Split assistant messages with mixed text and tool calls into separate messages
  * to comply with Anthropic's requirement that tool_use blocks must be immediately