feat(context)!: offload conditions decide granularity, and placeholders are typed (typescript)

strands-ts/src/__fixtures__/context-manager-helpers.ts

@@ -0,0 +1,84 @@
+/**
+ * Shared helpers for context-manager tests: messages, a character-based token count, and a
+ * {@link ContextState} over a mock agent.
+ */
+
+import { isOffloaded, Message, TextBlock, ToolResultBlock } from '../types/messages.js'
+import type { OffloadedBlockData, OffloadMethod } from '../types/messages.js'
+import type { Agent } from '../agent/agent.js'
+import type { ContextState } from '../context-manager/types.js'
+import type { Stash } from '../context-manager/stash.js'
+import { createMockAgent } from './agent-helpers.js'
+
+/** A user message holding one successful tool result with the given text. */
+export function makeToolResultMessage(text: string, toolUseId = 'tool-123'): Message {
+  return new Message({
+    role: 'user',
+    content: [new ToolResultBlock({ toolUseId, status: 'success', content: [new TextBlock(text)] })],
+  })
+}
+
+/** The offload method that left `block` as a placeholder, or undefined when it is not one. */
+export function offloadedBy(block: unknown): OffloadMethod | undefined {
+  return isOffloaded(block) ? block.metadata.offloaded.method : undefined
+}
+
+/** The metadata of every placeholder in `messages`, including those inside tool results, in history order. */
+export function placeholdersIn(messages: Message[]): OffloadedBlockData[] {
+  return messages
+    .flatMap((message) => message.content)
+    .flatMap((block): unknown[] => (block instanceof ToolResultBlock ? block.content : [block]))
+    .filter(isOffloaded)
+    .map((block) => block.metadata.offloaded)
+}
+
+/** Counts text at four characters per token and everything else at two characters per token of its JSON. */
+export function heuristicCountTokens(messages: Message[]): number {
+  let total = 0
+  for (const message of messages) {
+    for (const block of message.content) {
+      if (block instanceof TextBlock) {
+        total += Math.ceil(block.text.length / 4)
+      } else if (block instanceof ToolResultBlock) {
+        for (const content of block.content) {
+          if (content instanceof TextBlock) total += Math.ceil(content.text.length / 4)
+          else total += Math.ceil(JSON.stringify(content).length / 2)
+        }
+      } else {
+        total += Math.ceil(JSON.stringify(block).length / 2)
+      }
+    }
+  }
+  return total
+}
+
+export interface MakeContextOptions {
+  /** Stash the strategies may write to. */
+  stash?: Stash
+  /** Context window utilization reported to the strategies. Defaults to 0.5. */
+  utilization?: number
+  /** Whether the strategies run on an overflow. */
+  overflow?: boolean
+  /**
+   * Extra members of the agent's model, merged over {@link heuristicCountTokens}. `null` leaves the
+   * agent without a model.
+   */
+  model?: object | null
+}
+
+/** A {@link ContextState} over a mock agent whose model counts tokens with {@link heuristicCountTokens}. */
+export function makeContext(messages: Message[], options: MakeContextOptions = {}): ContextState {
+  const { stash, utilization = 0.5, overflow, model = {} } = options
+  const extra: Partial<Agent> =
+    model === null
+      ? {}
+      : ({ model: { countTokens: async (msgs: Message[]) => heuristicCountTokens(msgs), ...model } } as Partial<Agent>)
+  const agent = createMockAgent({ messages, extra })
+  return {
+    messages,
+    agent,
+    utilization,
+    ...(stash !== undefined && { stash }),
+    ...(overflow && { overflow: true }),
+  }
+}

strands-ts/src/a2a/__tests__/adapters.test.ts

@@ -102,6 +102,16 @@ describe('adapters', () => {
       ])
     })
 
+    it('converts an offloaded block to the text the model receives', () => {
+      const blocks: ContentBlock[] = [
+        new TextBlock('[Dropped] [ref: tu-1]', {
+          metadata: { offloaded: { method: 'drop', ref: 'tu-1', preview: '[Dropped]' } },
+        }),
+      ]
+
+      expect(contentBlocksToParts(blocks)).toStrictEqual([{ kind: 'text', text: '[Dropped] [ref: tu-1]' }])
+    })
+
     it.each([
       {
         desc: 'ImageBlock with bytes',

strands-ts/src/agent/__tests__/agent.context-manager.test.ts

@@ -1,9 +1,15 @@
-import { describe, expect, it } from 'vitest'
+import { describe, expect, it, vi } from 'vitest'
+import { z } from 'zod'
 import { Agent } from '../agent.js'
 import { MockMessageModel } from '../../__fixtures__/mock-message-model.js'
 import { SlidingWindowConversationManager } from '../../conversation-manager/sliding-window-conversation-manager.js'
 import { NullConversationManager } from '../../conversation-manager/null-conversation-manager.js'
 import { ContextManager } from '../../context-manager/context-manager.js'
+import { Offload } from '../../context-manager/strategies/offload/index.js'
+import { createRetrievalTool } from '../../context-manager/retrieval-tool.js'
+import { offloadPlaceholder } from '../../context-manager/placeholder.js'
+import { tool } from '../../tools/tool-factory.js'
+import { TextBlock, ToolResultBlock } from '../../types/messages.js'
 import type { ConversationManager } from '../../conversation-manager/conversation-manager.js'
 
 function internals(agent: Agent): any {
@@ -139,4 +145,41 @@ describe('Agent contextManager', () => {
       expect(plugins.get('strands:context-manager')).toBe(contextManager)
     })
   })
+
+  describe('offloaded tool results', () => {
+    it('stores an sends the model its text, and resolves the toolUseId to the original', async () => {
+      const rows = tool({
+        name: 'rows',
+        description: 'lists rows',
+        inputSchema: z.object({}),
+        callback: () => Array.from({ length: 5000 }, (_, index) => `row-${index + 1}`).join('\n'),
+      })
+      const model = new MockMessageModel()
+        .addTurn({ type: 'toolUseBlock', name: 'rows', toolUseId: 'tu-1', input: {} })
+        .addTurn({ type: 'textBlock', text: 'done' })
+      const sent = vi.spyOn(model, 'streamAggregated')
+      const agent = new Agent({
+        model,
+        tools: [rows],
+        printer: false,
+        contextManager: new ContextManager({
+          strategies: [Offload.truncate('toolResults', { previewTokens: 100 }).when({ blockThreshold: 200 })],
+        }),
+      })
+
+      await agent.invoke('list the rows')
+
+      const stored = (agent.messages[2]!.content[0] as ToolResultBlock).content[0] as TextBlock
+      expect(stored).toEqual(offloadPlaceholder('truncate', stored.metadata!.offloaded!.preview, 'tu-1'))
+      const received = (sent.mock.calls[1]![0][2]!.content[0] as ToolResultBlock).content[0]
+      expect(received).toBeInstanceOf(TextBlock)
+      expect((received as TextBlock).text).toBe(stored.text)
+
+      const retrieve = createRetrievalTool(agent.contextManager!.stash!) as unknown as {
+        invoke(input: unknown): Promise<unknown>
+      }
+      expect(await retrieve.invoke({ reference: 'tu-1', pattern: 'row-4321' })).toContain('row-4321')
+      expect(await retrieve.invoke({ reference: 'tu-1_0', pattern: 'row-4321' })).toContain('row-4321')
+    })
+  })
 })

strands-ts/src/agent/__tests__/agent.test.ts

@@ -1923,7 +1923,13 @@ describe('Agent', () => {
         new ToolResultBlock({
           toolUseId: 'tool-1',
           status: 'success',
-          content: [new TextBlock(expect.stringContaining('[Truncated:'))],
+          content: [
+            new TextBlock(expect.any(String), {
+              metadata: {
+                offloaded: { method: 'truncate', preview: expect.any(String), ref: 'tool-1' },
+              },
+            }),
+          ],
         }),
       ])
     })

strands-ts/src/context-manager/__tests__/context-manager.test.ts

@@ -127,7 +127,17 @@ describe('ContextManager', () => {
           new ToolResultBlock({
             toolUseId: 'tu-1',
             status: 'success',
-            content: [new TextBlock('[Truncated: 1 block, ~7,500 tokens] preview\n\n[Stashed] [ref: tu-1_0]')],
+            content: [
+              new TextBlock('[Truncated: 1 block, ~7,500 tokens] preview [ref: tu-1]', {
+                metadata: {
+                  offloaded: {
+                    method: 'truncate',
+                    ref: 'tu-1',
+                    preview: '[Truncated: 1 block, ~7,500 tokens] preview',
+                  },
+                },
+              }),
+            ],
           }),
         ],
       })

strands-ts/src/context-manager/__tests__/placeholder.test.ts

@@ -0,0 +1,58 @@
+import { describe, expect, it } from 'vitest'
+import {
+  formatDropped,
+  formatDroppedMessages,
+  formatElided,
+  formatOffloaded,
+  formatSummarized,
+  formatTruncated,
+  offloadPlaceholder,
+} from '../placeholder.js'
+import { TextBlock } from '../../types/messages.js'
+
+describe('offloadPlaceholder', () => {
+  it('appends the stash ref and records what the placeholder stands for', () => {
+    expect(offloadPlaceholder('drop', '[Dropped: 2 messages]', 'removed-abc', 2)).toStrictEqual(
+      new TextBlock('[Dropped: 2 messages] [ref: removed-abc]', {
+        metadata: {
+          offloaded: { method: 'drop', preview: '[Dropped: 2 messages]', ref: 'removed-abc', removedMessages: 2 },
+        },
+      })
+    )
+  })
+
+  it('leaves out the ref when the original was not stashed', () => {
+    expect(offloadPlaceholder('drop', '[Dropped]', undefined)).toStrictEqual(
+      new TextBlock('[Dropped]', { metadata: { offloaded: { method: 'drop', preview: '[Dropped]' } } })
+    )
+  })
+})
+
+describe('markers', () => {
+  it.each([
+    ['a dropped block', formatDropped(), '[Dropped]'],
+    ['dropped media', formatDropped('image'), '[Dropped: image, removed from context after it was viewed]'],
+    ['one dropped message', formatDroppedMessages(1), '[Dropped: 1 message]'],
+    ['dropped messages', formatDroppedMessages(3), '[Dropped: 3 messages]'],
+    ['an offloaded block', formatOffloaded(2525), '[Offloaded: ~2,525 tokens]'],
+    ['an unmeasured block', formatOffloaded(0), '[Offloaded]'],
+    [
+      'offloaded media',
+      formatOffloaded(2525, 'document'),
+      '[Offloaded: document, ~2,525 tokens, removed from context after it was viewed]',
+    ],
+    ['unmeasured media', formatOffloaded(0, 'image'), '[Offloaded: image, removed from context after it was viewed]'],
+    ['a summary', formatSummarized('tool result', 2525, 'short'), '[Summarized: tool result, ~2,525 tokens]\n\nshort'],
+    [
+      'an unmeasured summary',
+      formatSummarized('media block', 0, 'a red square'),
+      '[Summarized: media block]\n\na red square',
+    ],
+    ['a truncated block', formatTruncated(1, 7500, 'head'), '[Truncated: 1 block, ~7,500 tokens]\n\nhead'],
+    ['truncated blocks', formatTruncated(2, 7500, 'head'), '[Truncated: 2 blocks, ~7,500 tokens]\n\nhead'],
+    ['one elided message', formatElided(1), '[... 1 message elided ...]'],
+    ['elided messages', formatElided(7), '[... 7 messages elided ...]'],
+  ])('formats %s', (_case, text, expected) => {
+    expect(text).toBe(expected)
+  })
+})

strands-ts/src/context-manager/__tests__/presets.test.ts

@@ -1,6 +1,8 @@
 import { describe, it, expect } from 'vitest'
 import { resolvePreset, resolveStrategies, STRATEGY_PRESET_NAMES } from '../presets.js'
 import { ContextManager } from '../context-manager.js'
+import { Offload } from '../strategies/offload/index.js'
+import type { ContextStrategy } from '../types.js'
 
 describe('presets', () => {
   describe('resolvePreset', () => {
@@ -25,10 +27,10 @@ describe('presets', () => {
       expect(strategies[0]!.name).toBe('offload:truncate')
     })
 
-    it('overflowProtection resolves to a truncate strategy', () => {
+    it('overflowProtection resolves to a drop strategy', () => {
       const strategies = resolvePreset('overflowProtection')
       expect(strategies).toHaveLength(1)
-      expect(strategies[0]!.name).toBe('offload:truncate')
+      expect(strategies[0]!.name).toBe('offload:drop')
     })
 
     it('staleToolCleanup resolves to a drop strategy', () => {
@@ -57,7 +59,7 @@ describe('presets', () => {
       expect(result).toHaveLength(3)
       expect(result[0]!.name).toBe('offload:truncate')
       expect(result[1]!.name).toBe('custom')
-      expect(result[2]!.name).toBe('offload:truncate')
+      expect(result[2]!.name).toBe('offload:drop')
     })
 
     it('preserves order of preset expansion', () => {
@@ -72,6 +74,17 @@ describe('presets', () => {
         'Unknown strategy preset: "proactiveSummarisation"'
       )
     })
+
+    it.each([null, undefined])('rejects a %s entry', (entry) => {
+      expect(() => resolveStrategies([entry as unknown as ContextStrategy])).toThrow(
+        `Strategy entries must be a strategy or a preset name, got ${String(entry)}`
+      )
+    })
+
+    it('accepts an Offload strategy without .when()', () => {
+      const strategy = Offload.drop('toolResults')
+      expect(resolveStrategies([strategy])).toStrictEqual([strategy])
+    })
   })
 
   describe('ContextManager.from', () => {