Skip to content

Commit 385830d

Browse files
ThomasK33ThomasK33
authored
🤖 perf: context-efficient plan mode (#1072)
## Summary Optimizes context usage in plan mode by: 1. Removing redundant `planContent` from `propose_plan` tool results (plan is already visible via `file_edit_*` diffs) 2. Including the full plan in the mode transition message when switching plan → exec ## Changes - **`propose_plan` tool**: No longer returns `planContent` in the result, saving context during iterative planning sessions - **Mode transition (plan → exec)**: Now includes the full plan with soft framing: "evaluate whether it's relevant to the user's request" - **UI**: `ProposePlanToolCall` fetches content on-demand for the latest plan; shows path info for historical plans without embedded content - **Backwards compatibility**: Old chat history with `planContent` in results still renders correctly ## Context Flow | Phase | Before | After | |-------|--------|-------| | During planning | Plan in `file_edit_*` diffs + full plan in `propose_plan` result | Plan in `file_edit_*` diffs only | | On mode switch | Generic "mode switched" message | Full plan with "evaluate relevance" framing | ## Testing - Added 4 new tests for plan content in mode transitions - All existing tests pass --- <details> <summary>Plan</summary> # Plan: Context-Efficient Plan Mode ## Summary Improve the plan mode → exec mode transition to include the approved plan content in the model's context only when relevant, avoiding redundant context when the plan is already visible in conversation history or no plan was created. ## Current Behavior ### How Plan Mode Works Now 1. **System Prompt**: When in plan mode, `getPlanModeInstruction()` adds instructions telling the model to write its plan to `~/.mux/plans/{workspaceId}.md` 2. **propose_plan Tool**: When called, reads the plan file from disk and returns: ```typescript { success: true, planPath, planContent, message: "Plan proposed. Waiting for user approval." } ``` This content is stored in the tool result in chat history — **this is redundant** since the plan was already written via `file_edit_*` calls. 3. **Mode Transition**: When switching from plan → exec, `injectModeTransition()` inserts a synthetic user message: ``` [Mode switched from plan to exec. Follow exec mode instructions. Available tools: file_read, bash, ...] ``` ### Key Observation The plan content is duplicated in multiple places: 1. **`file_edit_*` tool calls** - The actual writes/edits to the plan file (as diffs) 2. **`propose_plan` tool result** - The full plan content (redundant!) 3. **Not included** - The mode transition message when switching to exec For iterative planning sessions, this means the plan content appears multiple times, wasting context: - Each plan revision has `file_edit_*` diffs ✓ (necessary, minimal) - Each `propose_plan` call duplicates the full content ✗ (redundant) - Final plan isn't surfaced when switching to exec ✗ (missing) ## Problem Statement When the user switches from plan mode to exec mode (implicitly by changing the mode selector), the model: 1. Doesn't receive explicit confirmation that the plan was approved 2. May not realize the plan from earlier in the conversation is what it should execute 3. Has to infer relevance from context rather than being told explicitly ## Proposed Solution Enhance the mode transition injection to optionally include the approved plan content when switching from plan → exec mode. ### Design Principles 1. **Only include plan when relevant** - The model should determine if the plan applies to the current user message 2. **Avoid redundancy** - Don't include plan if it's already visible in recent context 3. **Implicit approval** - Mux's mode switching is implicit (no explicit approve/reject), so we frame it as "plan available for reference" 4. **Graceful fallback** - If no plan file exists, proceed without it ### Implementation #### 1. Modify `injectModeTransition()` to accept plan content (~20 LoC) **File**: `src/browser/utils/messages/modelMessageTransform.ts` Add an optional `planContent` parameter that gets included when transitioning plan → exec: ```typescript export function injectModeTransition( messages: MuxMessage[], currentMode?: string, toolNames?: string[], planContent?: string // NEW: optional plan content for plan→exec transition ): MuxMessage[] { // ... existing logic ... // If transitioning from plan to exec AND plan content provided if (lastMode === "plan" && currentMode === "exec" && planContent) { transitionText += ` The following plan was developed in plan mode. Evaluate whether it's relevant to the user's request. If relevant, use it to guide your implementation: <approved-plan> ${planContent} </approved-plan>`; } // ... rest of function ... } ``` #### 2. Read plan content during stream preparation (~15 LoC) **File**: `src/node/services/aiService.ts` Before calling `injectModeTransition`, check if we're transitioning plan→exec and read the plan file: ```typescript // In prepareStream(), around line 959: let planContentForTransition: string | undefined; if (mode === "exec") { // Check if last assistant message was in plan mode const lastAssistantMessage = [...filteredMessages].reverse().find(m => m.role === "assistant"); if (lastAssistantMessage?.metadata?.mode === "plan") { // Read plan file for transition context const planFilePath = getPlanFilePath(workspaceId); try { planContentForTransition = await readFileString(runtime, planFilePath); } catch { // No plan file, proceed without } } } const messagesWithModeContext = injectModeTransition( messagesWithSentinel, mode, toolNamesForSentinel, planContentForTransition // NEW parameter ); ``` #### 3. Update test coverage (~40 LoC) **File**: `src/browser/utils/messages/modelMessageTransform.test.ts` Add tests for: - Plan content included when transitioning plan→exec with plan content - Plan content NOT included when transitioning exec→plan - Plan content NOT included when no plan content provided - Plan content NOT included when staying in same mode ### Alternative Considered: Include in System Prompt We could add the plan content to the system prompt during exec mode. However: - This would bust the system message cache on every plan change - Less contextually appropriate (plans are conversation-specific, not workspace-wide) - Mode transition injection is already the established pattern for mode-switching context ### Edge Cases 1. **No plan file exists**: Skip including plan content, use existing transition message 2. **Empty plan file**: Treat same as no plan 3. **Plan from previous conversation**: If user switches modes without a plan in current conversation, no plan content is included 4. **Very long plans**: Consider truncating after N characters with "... (plan truncated, see ~/.mux/plans/...)" #### 4. Remove planContent from propose_plan tool result (~5 LoC) **File**: `src/node/services/tools/propose_plan.ts` The tool currently returns the full plan content in the result. Since: - The plan is already in history via `file_edit_*` tool calls (as diffs) - The plan will be included in the mode transition message when switching to exec We can exclude it from the tool result to save context during iterative planning: ```typescript // Before: return { success: true as const, planPath, planContent, // REMOVE THIS message: "Plan proposed. Waiting for user approval.", }; // After: return { success: true as const, planPath, message: "Plan proposed. Waiting for user approval.", }; ``` #### 5. Update ProposePlanToolCall UI to fetch content on demand (~10 LoC) **File**: `src/browser/components/tools/ProposePlanToolCall.tsx` The UI component already has logic to fetch fresh content via `getPlanContent` for the latest plan. We need to ensure it falls back to this API call when `planContent` is not in the result: ```typescript // The component already handles this case - when isLatest is true AND freshContent is fetched // For historical plans (not latest), we can either: // 1. Fetch on demand when expanded (lazy load) // 2. Show "Plan content not available" with option to fetch ``` Since the UI already prioritizes `freshContent` from disk for the latest plan, this mostly works. For historical plans, we should show a minimal message indicating the plan exists at the path. ## Estimated LoC Changes | File | Change | LoC | |------|--------|-----| | `modelMessageTransform.ts` | Add planContent parameter | +20 | | `aiService.ts` | Read plan on transition | +15 | | `propose_plan.ts` | Remove planContent from result | -3 | | `ProposePlanToolCall.tsx` | Handle missing planContent | +10 | | `modelMessageTransform.test.ts` | New test cases | +40 | | **Total** | | **~82** | ## Design Decisions 1. **No truncation** - Include the full plan content. Can revisit if context limits become an issue. 2. **No "outdated" marking** - The existing file change notification system already handles external edits with diffs. 3. **Soft framing** - Use "developed in plan mode, evaluate relevance" rather than "approved" since Mux's approval is implicit. </details> --- _Generated with [mux](https://github.com/coder/mux)_ Signed-off-by: Thomas Kosiewski <[email protected]>
1 parent 92fd55f commit 385830d

File tree

6 files changed

+272
-63
lines changed

6 files changed

+272
-63
lines changed

‎src/browser/components/tools/ProposePlanToolCall.tsx‎

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,19 +27,26 @@ import { usePopoverError } from "@/browser/hooks/usePopoverError";
2727
import { PopoverError } from "../PopoverError";
2828

2929
/**
30-
* Check if the result is a successful file-based propose_plan result
30+
* Check if the result is a successful file-based propose_plan result.
31+
* Note: planContent may be absent in newer results (context optimization).
3132
*/
3233
function isProposePlanResult(result: unknown): result is ProposePlanToolResult {
3334
return (
3435
result !== null &&
3536
typeof result === "object" &&
3637
"success" in result &&
3738
result.success === true &&
38-
"planContent" in result &&
3939
"planPath" in result
4040
);
4141
}
4242

43+
/**
44+
* Result type that may have planContent (for backwards compatibility with old chat history)
45+
*/
46+
interface ProposePlanResultWithContent extends ProposePlanToolResult {
47+
planContent?: string;
48+
}
49+
4350
/**
4451
* Check if the result is an error from propose_plan tool
4552
*/
@@ -173,11 +180,20 @@ export const ProposePlanToolCall: React.FC<ProposePlanToolCallProps> = (props) =
173180
const titleMatch = /^#\s+(.+)$/m.exec(freshContent);
174181
planTitle = titleMatch ? titleMatch[1] : (planPath?.split("/").pop() ?? "Plan");
175182
} else if (isProposePlanResult(result)) {
176-
planContent = result.planContent;
183+
// New format: planContent may be absent (context optimization)
184+
// For backwards compatibility, check if planContent exists in old chat history
185+
const resultWithContent = result as ProposePlanResultWithContent;
177186
planPath = result.planPath;
178-
// Extract title from first markdown heading or use filename
179-
const titleMatch = /^#\s+(.+)$/m.exec(result.planContent);
180-
planTitle = titleMatch ? titleMatch[1] : (planPath.split("/").pop() ?? "Plan");
187+
if (resultWithContent.planContent) {
188+
// Old result with embedded content (backwards compatibility)
189+
planContent = resultWithContent.planContent;
190+
const titleMatch = /^#\s+(.+)$/m.exec(resultWithContent.planContent);
191+
planTitle = titleMatch ? titleMatch[1] : (planPath.split("/").pop() ?? "Plan");
192+
} else {
193+
// New result without content - show path info, content is fetched for latest
194+
planContent = `*Plan saved to ${planPath}*`;
195+
planTitle = planPath.split("/").pop() ?? "Plan";
196+
}
181197
} else if (isLegacyProposePlanResult(result)) {
182198
// Legacy format: title + plan passed directly (no file)
183199
planContent = result.plan;

‎src/browser/utils/messages/modelMessageTransform.test.ts‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -874,6 +874,155 @@ describe("injectModeTransition", () => {
874874
text: "[Mode switched from plan to exec. Follow exec mode instructions.]",
875875
});
876876
});
877+
878+
it("should include plan content when transitioning from plan to exec", () => {
879+
const messages: MuxMessage[] = [
880+
{
881+
id: "user-1",
882+
role: "user",
883+
parts: [{ type: "text", text: "Let's plan a feature" }],
884+
metadata: { timestamp: 1000 },
885+
},
886+
{
887+
id: "assistant-1",
888+
role: "assistant",
889+
parts: [{ type: "text", text: "Here's the plan..." }],
890+
metadata: { timestamp: 2000, mode: "plan" },
891+
},
892+
{
893+
id: "user-2",
894+
role: "user",
895+
parts: [{ type: "text", text: "Now execute it" }],
896+
metadata: { timestamp: 3000 },
897+
},
898+
];
899+
900+
const planContent = "# My Plan\n\n## Step 1\nDo something\n\n## Step 2\nDo more";
901+
const result = injectModeTransition(messages, "exec", undefined, planContent);
902+
903+
expect(result.length).toBe(4);
904+
const transitionMessage = result[2];
905+
expect(transitionMessage.role).toBe("user");
906+
expect(transitionMessage.metadata?.synthetic).toBe(true);
907+
908+
const textPart = transitionMessage.parts[0];
909+
expect(textPart.type).toBe("text");
910+
if (textPart.type === "text") {
911+
expect(textPart.text).toContain(
912+
"[Mode switched from plan to exec. Follow exec mode instructions.]"
913+
);
914+
expect(textPart.text).toContain("The following plan was developed in plan mode");
915+
expect(textPart.text).toContain("<plan>");
916+
expect(textPart.text).toContain(planContent);
917+
expect(textPart.text).toContain("</plan>");
918+
}
919+
});
920+
921+
it("should NOT include plan content when transitioning from exec to plan", () => {
922+
const messages: MuxMessage[] = [
923+
{
924+
id: "user-1",
925+
role: "user",
926+
parts: [{ type: "text", text: "Done with feature" }],
927+
metadata: { timestamp: 1000 },
928+
},
929+
{
930+
id: "assistant-1",
931+
role: "assistant",
932+
parts: [{ type: "text", text: "Feature complete" }],
933+
metadata: { timestamp: 2000, mode: "exec" },
934+
},
935+
{
936+
id: "user-2",
937+
role: "user",
938+
parts: [{ type: "text", text: "Let's plan the next one" }],
939+
metadata: { timestamp: 3000 },
940+
},
941+
];
942+
943+
const planContent = "# Old Plan\n\nSome content";
944+
const result = injectModeTransition(messages, "plan", undefined, planContent);
945+
946+
expect(result.length).toBe(4);
947+
const transitionMessage = result[2];
948+
const textPart = transitionMessage.parts[0];
949+
if (textPart.type === "text") {
950+
expect(textPart.text).toBe(
951+
"[Mode switched from exec to plan. Follow plan mode instructions.]"
952+
);
953+
expect(textPart.text).not.toContain("<plan>");
954+
}
955+
});
956+
957+
it("should NOT include plan content when no plan content provided", () => {
958+
const messages: MuxMessage[] = [
959+
{
960+
id: "user-1",
961+
role: "user",
962+
parts: [{ type: "text", text: "Let's plan" }],
963+
metadata: { timestamp: 1000 },
964+
},
965+
{
966+
id: "assistant-1",
967+
role: "assistant",
968+
parts: [{ type: "text", text: "Planning..." }],
969+
metadata: { timestamp: 2000, mode: "plan" },
970+
},
971+
{
972+
id: "user-2",
973+
role: "user",
974+
parts: [{ type: "text", text: "Execute" }],
975+
metadata: { timestamp: 3000 },
976+
},
977+
];
978+
979+
const result = injectModeTransition(messages, "exec", undefined, undefined);
980+
981+
expect(result.length).toBe(4);
982+
const transitionMessage = result[2];
983+
const textPart = transitionMessage.parts[0];
984+
if (textPart.type === "text") {
985+
expect(textPart.text).toBe(
986+
"[Mode switched from plan to exec. Follow exec mode instructions.]"
987+
);
988+
expect(textPart.text).not.toContain("<plan>");
989+
}
990+
});
991+
992+
it("should include both tools and plan content in transition message", () => {
993+
const messages: MuxMessage[] = [
994+
{
995+
id: "user-1",
996+
role: "user",
997+
parts: [{ type: "text", text: "Plan done" }],
998+
metadata: { timestamp: 1000 },
999+
},
1000+
{
1001+
id: "assistant-1",
1002+
role: "assistant",
1003+
parts: [{ type: "text", text: "Plan ready" }],
1004+
metadata: { timestamp: 2000, mode: "plan" },
1005+
},
1006+
{
1007+
id: "user-2",
1008+
role: "user",
1009+
parts: [{ type: "text", text: "Go" }],
1010+
metadata: { timestamp: 3000 },
1011+
},
1012+
];
1013+
1014+
const toolNames = ["file_read", "bash"];
1015+
const planContent = "# Plan\n\nDo stuff";
1016+
const result = injectModeTransition(messages, "exec", toolNames, planContent);
1017+
1018+
expect(result.length).toBe(4);
1019+
const textPart = result[2].parts[0];
1020+
if (textPart.type === "text") {
1021+
expect(textPart.text).toContain("Available tools: file_read, bash.]");
1022+
expect(textPart.text).toContain("<plan>");
1023+
expect(textPart.text).toContain(planContent);
1024+
}
1025+
});
8771026
});
8781027

8791028
describe("filterEmptyAssistantMessages", () => {

‎src/browser/utils/messages/modelMessageTransform.ts‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,15 +112,20 @@ export function addInterruptedSentinel(messages: MuxMessage[]): MuxMessage[] {
112112
* Inserts a synthetic user message before the final user message to signal the mode switch.
113113
* This provides temporal context that helps models understand they should follow new mode instructions.
114114
*
115+
* When transitioning from plan → exec mode with plan content, includes the plan so the model
116+
* can evaluate its relevance to the current request.
117+
*
115118
* @param messages The conversation history
116119
* @param currentMode The mode for the upcoming assistant response (e.g., "plan", "exec")
117120
* @param toolNames Optional list of available tool names to include in transition message
121+
* @param planContent Optional plan content to include when transitioning plan → exec
118122
* @returns Messages with mode transition context injected if needed
119123
*/
120124
export function injectModeTransition(
121125
messages: MuxMessage[],
122126
currentMode?: string,
123-
toolNames?: string[]
127+
toolNames?: string[],
128+
planContent?: string
124129
): MuxMessage[] {
125130
// No mode specified, nothing to do
126131
if (!currentMode) {
@@ -175,6 +180,17 @@ export function injectModeTransition(
175180
transitionText += "]";
176181
}
177182

183+
// When transitioning plan → exec with plan content, include the plan for context
184+
if (lastMode === "plan" && currentMode === "exec" && planContent) {
185+
transitionText += `
186+
187+
The following plan was developed in plan mode. Based on the user's message, determine if they have accepted the plan. If accepted and relevant, use it to guide your implementation:
188+
189+
<plan>
190+
${planContent}
191+
</plan>`;
192+
}
193+
178194
const transitionMessage: MuxMessage = {
179195
id: `mode-transition-${Date.now()}`,
180196
role: "user",

‎src/common/types/tools.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -170,11 +170,12 @@ export type FileEditToolArgs =
170170
// Args derived from schema
171171
export type ProposePlanToolArgs = z.infer<typeof TOOL_DEFINITIONS.propose_plan.schema>;
172172

173-
// Result type for new file-based propose_plan tool
173+
// Result type for file-based propose_plan tool
174+
// Note: planContent is NOT included to save context - plan is visible via file_edit_* diffs
175+
// and will be included in mode transition message when switching to exec mode
174176
export interface ProposePlanToolResult {
175177
success: true;
176178
planPath: string;
177-
planContent: string;
178179
message: string;
179180
}
180181

0 commit comments

Comments
 (0)