diff --git a/lib/commands/context.ts b/lib/commands/context.ts index 9923c53d..bd2e8661 100644 --- a/lib/commands/context.ts +++ b/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 || "" diff --git a/lib/commands/sweep.ts b/lib/commands/sweep.ts index 961976af..a241b684 100644 --- a/lib/commands/sweep.ts +++ b/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) } diff --git a/lib/messages/inject.ts b/lib/messages/inject.ts index c4218260..5920566a 100644 --- a/lib/messages/inject.ts +++ b/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), + ) + } } diff --git a/lib/messages/prune.ts b/lib/messages/prune.ts index f224ce1c..fb86036e 100644 --- a/lib/messages/prune.ts +++ b/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 } diff --git a/lib/messages/utils.ts b/lib/messages/utils.ts index 219027c6..48ae0e6c 100644 --- a/lib/messages/utils.ts +++ b/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 } diff --git a/lib/prompts/discard-tool-spec.ts b/lib/prompts/discard-tool-spec.ts index dcd21358..e5084212 100644 --- a/lib/prompts/discard-tool-spec.ts +++ b/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 diff --git a/lib/prompts/system/both.ts b/lib/prompts/system/both.ts index f5551aaa..9c53a748 100644 --- a/lib/prompts/system/both.ts +++ b/lib/prompts/system/both.ts @@ -2,7 +2,7 @@ export const SYSTEM_PROMPT_BOTH = ` 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 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 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 -After each assistant turn, the environment calls the \`context_info\` tool to inject an assistant message containing a 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 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. diff --git a/lib/prompts/system/discard.ts b/lib/prompts/system/discard.ts index 1bf661fc..e5cd77da 100644 --- a/lib/prompts/system/discard.ts +++ b/lib/prompts/system/discard.ts @@ -2,7 +2,7 @@ export const SYSTEM_PROMPT_DISCARD = ` 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 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 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 -After each assistant turn, the environment calls the \`context_info\` tool to inject an assistant message containing a 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 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. diff --git a/lib/prompts/system/extract.ts b/lib/prompts/system/extract.ts index 859f36dd..3f225e1e 100644 --- a/lib/prompts/system/extract.ts +++ b/lib/prompts/system/extract.ts @@ -2,7 +2,7 @@ export const SYSTEM_PROMPT_EXTRACT = ` 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 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 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 -After each assistant turn, the environment calls the \`context_info\` tool to inject an assistant message containing a 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 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. diff --git a/lib/state/state.ts b/lib/state/state.ts index e68ecf89..69add020 100644 --- a/lib/state/state.ts +++ b/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++ } diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index 057bcf10..38d3b54b 100644 --- a/lib/state/tool-cache.ts +++ b/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 diff --git a/lib/strategies/utils.ts b/lib/strategies/utils.ts index 43bfdf7a..7ae04154 100644 --- a/lib/strategies/utils.ts +++ b/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 }