From 900114929cc75ca19088c662277f6e9f2a641130 Mon Sep 17 00:00:00 2001 From: Liz <91279165+lizradway@users.noreply.github.com> Date: Tue, 26 May 2026 14:54:10 -0400 Subject: [PATCH] feat(context): add message pinning to conversation managers --- strands-ts/package.json | 8 + .../__tests__/pin-message-tool.test.ts | 66 ++++++++ .../compression/__tests__/pin.test.ts | 144 ++++++++++++++++++ .../src/context-manager/compression/index.ts | 1 + .../compression/pin-message.ts | 117 ++++++++++++++ ...liding-window-conversation-manager.test.ts | 77 ++++++++++ .../summarizing-conversation-manager.test.ts | 104 +++++++++++++ .../sliding-window-conversation-manager.ts | 43 +++++- .../summarizing-conversation-manager.ts | 41 ++++- 9 files changed, 592 insertions(+), 9 deletions(-) create mode 100644 strands-ts/src/context-manager/compression/__tests__/pin-message-tool.test.ts create mode 100644 strands-ts/src/context-manager/compression/__tests__/pin.test.ts create mode 100644 strands-ts/src/context-manager/compression/index.ts create mode 100644 strands-ts/src/context-manager/compression/pin-message.ts diff --git a/strands-ts/package.json b/strands-ts/package.json index 4060223f9..5d8ec0dad 100644 --- a/strands-ts/package.json +++ b/strands-ts/package.json @@ -95,6 +95,14 @@ "./vended-plugins": { "types": "./dist/src/vended-plugins/index.d.ts", "default": "./dist/src/vended-plugins/index.js" + }, + "./context-manager/compression": { + "types": "./dist/src/context-manager/compression/index.d.ts", + "default": "./dist/src/context-manager/compression/index.js" + }, + "./context-manager/compression/pin-message": { + "types": "./dist/src/context-manager/compression/pin-message.d.ts", + "default": "./dist/src/context-manager/compression/pin-message.js" } }, "scripts": { diff --git a/strands-ts/src/context-manager/compression/__tests__/pin-message-tool.test.ts b/strands-ts/src/context-manager/compression/__tests__/pin-message-tool.test.ts new file mode 100644 index 000000000..6c134cdf8 --- /dev/null +++ b/strands-ts/src/context-manager/compression/__tests__/pin-message-tool.test.ts @@ -0,0 +1,66 @@ +import { describe, it, expect } from 'vitest' +import { pinMessageTool, isPinned, pinMessage } from '../pin-message.js' +import { Message, TextBlock } from '../../../types/messages.js' +import type { Agent } from '../../../agent/agent.js' + +function makeAgent(messages: Message[]): Agent { + return { messages } as unknown as Agent +} + +function makeMessage(text: string): Message { + return new Message({ role: 'user', content: [new TextBlock(text)] }) +} + +describe('pinMessageTool', () => { + it('has the correct name and description', () => { + expect(pinMessageTool.name).toBe('pin_message') + expect(pinMessageTool.description).toContain('Pin or unpin') + }) + + it('pins a message at a valid index', async () => { + const messages = [makeMessage('first'), makeMessage('second'), makeMessage('third')] + const agent = makeAgent(messages) + + const result = await pinMessageTool.invoke({ index: 1, action: 'pin' }, { agent } as any) + + expect(result).toBe('Pinned message at index 1.') + expect(isPinned(agent.messages[1]!)).toBe(true) + expect(isPinned(agent.messages[0]!)).toBe(false) + }) + + it('defaults action to pin', async () => { + const messages = [makeMessage('first')] + const agent = makeAgent(messages) + + const result = await pinMessageTool.invoke({ index: 0 } as any, { agent } as any) + + expect(result).toBe('Pinned message at index 0.') + expect(isPinned(agent.messages[0]!)).toBe(true) + }) + + it('unpins a pinned message', async () => { + const messages = [pinMessage(makeMessage('pinned'))] + const agent = makeAgent(messages) + + expect(isPinned(agent.messages[0]!)).toBe(true) + + const result = await pinMessageTool.invoke({ index: 0, action: 'unpin' }, { agent } as any) + + expect(result).toBe('Unpinned message at index 0.') + expect(isPinned(agent.messages[0]!)).toBe(false) + }) + + it('rejects negative index via schema validation', async () => { + const agent = makeAgent([makeMessage('only')]) + + await expect(pinMessageTool.invoke({ index: -1, action: 'pin' }, { agent } as any)).rejects.toThrow() + }) + + it('returns error for out-of-bounds index', async () => { + const agent = makeAgent([makeMessage('only')]) + + const result = await pinMessageTool.invoke({ index: 5, action: 'pin' }, { agent } as any) + + expect(result).toContain('Invalid index 5') + }) +}) diff --git a/strands-ts/src/context-manager/compression/__tests__/pin.test.ts b/strands-ts/src/context-manager/compression/__tests__/pin.test.ts new file mode 100644 index 000000000..1ccff2931 --- /dev/null +++ b/strands-ts/src/context-manager/compression/__tests__/pin.test.ts @@ -0,0 +1,144 @@ +import { describe, it, expect } from 'vitest' +import { pinMessage, unpinMessage, isPinned } from '../pin-message.js' +import { Message, TextBlock, ToolUseBlock, ToolResultBlock } from '../../../types/messages.js' + +function makeMessage(text: string, metadata?: Record): Message { + return new Message({ + role: 'user', + content: [new TextBlock(text)], + ...(metadata !== undefined ? { metadata: metadata as any } : {}), + }) +} + +describe('isPinned', () => { + it('returns false for message without metadata', () => { + expect(isPinned(makeMessage('hello'))).toBe(false) + }) + + it('returns false for message with empty custom', () => { + expect(isPinned(makeMessage('hello', { custom: {} }))).toBe(false) + }) + + it('returns true for message with custom.pinned = true', () => { + expect(isPinned(makeMessage('hello', { custom: { pinned: true } }))).toBe(true) + }) + + it('returns false for message with custom.pinned = false', () => { + expect(isPinned(makeMessage('hello', { custom: { pinned: false } }))).toBe(false) + }) +}) + +describe('pinMessage', () => { + it('returns a new message with pinned = true in custom metadata', () => { + const original = makeMessage('important') + const pinned = pinMessage(original) + + expect(isPinned(pinned)).toBe(true) + expect(pinned.role).toBe('user') + expect(pinned.content).toEqual(original.content) + }) + + it('preserves existing metadata', () => { + const original = makeMessage('important', { usage: { inputTokens: 10, outputTokens: 5 } }) + const pinned = pinMessage(original) + + expect(pinned.metadata?.usage).toEqual({ inputTokens: 10, outputTokens: 5 }) + expect(isPinned(pinned)).toBe(true) + }) + + it('preserves existing custom fields', () => { + const original = makeMessage('important', { custom: { myField: 'value' } }) + const pinned = pinMessage(original) + + expect(pinned.metadata?.custom?.myField).toBe('value') + expect(isPinned(pinned)).toBe(true) + }) + + it('does not mutate the original message', () => { + const original = makeMessage('important') + pinMessage(original) + + expect(isPinned(original)).toBe(false) + }) +}) + +describe('unpinMessage', () => { + it('removes pinned from custom metadata', () => { + const pinned = pinMessage(makeMessage('important')) + const unpinned = unpinMessage(pinned) + + expect(isPinned(unpinned)).toBe(false) + }) + + it('preserves other custom fields', () => { + const original = makeMessage('important', { custom: { pinned: true, other: 'keep' } }) + const unpinned = unpinMessage(original) + + expect(isPinned(unpinned)).toBe(false) + expect(unpinned.metadata?.custom?.other).toBe('keep') + }) + + it('removes metadata entirely when nothing remains', () => { + const pinned = pinMessage(makeMessage('hello')) + const unpinned = unpinMessage(pinned) + + expect(unpinned.metadata).toBeUndefined() + }) + + it('preserves non-custom metadata fields', () => { + const original = makeMessage('important', { usage: { inputTokens: 10, outputTokens: 5 }, custom: { pinned: true } }) + const unpinned = unpinMessage(original) + + expect(unpinned.metadata?.usage).toEqual({ inputTokens: 10, outputTokens: 5 }) + expect(isPinned(unpinned)).toBe(false) + }) +}) + +describe('isPinned', () => { + it('returns false for unpinned message', () => { + const messages = [makeMessage('a'), makeMessage('b')] + expect(isPinned(messages, 0)).toBe(false) + }) + + it('returns true for pinned message', () => { + const messages = [pinMessage(makeMessage('a')), makeMessage('b')] + expect(isPinned(messages, 0)).toBe(true) + }) + + it('returns true for toolResult whose toolUse partner is pinned', () => { + const toolUseMsg = pinMessage( + new Message({ + role: 'assistant', + content: [new ToolUseBlock({ toolUseId: 'id-1', name: 'test', input: {} })], + }) + ) + const toolResultMsg = new Message({ + role: 'user', + content: [new ToolResultBlock({ toolUseId: 'id-1', content: [new TextBlock('result')], status: 'success' })], + }) + const messages = [toolUseMsg, toolResultMsg, makeMessage('other')] + + expect(isPinned(messages, 1)).toBe(true) + }) + + it('returns true for toolUse whose toolResult partner is pinned', () => { + const toolUseMsg = new Message({ + role: 'assistant', + content: [new ToolUseBlock({ toolUseId: 'id-1', name: 'test', input: {} })], + }) + const toolResultMsg = pinMessage( + new Message({ + role: 'user', + content: [new ToolResultBlock({ toolUseId: 'id-1', content: [new TextBlock('result')], status: 'success' })], + }) + ) + const messages = [toolUseMsg, toolResultMsg, makeMessage('other')] + + expect(isPinned(messages, 0)).toBe(true) + }) + + it('returns false for unrelated message next to pinned', () => { + const messages = [pinMessage(makeMessage('a')), makeMessage('b')] + expect(isPinned(messages, 1)).toBe(false) + }) +}) diff --git a/strands-ts/src/context-manager/compression/index.ts b/strands-ts/src/context-manager/compression/index.ts new file mode 100644 index 000000000..15638e894 --- /dev/null +++ b/strands-ts/src/context-manager/compression/index.ts @@ -0,0 +1 @@ +export { pinMessage, unpinMessage, isPinned, pinMessageTool } from './pin-message.js' diff --git a/strands-ts/src/context-manager/compression/pin-message.ts b/strands-ts/src/context-manager/compression/pin-message.ts new file mode 100644 index 000000000..ac077eae3 --- /dev/null +++ b/strands-ts/src/context-manager/compression/pin-message.ts @@ -0,0 +1,117 @@ +import { z } from 'zod' +import { Message, type ToolUseBlock, type ToolResultBlock } from '../../types/messages.js' +import { tool } from '../../tools/tool-factory.js' + +/** + * Check if a single message is pinned. + * + * @param message - The message to check + * @returns `true` if the message has `metadata.custom.pinned === true` + */ +export function isPinned(message: Message): boolean +/** + * Check if a message is pinned, including tool-pair partner protection. + * Returns `true` if the message at `index` is pinned, or if it is the + * adjacent tool-pair partner (toolUse/toolResult) of a pinned message, + * matched by toolUseId. + * + * @param messages - The full messages array + * @param index - The index to check + * @returns `true` if the message or its tool-pair partner is pinned + */ +export function isPinned(messages: Message[], index: number): boolean +export function isPinned(messageOrMessages: Message | Message[], index?: number): boolean { + if (index === undefined) { + return (messageOrMessages as Message).metadata?.custom?.pinned === true + } + + const messages = messageOrMessages as Message[] + const msg = messages[index]! + if (msg.metadata?.custom?.pinned === true) return true + + const toolResultBlocks = msg.content.filter((b): b is ToolResultBlock => b.type === 'toolResultBlock') + if (toolResultBlocks.length > 0 && index > 0) { + const prev = messages[index - 1]! + if (prev.metadata?.custom?.pinned === true) { + const resultIds = new Set(toolResultBlocks.map((b) => b.toolUseId)) + if (prev.content.some((b) => b.type === 'toolUseBlock' && resultIds.has((b as ToolUseBlock).toolUseId))) { + return true + } + } + } + + const toolUseBlocks = msg.content.filter((b): b is ToolUseBlock => b.type === 'toolUseBlock') + if (toolUseBlocks.length > 0 && index + 1 < messages.length) { + const next = messages[index + 1]! + if (next.metadata?.custom?.pinned === true) { + const useIds = new Set(toolUseBlocks.map((b) => b.toolUseId)) + if (next.content.some((b) => b.type === 'toolResultBlock' && useIds.has((b as ToolResultBlock).toolUseId))) { + return true + } + } + } + + return false +} + +/** + * Returns a new Message marked as pinned (protected from eviction during context reduction). + * + * @param message - The message to pin + * @returns A new Message with `metadata.custom.pinned` set to `true` + */ +export function pinMessage(message: Message): Message { + return new Message({ + role: message.role, + content: message.content, + metadata: { + ...message.metadata, + custom: { ...message.metadata?.custom, pinned: true }, + }, + }) +} + +/** + * Returns a new Message with pinning removed. + * + * @param message - The message to unpin + * @returns A new Message without the `pinned` flag in `metadata.custom` + */ +export function unpinMessage(message: Message): Message { + const { pinned: _, ...restCustom } = message.metadata?.custom ?? {} + const { custom: __, ...restMetadata } = message.metadata ?? {} + const hasCustom = Object.keys(restCustom).length > 0 + const hasMetadata = hasCustom || Object.keys(restMetadata).length > 0 + const metadata = hasMetadata ? { ...restMetadata, ...(hasCustom ? { custom: restCustom } : {}) } : undefined + + return new Message({ + role: message.role, + content: message.content, + ...(metadata !== undefined ? { metadata } : {}), + }) +} + +/** + * Agent-invokable tool that pins or unpins a message in the conversation history. + * When added to an agent's tools array, allows the agent to protect important + * messages from eviction during context reduction. + */ +export const pinMessageTool = tool({ + name: 'pin_message', + description: + 'Pin or unpin a message in the conversation history. ' + + 'Pinned messages are protected from eviction during context reduction. ' + + 'Use this to preserve important context that should not be summarized or trimmed away.', + inputSchema: z.object({ + index: z.number().int().min(0).describe('The zero-based index of the message in the conversation history.'), + action: z.enum(['pin', 'unpin']).default('pin').describe('Whether to pin or unpin the message.'), + }), + callback: ({ index, action }, context) => { + const messages = context!.agent.messages + if (index >= messages.length) { + return `Invalid index ${index}. Conversation has ${messages.length} messages (indices 0-${messages.length - 1}).` + } + messages[index] = action === 'pin' ? pinMessage(messages[index]!) : unpinMessage(messages[index]!) + return `${action === 'pin' ? 'Pinned' : 'Unpinned'} message at index ${index}.` + }, +}) diff --git a/strands-ts/src/conversation-manager/__tests__/sliding-window-conversation-manager.test.ts b/strands-ts/src/conversation-manager/__tests__/sliding-window-conversation-manager.test.ts index 836f2c66f..914b7cf84 100644 --- a/strands-ts/src/conversation-manager/__tests__/sliding-window-conversation-manager.test.ts +++ b/strands-ts/src/conversation-manager/__tests__/sliding-window-conversation-manager.test.ts @@ -1211,4 +1211,81 @@ describe('SlidingWindowConversationManager', () => { expect(mockAgent.messages).toHaveLength(2) }) }) + + describe('protectedMessageRange', () => { + it('protects first N messages from trimming (positive)', async () => { + const manager = new SlidingWindowConversationManager({ windowSize: 4, protectedMessageRange: 2 }) + const mockAgent = createMockAgent({ + messages: [ + new Message({ role: 'user', content: [new TextBlock('first')] }), + new Message({ role: 'assistant', content: [new TextBlock('second')] }), + new Message({ role: 'user', content: [new TextBlock('third')] }), + new Message({ role: 'assistant', content: [new TextBlock('fourth')] }), + new Message({ role: 'user', content: [new TextBlock('fifth')] }), + new Message({ role: 'assistant', content: [new TextBlock('sixth')] }), + ], + }) as unknown as Agent + + await triggerSlidingWindow(manager, mockAgent) + + const texts = mockAgent.messages.map((m) => (m.content[0] as TextBlock).text) + expect(texts).toContain('first') + expect(texts).toContain('second') + }) + + it('protects last N messages from trimming (negative)', async () => { + const manager = new SlidingWindowConversationManager({ windowSize: 4, protectedMessageRange: -2 }) + const mockAgent = createMockAgent({ + messages: [ + new Message({ role: 'user', content: [new TextBlock('first')] }), + new Message({ role: 'assistant', content: [new TextBlock('second')] }), + new Message({ role: 'user', content: [new TextBlock('third')] }), + new Message({ role: 'assistant', content: [new TextBlock('fourth')] }), + new Message({ role: 'user', content: [new TextBlock('fifth')] }), + new Message({ role: 'assistant', content: [new TextBlock('sixth')] }), + ], + }) as unknown as Agent + + await triggerSlidingWindow(manager, mockAgent) + + const texts = mockAgent.messages.map((m) => (m.content[0] as TextBlock).text) + expect(texts).toContain('fifth') + expect(texts).toContain('sixth') + }) + + it('returns false when all messages in trim range are protected', () => { + const manager = new SlidingWindowConversationManager({ windowSize: 2, protectedMessageRange: 4 }) + const mockAgent = createMockAgent({ + messages: [ + new Message({ role: 'user', content: [new TextBlock('a')] }), + new Message({ role: 'assistant', content: [new TextBlock('b')] }), + new Message({ role: 'user', content: [new TextBlock('c')] }), + new Message({ role: 'assistant', content: [new TextBlock('d')] }), + ], + }) as unknown as Agent + + const result = manager.reduce({ agent: mockAgent, model: {} as any }) + expect(result).toBe(false) + }) + + it('pinned message in middle of window survives trimming', async () => { + const { pinMessage } = await import('../../context-manager/compression/pin-message.js') + const manager = new SlidingWindowConversationManager({ windowSize: 4 }) + const mockAgent = createMockAgent({ + messages: [ + new Message({ role: 'user', content: [new TextBlock('first')] }), + new Message({ role: 'assistant', content: [new TextBlock('second')] }), + pinMessage(new Message({ role: 'user', content: [new TextBlock('pinned-middle')] })), + new Message({ role: 'assistant', content: [new TextBlock('fourth')] }), + new Message({ role: 'user', content: [new TextBlock('fifth')] }), + new Message({ role: 'assistant', content: [new TextBlock('sixth')] }), + ], + }) as unknown as Agent + + await triggerSlidingWindow(manager, mockAgent) + + const texts = mockAgent.messages.map((m) => (m.content[0] as TextBlock).text) + expect(texts).toContain('pinned-middle') + }) + }) }) diff --git a/strands-ts/src/conversation-manager/__tests__/summarizing-conversation-manager.test.ts b/strands-ts/src/conversation-manager/__tests__/summarizing-conversation-manager.test.ts index cf2ea0794..ef09f44a0 100644 --- a/strands-ts/src/conversation-manager/__tests__/summarizing-conversation-manager.test.ts +++ b/strands-ts/src/conversation-manager/__tests__/summarizing-conversation-manager.test.ts @@ -406,4 +406,108 @@ describe('SummarizingConversationManager', () => { expect(mockAgent.messages).toHaveLength(20) }) }) + + describe('protectedMessageRange', () => { + it('protects first N messages from summarization (positive)', async () => { + const model = new MockMessageModel() + model.addTurn({ type: 'textBlock', text: 'Summary' }) + + const manager = new SummarizingConversationManager({ + summaryRatio: 0.5, + preserveRecentMessages: 2, + protectedMessageRange: 2, + }) + + const agent = createMockAgent({ + messages: [ + textMsg('user', 'protected-1'), + textMsg('assistant', 'protected-2'), + textMsg('user', 'summarize-me'), + textMsg('assistant', 'summarize-me-too'), + textMsg('user', 'recent-1'), + textMsg('assistant', 'recent-2'), + ], + }) + + await manager.reduce({ agent, model }) + + const texts = agent.messages.map((m) => (m.content[0] as TextBlock).text) + expect(texts).toContain('protected-1') + expect(texts).toContain('protected-2') + expect(texts).toContain('Summary') + expect(texts).not.toContain('summarize-me') + }) + + it('protects last N messages from summarization (negative)', async () => { + const model = new MockMessageModel() + model.addTurn({ type: 'textBlock', text: 'Summary' }) + + const manager = new SummarizingConversationManager({ + summaryRatio: 0.8, + preserveRecentMessages: 0, + protectedMessageRange: -2, + }) + + const agent = createMockAgent({ + messages: [ + textMsg('user', 'old-1'), + textMsg('assistant', 'old-2'), + textMsg('user', 'old-3'), + textMsg('assistant', 'old-4'), + textMsg('user', 'protected-last-1'), + textMsg('assistant', 'protected-last-2'), + ], + }) + + await manager.reduce({ agent, model }) + + const texts = agent.messages.map((m) => (m.content[0] as TextBlock).text) + expect(texts).toContain('protected-last-1') + expect(texts).toContain('protected-last-2') + expect(texts).toContain('Summary') + }) + + it('returns false when all messages in summary range are protected', async () => { + const model = new MockMessageModel() + model.addTurn({ type: 'textBlock', text: 'Summary' }) + + const manager = new SummarizingConversationManager({ + summaryRatio: 0.3, + preserveRecentMessages: 2, + protectedMessageRange: 10, + }) + + const agent = createMockAgent({ messages: makeMessages(6) }) + const result = await manager.reduce({ agent, model }) + expect(result).toBe(false) + }) + + it('pinned message in middle survives summarization', async () => { + const { pinMessage } = await import('../../context-manager/compression/pin-message.js') + const model = new MockMessageModel() + model.addTurn({ type: 'textBlock', text: 'Summary' }) + + const manager = new SummarizingConversationManager({ + summaryRatio: 0.5, + preserveRecentMessages: 2, + }) + + const agent = createMockAgent({ + messages: [ + textMsg('user', 'old-1'), + pinMessage(textMsg('assistant', 'pinned-middle')), + textMsg('user', 'old-3'), + textMsg('assistant', 'old-4'), + textMsg('user', 'recent-1'), + textMsg('assistant', 'recent-2'), + ], + }) + + await manager.reduce({ agent, model }) + + const texts = agent.messages.map((m) => (m.content[0] as TextBlock).text) + expect(texts).toContain('pinned-middle') + expect(texts).toContain('Summary') + }) + }) }) diff --git a/strands-ts/src/conversation-manager/sliding-window-conversation-manager.ts b/strands-ts/src/conversation-manager/sliding-window-conversation-manager.ts index a32f4282b..cf7dac49a 100644 --- a/strands-ts/src/conversation-manager/sliding-window-conversation-manager.ts +++ b/strands-ts/src/conversation-manager/sliding-window-conversation-manager.ts @@ -14,6 +14,7 @@ import { type ProactiveCompressionConfig, type ConversationManagerReduceOptions, } from './conversation-manager.js' +import { isPinned } from '../context-manager/compression/pin-message.js' import { logger } from '../logging/logger.js' const PRESERVE_CHARS = 200 @@ -103,6 +104,14 @@ export type SlidingWindowConversationManagerConfig = { * - `false` or omitted: disabled, only reactive overflow recovery is used. */ proactiveCompression?: boolean | ProactiveCompressionConfig + + /** + * Protect messages from eviction during reduction. + * Positive values protect the first N messages; negative values protect the last N. + * + * For agent-controlled pinning, add `pinMessageTool` to the agent's tools array. + */ + protectedMessageRange?: number } /** @@ -121,6 +130,7 @@ export type SlidingWindowConversationManagerConfig = { export class SlidingWindowConversationManager extends ConversationManager { private readonly _windowSize: number private readonly _shouldTruncateResults: boolean + private readonly _protectedMessageRange: number | undefined /** * Unique identifier for this conversation manager. @@ -136,6 +146,7 @@ export class SlidingWindowConversationManager extends ConversationManager { super(config) this._windowSize = config?.windowSize ?? 40 this._shouldTruncateResults = config?.shouldTruncateResults ?? true + this._protectedMessageRange = config?.protectedMessageRange } /** @@ -264,8 +275,24 @@ export class SlidingWindowConversationManager extends ConversationManager { return false } - // trimIndex is guaranteed to be < messages.length here, so splice always removes at least one message - messages.splice(0, trimIndex) + // Collect non-protected indices in [0, trimIndex) to remove + const indicesToRemove: number[] = [] + for (let i = 0; i < trimIndex; i++) { + if (this._isProtected(messages, i)) continue + indicesToRemove.push(i) + } + + if (indicesToRemove.length === 0) { + logger.warn( + `window_size=<${this._windowSize}>, messages=<${messages.length}> | all messages in trim range are protected, unable to reduce` + ) + return false + } + + // Remove in reverse order to keep indices stable + for (let i = indicesToRemove.length - 1; i >= 0; i--) { + messages.splice(indicesToRemove[i]!, 1) + } return true } @@ -449,10 +476,9 @@ export class SlidingWindowConversationManager extends ConversationManager { */ private _findOldestMessageWithToolResults(messages: Message[]): number | undefined { for (let idx = 0; idx < messages.length; idx++) { - const currentMessage = messages[idx]! - - const hasToolResult = currentMessage.content.some((block) => block.type === 'toolResultBlock') + if (this._isProtected(messages, idx)) continue + const hasToolResult = messages[idx]!.content.some((block) => block.type === 'toolResultBlock') if (hasToolResult) { return idx } @@ -460,4 +486,11 @@ export class SlidingWindowConversationManager extends ConversationManager { return undefined } + + private _isProtected(messages: Message[], index: number): boolean { + if (isPinned(messages, index)) return true + if (this._protectedMessageRange === undefined || this._protectedMessageRange === 0) return false + if (this._protectedMessageRange > 0) return index < this._protectedMessageRange + return index >= messages.length + this._protectedMessageRange + } } diff --git a/strands-ts/src/conversation-manager/summarizing-conversation-manager.ts b/strands-ts/src/conversation-manager/summarizing-conversation-manager.ts index 710311a9a..af80717fe 100644 --- a/strands-ts/src/conversation-manager/summarizing-conversation-manager.ts +++ b/strands-ts/src/conversation-manager/summarizing-conversation-manager.ts @@ -13,6 +13,7 @@ import { type ProactiveCompressionConfig, type ConversationManagerReduceOptions, } from './conversation-manager.js' +import { isPinned } from '../context-manager/compression/pin-message.js' import { logger } from '../logging/logger.js' import { normalizeError } from '../errors.js' import type { Model } from '../models/model.js' @@ -83,6 +84,15 @@ export type SummarizingConversationManagerConfig = { * - `false` or omitted: disabled, only reactive overflow recovery is used. */ proactiveCompression?: boolean | ProactiveCompressionConfig + + /** + * Protect messages from eviction during summarization. + * Positive values protect the first N messages; negative values protect the last N. + * + * Protected messages are compacted to the front of the array before the summary message. + * For agent-controlled pinning, add `pinMessageTool` to the agent's tools array. + */ + protectedMessageRange?: number } /** @@ -99,6 +109,7 @@ export class SummarizingConversationManager extends ConversationManager { private readonly _summaryRatio: number private readonly _preserveRecentMessages: number private readonly _summarizationSystemPrompt: string + private readonly _protectedMessageRange: number | undefined constructor(config?: SummarizingConversationManagerConfig) { super(config) @@ -107,6 +118,7 @@ export class SummarizingConversationManager extends ConversationManager { this._summaryRatio = Math.max(0.1, Math.min(0.8, config?.summaryRatio ?? 0.3)) this._preserveRecentMessages = config?.preserveRecentMessages ?? 10 this._summarizationSystemPrompt = config?.summarizationSystemPrompt ?? DEFAULT_SUMMARIZATION_PROMPT + this._protectedMessageRange = config?.protectedMessageRange } /** @@ -164,13 +176,27 @@ export class SummarizingConversationManager extends ConversationManager { // Adjust split point to avoid breaking tool use/result pairs messagesToSummarizeCount = this._adjustSplitPointForToolPairs(messages, messagesToSummarizeCount) - const messagesToSummarize = messages.slice(0, messagesToSummarizeCount) + // Partition [0, messagesToSummarizeCount) into protected (preserve) and non-protected (summarize) + const protectedToPreserve: Message[] = [] + const toSummarize: Message[] = [] + for (let i = 0; i < messagesToSummarizeCount; i++) { + if (this._isProtected(messages, i)) { + protectedToPreserve.push(messages[i]!) + } else { + toSummarize.push(messages[i]!) + } + } + + if (toSummarize.length === 0) { + logger.warn(`messages=<${messages.length}> | all messages in summarize range are protected, unable to reduce`) + return false + } // Generate summary via model call - const summaryMessage = await this._generateSummary(messagesToSummarize, model) + const summaryMessage = await this._generateSummary(toSummarize, model) - // Replace summarized messages with the summary - messages.splice(0, messagesToSummarizeCount, summaryMessage) + // Replace summarized range with protected messages + summary + messages.splice(0, messagesToSummarizeCount, ...protectedToPreserve, summaryMessage) return true } @@ -260,4 +286,11 @@ export class SummarizingConversationManager extends ConversationManager { return splitPoint } + + private _isProtected(messages: Message[], index: number): boolean { + if (isPinned(messages, index)) return true + if (this._protectedMessageRange === undefined || this._protectedMessageRange === 0) return false + if (this._protectedMessageRange > 0) return index < this._protectedMessageRange + return index >= messages.length + this._protectedMessageRange + } }