keep specific messages during truncation

Author: charleslien

backend/src/run-agent-step.ts

@@ -107,15 +107,16 @@ export const runAgentStep = async (
   })
 
   let messageHistory = agentState.messageHistory
-  const messagesWithUserPrompt = buildArray<CodebuffMessage>(
-    ...messageHistory,
-    prompt && [
-      {
-        role: 'user' as const,
-        content: asUserMessage(prompt),
-      },
-    ],
-  )
+  const messagesWithUserPrompt = prompt
+    ? [
+        ...messageHistory.map((m) => ({ ...m, keepDuringTruncation: false })),
+        {
+          role: 'user' as const,
+          content: asUserMessage(prompt),
+          keepDuringTruncation: true,
+        },
+      ]
+    : messageHistory
 
   // Check number of assistant messages since last user message with prompt
   if (agentState.stepsRemaining <= 0) {
@@ -252,7 +253,9 @@ export const runAgentStep = async (
     : undefined
 
   const agentMessagesUntruncated = buildArray<CodebuffMessage>(
-    ...expireMessages(messageHistory, prompt ? 'userPrompt' : 'agentStep'),
+    ...expireMessages(messageHistory, prompt ? 'userPrompt' : 'agentStep').map(
+      (m) => (prompt ? { ...m, keepDuringTruncation: false } : m),
+    ),
 
     toolResults.length > 0 && {
       role: 'user' as const,
@@ -282,12 +285,14 @@ export const runAgentStep = async (
       role: 'user' as const,
       content: instructionsPrompt,
       timeToLive: 'userPrompt' as const,
+      keepDuringTruncation: true,
     },
 
     stepPrompt && {
       role: 'user' as const,
       content: stepPrompt,
       timeToLive: 'agentStep' as const,
+      keepDuringTruncation: true,
     },
   )
 

backend/src/util/__tests__/messages.test.ts

@@ -159,31 +159,16 @@ describe('trimMessagesToFitTokenLimit', () => {
       maxTotalTokens,
     )
 
-    // Verify the first message was dropped
-    expect(result).toHaveLength(testMessages.length - 1)
-
-    // Regular messages should be unchanged
-    expect(result[0].content).toBe(testMessages[1].content)
-    expect(result[6].content).toEqual(testMessages[7].content)
-
-    // 0th and second terminal outputs should be simplified
-    expect(result[1].role).toEqual(testMessages[2].role)
-    expect(Array.isArray(result[1].content)).toBe(true)
-    expect((result[1].content[0] as any).text).toContain(
-      '<result>[Output omitted]</result>',
-    )
-    expect((result[1].content[1] as any).text).toBe(
-      (testMessages[2].content[1] as any).text,
-    )
-
-    expect(result[2].role).toEqual(testMessages[3].role)
-    expect(result[2].content).toContain('<result>[Output omitted]</result>')
+    // Should have replacement message for omitted content
+    expect(result.length).toBeGreaterThan(0)
 
-    // Terminal outputs 3-7 should be preserved exactly
-    expect(result[3].content).toBe(testMessages[4].content)
-    expect(result[4].content).toEqual(testMessages[5].content)
-    expect(result[5].content).toBe(testMessages[6].content)
-    expect(result[6].content).toBe(testMessages[7].content)
+    // Should contain a replacement message for omitted content
+    const hasReplacementMessage = result.some(
+      (msg) =>
+        typeof msg.content === 'string' &&
+        msg.content.includes('Previous message(s) omitted due to length'),
+    )
+    expect(hasReplacementMessage).toBe(true)
 
     // Verify total tokens are under limit
     const finalTokens = tokenCounter.countTokensJson(result)
@@ -199,31 +184,16 @@ describe('trimMessagesToFitTokenLimit', () => {
       maxTotalTokens,
     )
 
-    // Verify the first message was dropped
-    expect(result).toHaveLength(testMessages.length - 1)
-
-    // Regular messages should be unchanged
-    expect(result[0].content).toBe(testMessages[1].content)
-    expect(result[6].content).toEqual(testMessages[7].content)
+    // Should have replacement message for omitted content
+    expect(result.length).toBeGreaterThan(0)
 
-    // 0th and second terminal outputs should be simplified
-    expect(result[1].role).toEqual(testMessages[2].role)
-    expect(Array.isArray(result[1].content)).toBe(true)
-    expect((result[1].content[0] as any).text).toContain(
-      '<result>[Output omitted]</result>',
-    )
-    expect((result[1].content[1] as any).text).toBe(
-      (testMessages[2].content[1] as any).text,
+    // Should contain a replacement message for omitted content
+    const hasReplacementMessage = result.some(
+      (msg) =>
+        typeof msg.content === 'string' &&
+        msg.content.includes('Previous message(s) omitted due to length'),
     )
-
-    expect(result[2].role).toEqual(testMessages[3].role)
-    expect(result[2].content).toContain('<result>[Output omitted]</result>')
-
-    // Terminal outputs 3-7 should be preserved exactly
-    expect(result[3].content).toBe(testMessages[4].content)
-    expect(result[4].content).toEqual(testMessages[5].content)
-    expect(result[5].content).toBe(testMessages[6].content)
-    expect(result[6].content).toBe(testMessages[7].content)
+    expect(hasReplacementMessage).toBe(true)
 
     // Verify total tokens are under limit
     const finalTokens = tokenCounter.countTokensJson(result)
@@ -258,4 +228,142 @@ describe('trimMessagesToFitTokenLimit', () => {
 
     expect(result).toEqual([])
   })
+
+  describe('keepDuringTruncation functionality', () => {
+    it('preserves messages marked with keepDuringTruncation=true', () => {
+      const messages = [
+        { role: 'user', content: 'A'.repeat(500) }, // Large message to force truncation
+        { role: 'user', content: 'B'.repeat(500) }, // Large message to force truncation
+        {
+          role: 'user',
+          content: 'Message 3 - keep me!',
+          keepDuringTruncation: true,
+        },
+        { role: 'assistant', content: 'C'.repeat(500) }, // Large message to force truncation
+        {
+          role: 'user',
+          content: 'Message 5 - keep me too!',
+          keepDuringTruncation: true,
+        },
+      ] as CodebuffMessage[]
+
+      const result = trimMessagesToFitTokenLimit(messages, 0, 1000)
+
+      // Should contain the kept messages
+      const keptMessages = result.filter(
+        (msg) =>
+          typeof msg.content === 'string' &&
+          (msg.content.includes('keep me!') ||
+            msg.content.includes('keep me too!')),
+      )
+      expect(keptMessages).toHaveLength(2)
+
+      // Should have replacement message for omitted content
+      const hasReplacementMessage = result.some(
+        (msg) =>
+          typeof msg.content === 'string' &&
+          msg.content.includes('Previous message(s) omitted due to length'),
+      )
+      expect(hasReplacementMessage).toBe(true)
+    })
+
+    it('does not add replacement message when no messages are removed', () => {
+      const messages = [
+        { role: 'user', content: 'Short message 1' },
+        {
+          role: 'user',
+          content: 'Short message 2',
+          keepDuringTruncation: true,
+        },
+      ] as CodebuffMessage[]
+
+      const result = trimMessagesToFitTokenLimit(messages, 0, 10000)
+
+      // Should be unchanged when under token limit
+      expect(result).toHaveLength(2)
+      expect(result[0].content).toBe('Short message 1')
+      expect(result[1].content).toBe('Short message 2')
+    })
+
+    it('handles consecutive replacement messages correctly', () => {
+      const messages = [
+        { role: 'user', content: 'A'.repeat(1000) }, // Large message to be removed
+        { role: 'user', content: 'B'.repeat(1000) }, // Large message to be removed
+        { role: 'user', content: 'C'.repeat(1000) }, // Large message to be removed
+        { role: 'user', content: 'Keep this', keepDuringTruncation: true },
+      ] as CodebuffMessage[]
+
+      const result = trimMessagesToFitTokenLimit(messages, 0, 1000)
+
+      // Should only have one replacement message for consecutive removals
+      const replacementMessages = result.filter(
+        (msg) =>
+          typeof msg.content === 'string' &&
+          msg.content.includes('Previous message(s) omitted due to length'),
+      )
+      expect(replacementMessages).toHaveLength(1)
+
+      // Should keep the marked message
+      const keptMessage = result.find(
+        (msg) =>
+          typeof msg.content === 'string' && msg.content.includes('Keep this'),
+      )
+      expect(keptMessage).toBeDefined()
+    })
+
+    it('calculates token removal correctly with keepDuringTruncation', () => {
+      const messages = [
+        { role: 'user', content: 'A'.repeat(500) }, // Will be removed
+        { role: 'user', content: 'B'.repeat(500) }, // Will be removed
+        {
+          role: 'user',
+          content: 'Keep this short message',
+          keepDuringTruncation: true,
+        },
+        { role: 'user', content: 'C'.repeat(100) }, // Might be kept
+      ] as CodebuffMessage[]
+
+      const result = trimMessagesToFitTokenLimit(messages, 0, 2000)
+
+      // Should preserve the keepDuringTruncation message
+      const keptMessage = result.find(
+        (msg) =>
+          typeof msg.content === 'string' &&
+          msg.content.includes('Keep this short message'),
+      )
+      expect(keptMessage).toBeDefined()
+
+      // Total tokens should be under limit
+      const finalTokens = tokenCounter.countTokensJson(result)
+      expect(finalTokens).toBeLessThan(2000)
+    })
+
+    it('handles mixed keepDuringTruncation and regular messages', () => {
+      const messages = [
+        { role: 'user', content: 'A'.repeat(800) }, // Large message to force truncation
+        { role: 'user', content: 'Keep 1', keepDuringTruncation: true },
+        { role: 'user', content: 'B'.repeat(800) }, // Large message to force truncation
+        { role: 'user', content: 'Keep 2', keepDuringTruncation: true },
+        { role: 'user', content: 'C'.repeat(800) }, // Large message to force truncation
+      ] as CodebuffMessage[]
+
+      const result = trimMessagesToFitTokenLimit(messages, 0, 500)
+
+      // Should keep both marked messages
+      const keptMessages = result.filter(
+        (msg) =>
+          typeof msg.content === 'string' &&
+          (msg.content.includes('Keep 1') || msg.content.includes('Keep 2')),
+      )
+      expect(keptMessages).toHaveLength(2)
+
+      // Should have replacement messages for removed content
+      const replacementMessages = result.filter(
+        (msg) =>
+          typeof msg.content === 'string' &&
+          msg.content.includes('Previous message(s) omitted due to length'),
+      )
+      expect(replacementMessages.length).toBeGreaterThan(0)
+    })
+  })
 })

backend/src/util/messages.ts

@@ -113,6 +113,10 @@ function simplifyTerminalHelper(
 
 // Factor to reduce token count target by, to leave room for new messages
 const shortenedMessageTokenFactor = 0.5
+const replacementMessage = {
+  role: 'user',
+  content: asSystemMessage('Previous message(s) omitted due to length'),
+} satisfies CodebuffMessage
 
 /**
  * Trims messages from the beginning to fit within token limits while preserving
@@ -133,18 +137,16 @@ export function trimMessagesToFitTokenLimit(
   systemTokens: number,
   maxTotalTokens: number = 190_000,
 ): CodebuffMessage[] {
-  const MAX_MESSAGE_TOKENS = maxTotalTokens - systemTokens
+  const maxMessageTokens = maxTotalTokens - systemTokens
 
   // Check if we're already under the limit
   const initialTokens = countTokensJson(messages)
 
-  if (initialTokens < MAX_MESSAGE_TOKENS) {
+  if (initialTokens < maxMessageTokens) {
     return messages
   }
 
-  let totalTokens = 0
-  const targetTokens = MAX_MESSAGE_TOKENS * shortenedMessageTokenFactor
-  const results: CodebuffMessage[] = []
+  const shortenedMessages: CodebuffMessage[] = []
   let numKept = 0
 
   // Process messages from newest to oldest
@@ -208,22 +210,41 @@ export function trimMessagesToFitTokenLimit(
         message = { ...m, content: newContent }
       }
     } else {
+      m satisfies never
       throw new AssertionError({ message: 'Not a valid role' })
     }
 
-    // Check if adding this message would exceed our token target
-    const messageTokens = countTokensJson(message)
+    shortenedMessages.push(message)
+  }
+  shortenedMessages.reverse()
 
-    if (totalTokens + messageTokens <= targetTokens) {
-      results.push(message)
-      totalTokens += messageTokens
-    } else {
-      break
+  const requiredTokens = countTokensJson(
+    shortenedMessages.filter((m) => m.keepDuringTruncation),
+  )
+  let removedTokens = 0
+  const tokensToRemove =
+    (maxMessageTokens - requiredTokens) * (1 - shortenedMessageTokenFactor)
+
+  const placeholder = 'deleted'
+  const filteredMessages: (CodebuffMessage | typeof placeholder)[] = []
+  for (const message of shortenedMessages) {
+    if (removedTokens >= tokensToRemove || message.keepDuringTruncation) {
+      filteredMessages.push(message)
+      continue
+    }
+    removedTokens += countTokensJson(message)
+    if (
+      filteredMessages.length === 0 ||
+      filteredMessages[filteredMessages.length - 1] !== placeholder
+    ) {
+      filteredMessages.push(placeholder)
+      removedTokens -= countTokensJson(replacementMessage)
     }
   }
 
-  results.reverse()
-  return results
+  return filteredMessages.map((m) =>
+    m === placeholder ? replacementMessage : m,
+  )
 }
 
 export function getMessagesSubset(

common/src/types/message.ts

@@ -61,6 +61,7 @@ export const CodebuffMessageSchema = z.intersection(
     timeToLive: z
       .union([z.literal('agentStep'), z.literal('userPrompt')])
       .optional(),
+    keepDuringTruncation: z.boolean().optional(),
   }),
 )