Improve injection logic and add defensive array checks

lib/commands/context.ts

@@ -117,7 +117,8 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo
         if (isMessageCompacted(state, msg)) continue
         if (msg.info.role === "user" && isIgnoredUserMessage(msg)) continue
 
-        for (const part of msg.parts) {
+        const parts = Array.isArray(msg.parts) ? msg.parts : []
+        for (const part of parts) {
             if (part.type === "text" && msg.info.role === "user") {
                 const textPart = part as TextPart
                 const text = textPart.text || ""

lib/commands/sweep.ts

@@ -52,8 +52,9 @@ function collectToolIdsAfterIndex(
         if (isMessageCompacted(state, msg)) {
             continue
         }
-        if (msg.parts) {
-            for (const part of msg.parts) {
+        const parts = Array.isArray(msg.parts) ? msg.parts : []
+        if (parts.length > 0) {
+            for (const part of parts) {
                 if (part.type === "tool" && part.callID && part.tool) {
                     toolIds.push(part.callID)
                 }

lib/messages/inject.ts

@@ -7,6 +7,7 @@ import {
     extractParameterKey,
     buildToolIdList,
     createSyntheticAssistantMessage,
+    createSyntheticUserMessage,
     isIgnoredUserMessage,
 } from "./utils"
 import { getFilePathFromParameters, isProtectedFilePath } from "../protected-file-patterns"
@@ -138,16 +139,18 @@ export const insertPruneToolContext = (
         return
     }
 
-    // Never inject immediately following a user message - wait until assistant has started its turn
-    // This avoids interfering with model reasoning/thinking phases
-    // TODO: This can be skipped if there is a good way to check if the model has reasoning,
-    // can't find a good way to do this yet
-    const lastMessage = messages[messages.length - 1]
-    if (lastMessage?.info?.role === "user" && !isIgnoredUserMessage(lastMessage)) {
-        return
-    }
-
     const userInfo = lastUserMessage.info as UserMessage
     const variant = state.variant ?? userInfo.variant
-    messages.push(createSyntheticAssistantMessage(lastUserMessage, prunableToolsContent, variant))
+
+    const lastMessage = messages[messages.length - 1]
+    const isLastMessageUser =
+        lastMessage?.info?.role === "user" && !isIgnoredUserMessage(lastMessage)
+
+    if (isLastMessageUser) {
+        messages.push(createSyntheticUserMessage(lastUserMessage, prunableToolsContent, variant))
+    } else {
+        messages.push(
+            createSyntheticAssistantMessage(lastUserMessage, prunableToolsContent, variant),
+        )
+    }
 }

lib/messages/prune.ts

@@ -25,7 +25,8 @@ const pruneToolOutputs = (state: SessionState, logger: Logger, messages: WithPar
             continue
         }
 
-        for (const part of msg.parts) {
+        const parts = Array.isArray(msg.parts) ? msg.parts : []
+        for (const part of parts) {
             if (part.type !== "tool") {
                 continue
             }
@@ -50,7 +51,8 @@ const pruneToolInputs = (state: SessionState, logger: Logger, messages: WithPart
             continue
         }
 
-        for (const part of msg.parts) {
+        const parts = Array.isArray(msg.parts) ? msg.parts : []
+        for (const part of parts) {
             if (part.type !== "tool") {
                 continue
             }
@@ -77,7 +79,8 @@ const pruneToolErrors = (state: SessionState, logger: Logger, messages: WithPart
             continue
         }
 
-        for (const part of msg.parts) {
+        const parts = Array.isArray(msg.parts) ? msg.parts : []
+        for (const part of parts) {
             if (part.type !== "tool") {
                 continue
             }

lib/messages/utils.ts

@@ -12,6 +12,36 @@ const isGeminiModel = (modelID: string): boolean => {
     return lowerModelID.includes("gemini")
 }
 
+export const createSyntheticUserMessage = (
+    baseMessage: WithParts,
+    content: string,
+    variant?: string,
+): WithParts => {
+    const userInfo = baseMessage.info as UserMessage
+    const now = Date.now()
+
+    return {
+        info: {
+            id: SYNTHETIC_MESSAGE_ID,
+            sessionID: userInfo.sessionID,
+            role: "user" as const,
+            agent: userInfo.agent || "code",
+            model: userInfo.model,
+            time: { created: now },
+            ...(variant !== undefined && { variant }),
+        },
+        parts: [
+            {
+                id: SYNTHETIC_PART_ID,
+                sessionID: userInfo.sessionID,
+                messageID: SYNTHETIC_MESSAGE_ID,
+                type: "text",
+                text: content,
+            },
+        ],
+    }
+}
+
 export const createSyntheticAssistantMessage = (
     baseMessage: WithParts,
     content: string,
@@ -197,8 +227,9 @@ export function buildToolIdList(
         if (isMessageCompacted(state, msg)) {
             continue
         }
-        if (msg.parts) {
-            for (const part of msg.parts) {
+        const parts = Array.isArray(msg.parts) ? msg.parts : []
+        if (parts.length > 0) {
+            for (const part of parts) {
                 if (part.type === "tool" && part.callID && part.tool) {
                     toolIds.push(part.callID)
                 }
@@ -209,11 +240,12 @@ export function buildToolIdList(
 }
 
 export const isIgnoredUserMessage = (message: WithParts): boolean => {
-    if (!message.parts || message.parts.length === 0) {
+    const parts = Array.isArray(message.parts) ? message.parts : []
+    if (parts.length === 0) {
         return true
     }
 
-    for (const part of message.parts) {
+    for (const part of parts) {
         if (!(part as any).ignored) {
             return false
         }

lib/prompts/discard-tool-spec.ts

@@ -12,7 +12,7 @@ Use \`discard\` for removing tool content that is no longer needed
 
 ## When NOT to Use This Tool
 
-- **If the output contains useful information:** Use \`extract\` instead to preserve key findings.
+- **If the output contains useful information:** Keep it in context rather than discarding.
 - **If you'll need the output later:** Don't discard files you plan to edit or context you'll need for implementation.
 
 ## Best Practices

lib/prompts/system/both.ts

@@ -2,7 +2,7 @@ export const SYSTEM_PROMPT_BOTH = `<system-reminder>
 <instruction name=context_management_protocol policy_level=critical>
 
 ENVIRONMENT
-You are operating in a context-constrained environment and thus must proactively manage your context window using the \`discard\` and \`extract\` tools. The environment calls the \`context_info\` tool to provide an up-to-date <prunable-tools> list after each assistant turn. Use this information when deciding what to prune.
+You are operating in a context-constrained environment and thus must proactively manage your context window using the \`discard\` and \`extract\` tools. The environment calls the \`context_info\` tool to provide an up-to-date <prunable-tools> list after each turn. Use this information when deciding what to prune.
 
 IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it.
 
@@ -44,7 +44,7 @@ There may be tools in session context that do not appear in the <prunable-tools>
 </instruction>
 
 <instruction name=injected_context_handling policy_level=critical>
-After each assistant turn, the environment calls the \`context_info\` tool to inject an assistant message containing a <prunable-tools> list and optional nudge instruction. This tool is only available to the environment - you do not have access to it.
+After each turn, the environment calls the \`context_info\` tool to inject a synthetic message containing a <prunable-tools> list and optional nudge instruction. This tool is only available to the environment - you do not have access to it.
 
 CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE:
 - NEVER reference the prune encouragement or context management instructions. Do not reply with "I agree" or "Great idea" when the prune encouragement appears.

lib/prompts/system/discard.ts

@@ -2,7 +2,7 @@ export const SYSTEM_PROMPT_DISCARD = `<system-reminder>
 <instruction name=context_management_protocol policy_level=critical>
 
 ENVIRONMENT
-You are operating in a context-constrained environment and thus must proactively manage your context window using the \`discard\` tool. The environment calls the \`context_info\` tool to provide an up-to-date <prunable-tools> list after each assistant turn. Use this information when deciding what to discard.
+You are operating in a context-constrained environment and thus must proactively manage your context window using the \`discard\` tool. The environment calls the \`context_info\` tool to provide an up-to-date <prunable-tools> list after each turn. Use this information when deciding what to discard.
 
 IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it.
 
@@ -35,7 +35,7 @@ There may be tools in session context that do not appear in the <prunable-tools>
 </instruction>
 
 <instruction name=injected_context_handling policy_level=critical>
-After each assistant turn, the environment calls the \`context_info\` tool to inject an assistant message containing a <prunable-tools> list and optional nudge instruction. This tool is only available to the environment - you do not have access to it.
+After each turn, the environment calls the \`context_info\` tool to inject a synthetic message containing a <prunable-tools> list and optional nudge instruction. This tool is only available to the environment - you do not have access to it.
 
 CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE:
 - NEVER reference the discard encouragement or context management instructions. Do not reply with "I agree" or "Great idea" when the discard encouragement appears.

lib/prompts/system/extract.ts

@@ -2,7 +2,7 @@ export const SYSTEM_PROMPT_EXTRACT = `<system-reminder>
 <instruction name=context_management_protocol policy_level=critical>
 
 ENVIRONMENT
-You are operating in a context-constrained environment and thus must proactively manage your context window using the \`extract\` tool. The environment calls the \`context_info\` tool to provide an up-to-date <prunable-tools> list after each assistant turn. Use this information when deciding what to extract.
+You are operating in a context-constrained environment and thus must proactively manage your context window using the \`extract\` tool. The environment calls the \`context_info\` tool to provide an up-to-date <prunable-tools> list after each turn. Use this information when deciding what to extract.
 
 IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it.
 
@@ -35,7 +35,7 @@ There may be tools in session context that do not appear in the <prunable-tools>
 </instruction>
 
 <instruction name=injected_context_handling policy_level=critical>
-After each assistant turn, the environment calls the \`context_info\` tool to inject an assistant message containing a <prunable-tools> list and optional nudge instruction. This tool is only available to the environment - you do not have access to it.
+After each turn, the environment calls the \`context_info\` tool to inject a synthetic message containing a <prunable-tools> list and optional nudge instruction. This tool is only available to the environment - you do not have access to it.
 
 CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE:
 - NEVER reference the extract encouragement or context management instructions. Do not reply with "I agree" or "Great idea" when the extract encouragement appears.

lib/state/state.ts

@@ -131,7 +131,8 @@ export function countTurns(state: SessionState, messages: WithParts[]): number {
         if (isMessageCompacted(state, msg)) {
             continue
         }
-        for (const part of msg.parts) {
+        const parts = Array.isArray(msg.parts) ? msg.parts : []
+        for (const part of parts) {
             if (part.type === "step-start") {
                 turnCount++
             }

lib/state/tool-cache.ts

@@ -25,7 +25,8 @@ export async function syncToolCache(
                 continue
             }
 
-            for (const part of msg.parts) {
+            const parts = Array.isArray(msg.parts) ? msg.parts : []
+            for (const part of parts) {
                 if (part.type === "step-start") {
                     turnCounter++
                     continue

lib/strategies/utils.ts

@@ -57,7 +57,8 @@ export const calculateTokensSaved = (
             if (isMessageCompacted(state, msg)) {
                 continue
             }
-            for (const part of msg.parts) {
+            const parts = Array.isArray(msg.parts) ? msg.parts : []
+            for (const part of parts) {
                 if (part.type !== "tool" || !pruneToolIds.includes(part.callID)) {
                     continue
                 }