feat(context): add message pinning to conversation managers

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": {

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')
+  })
+})

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<string, unknown>): 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)
+  })
+})

strands-ts/src/context-manager/compression/index.ts

@@ -0,0 +1 @@
+export { pinMessage, unpinMessage, isPinned, pinMessageTool } from './pin-message.js'

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}.`
+  },
+})