Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 39 additions & 4 deletions strands-ts/src/agent/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,8 @@ import { PluginRegistry } from '../plugins/registry.js'
import { SlidingWindowConversationManager } from '../conversation-manager/sliding-window-conversation-manager.js'
import { NullConversationManager } from '../conversation-manager/null-conversation-manager.js'
import { ConversationManager } from '../conversation-manager/conversation-manager.js'
import type { ContextManagerParam } from '../context-manager/context-manager.js'
import { resolveContextManager } from '../context-manager/context-manager.js'
import { HookRegistryImplementation } from '../hooks/registry.js'
import type { HookableEventConstructor, HookCallback, HookCallbackOptions, HookCleanup } from '../hooks/types.js'
import {
Expand Down Expand Up @@ -167,9 +169,24 @@ export type AgentConfig = {
* Defaults to true.
*/
printer?: boolean
/**
* Pre-composed context management strategy.
*
* - `"auto"`: enables tool result caching and proactive compression with defaults.
* - Object: fine-grained control over strategy, storage, caching, and compression settings.
* - `undefined` (default): no context management facade; use `conversationManager`
* and `plugins` directly.
*
* When set, takes priority over `conversationManager` — `NullConversationManager` is used.
*/
contextManager?: ContextManagerParam
Comment thread
lizradway marked this conversation as resolved.
/**
* Conversation manager for handling message history and context overflow.
* Defaults to SlidingWindowConversationManager with windowSize of 40.
*
* @remarks Pending deprecation — use `contextManager` instead. The `contextManager` parameter
Comment thread
lizradway marked this conversation as resolved.
* composes compression, tool result caching, and token estimation into a single
* configuration surface. This field will be deprecated in a future version.
*/
conversationManager?: ConversationManager
/**
Expand Down Expand Up @@ -331,15 +348,29 @@ export class Agent implements LocalAgent, InvokableAgent {
this.model = config?.model ?? new BedrockModel()
}

// Validate and assign conversation manager
let contextManagerPlugin: Plugin | undefined
if (config?.contextManager) {
contextManagerPlugin = resolveContextManager(config.contextManager, config.plugins)
}

// Validate and assign conversation manager.
// When contextManager is set, ContextCompression owns compression — use NullConversationManager.
if (this.model.stateful) {
if (config?.conversationManager) {
if (config?.conversationManager || config?.contextManager) {
throw new Error(
'Cannot use a conversationManager with a stateful model. The model manages conversation state server-side.'
'Cannot use a conversationManager or contextManager with a stateful model. The model manages conversation state server-side.'
)
}
this._conversationManager = new NullConversationManager()
} else if (contextManagerPlugin) {
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Issue: When a non-stateful model has both contextManager and conversationManager set, the conversationManager is silently ignored (line 365-366 takes priority). This could confuse users who set both accidentally.

Suggestion: Consider logging a warning when both are provided, e.g.:

} else if (contextManagerPlugin) {
  if (config?.conversationManager) {
    logger.warn('contextManager takes priority over conversationManager — conversationManager will be ignored')
  }
  this._conversationManager = new NullConversationManager()
}

if (config?.conversationManager) {
logger.warn('contextManager takes priority over conversationManager — conversationManager will be ignored')
}
this._conversationManager = new NullConversationManager()
} else {
if (config?.conversationManager) {
logger.warn('conversationManager is deprecated and will be removed in v2. Use contextManager instead.')
}
this._conversationManager =
config?.conversationManager ?? new SlidingWindowConversationManager({ windowSize: 40 })
}
Expand Down Expand Up @@ -372,9 +403,12 @@ export class Agent implements LocalAgent, InvokableAgent {
// - Retry-strategy ordering is not load-bearing for correctness: `DefaultModelRetryStrategy`
// guards on `event.retry`, so a user hook that already set it short-circuits
// the strategy regardless of registration order.
// - contextManager plugin goes before user plugins so the offloader's AfterToolCallEvent
// hook fires first, ensuring large results are cached before user hooks see the event.
this._pluginRegistry = new PluginRegistry([
this._conversationManager,
...retryStrategies,
...(contextManagerPlugin ? [contextManagerPlugin] : []),
...(config?.plugins ?? []),
...(config?.sessionManager ? [config.sessionManager] : []),
new ModelPlugin(this.model),
Expand Down Expand Up @@ -1397,7 +1431,8 @@ export class Agent implements LocalAgent, InvokableAgent {

let attemptCount = 1
while (true) {
// Estimate input tokens for the upcoming model call (non-fatal if estimation fails)
// Pending deprecation: token estimation will move fully to ContextManager.
// This remains for backward compat with standalone ConversationManager.proactiveCompression.
let projectedInputTokens: number | undefined
try {
projectedInputTokens = await this._estimateInputTokens(streamOptions)
Expand Down
164 changes: 164 additions & 0 deletions strands-ts/src/context-manager/__tests__/context-manager.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,164 @@
import { describe, it, expect } from 'vitest'
import { ContextManager, resolveContextManager } from '../context-manager.js'
import { ContextCompression } from '../compression/context-compression.js'
import { InMemoryStorage } from '../../vended-plugins/context-offloader/storage.js'
import { createMockAgent } from '../../__fixtures__/agent-helpers.js'
import type { Plugin } from '../../plugins/plugin.js'

describe('resolveContextManager', () => {
it('with "auto" enables both compression and offloader', () => {
const cm = resolveContextManager('auto')
const subPlugins = (cm as any)._subPlugins as Plugin[]

expect(subPlugins).toHaveLength(2)
const names = subPlugins.map((p) => p.name)
expect(names).toContain('strands:context-compression')
expect(names).toContain('strands:context-offloader')
})

it('with config object is additive (omitted = disabled)', () => {
const cm = resolveContextManager({ compression: true })
const subPlugins = (cm as any)._subPlugins as Plugin[]

const names = subPlugins.map((p) => p.name)
expect(names).toContain('strands:context-compression')
expect(names).not.toContain('strands:context-offloader')
})

it('with strategy: "auto" applies override semantics (omitted features stay enabled)', () => {
const cm = resolveContextManager({ strategy: 'auto', compression: 'summarize' })
const subPlugins = (cm as any)._subPlugins as Plugin[]

const names = subPlugins.map((p) => p.name)
// Both should be enabled because strategy: 'auto' defaults include offloader: true
expect(names).toContain('strands:context-compression')
expect(names).toContain('strands:context-offloader')

// Compression should use summarize method
const compression = subPlugins.find((p) => p.name === 'strands:context-compression') as ContextCompression
expect((compression as any)._method).toBe('summarize')
})

it('with offloader: true uses default thresholds', () => {
const cm = resolveContextManager({ offloader: true })
const subPlugins = (cm as any)._subPlugins as Plugin[]

const offloader = subPlugins.find((p) => p.name === 'strands:context-offloader')
expect(offloader).toBeDefined()
})

it('with offloader config applies custom settings', () => {
const cm = resolveContextManager({ offloader: { threshold: 5000, previewTokens: 1000 } })
const subPlugins = (cm as any)._subPlugins as Plugin[]

const offloader = subPlugins.find((p) => p.name === 'strands:context-offloader')
expect(offloader).toBeDefined()
})
})

describe('ContextManager._buildSubPlugins', () => {
it('skips compression plugin when user already provides one', () => {
const userCompression: Plugin = {
name: 'strands:context-compression',
initAgent: () => {},
getTools: () => [],
}

const cm = resolveContextManager('auto', [userCompression])
const subPlugins = (cm as any)._subPlugins as Plugin[]

const compressionPlugins = subPlugins.filter((p) => p.name === 'strands:context-compression')
expect(compressionPlugins).toHaveLength(0)
// Offloader should still be present
const offloaderPlugins = subPlugins.filter((p) => p.name === 'strands:context-offloader')
expect(offloaderPlugins).toHaveLength(1)
})

it('skips offloader plugin when user already provides one', () => {
const userOffloader: Plugin = {
name: 'strands:context-offloader',
initAgent: () => {},
getTools: () => [],
}

const cm = resolveContextManager('auto', [userOffloader])
const subPlugins = (cm as any)._subPlugins as Plugin[]

const offloaderPlugins = subPlugins.filter((p) => p.name === 'strands:context-offloader')
expect(offloaderPlugins).toHaveLength(0)
// Compression should still be present
const compressionPlugins = subPlugins.filter((p) => p.name === 'strands:context-compression')
expect(compressionPlugins).toHaveLength(1)
})
})

describe('ContextManager', () => {
describe('constructor', () => {
it('uses InMemoryStorage by default', () => {
const cm = new ContextManager()
expect(cm.storage).toBeInstanceOf(InMemoryStorage)
})

it('accepts custom storage', () => {
const storage = new InMemoryStorage()
const cm = new ContextManager({ storage })
expect(cm.storage).toBe(storage)
})

it('has correct plugin name', () => {
const cm = new ContextManager()
expect(cm.name).toBe('strands:context-manager')
})
})

describe('initAgent', () => {
it('initializes sub-plugins', () => {
const cm = new ContextManager({ compression: true, offloader: true })
const agent = createMockAgent()

cm.initAgent(agent)

// Should have registered hooks from both sub-plugins
expect(agent.trackedHooks.length).toBeGreaterThan(0)
})

it('builds sub-plugins if not already resolved', () => {
const cm = new ContextManager({ compression: true })
const agent = createMockAgent()

// Don't call _resolveSubPlugins first
cm.initAgent(agent)

// Should still work and register hooks
expect(agent.trackedHooks.length).toBeGreaterThan(0)
})
})

describe('getTools', () => {
it('returns tools from sub-plugins', () => {
const cm = new ContextManager({ offloader: true })
cm._resolveSubPlugins()

const tools = cm.getTools()
// ContextOffloader provides retrieval tool by default
expect(tools.length).toBeGreaterThan(0)
expect(tools[0]!.name).toBe('retrieve_offloaded_content')
})

it('returns empty array when no sub-plugins configured', () => {
const cm = new ContextManager({})
cm._resolveSubPlugins()

const tools = cm.getTools()
expect(tools).toHaveLength(0)
})

it('returns empty array when sub-plugins are not resolved yet', () => {
const cm = new ContextManager({ offloader: true })
// Don't resolve sub-plugins

const tools = cm.getTools()
expect(tools).toHaveLength(0)
})
})
})
104 changes: 104 additions & 0 deletions strands-ts/src/context-manager/__tests__/token-estimation.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
import { describe, it, expect, vi } from 'vitest'
import { estimateInputTokens } from '../token-estimation.js'
import { Message, TextBlock } from '../../types/messages.js'
import type { Model } from '../../models/model.js'

function userMsg(text: string): Message {
return new Message({ role: 'user', content: [new TextBlock(text)] })
}

function assistantMsg(
text: string,
usage?: { inputTokens: number; outputTokens: number; totalTokens: number }
): Message {
return new Message({
role: 'assistant',
content: [new TextBlock(text)],
...(usage && { metadata: { usage } }),
})
}

function mockModel(countTokens?: (messages: Message[]) => Promise<number>): Model {
return {
countTokens: countTokens ?? vi.fn().mockResolvedValue(100),
} as unknown as Model
}

describe('estimateInputTokens', () => {
it('returns baseline from last assistant message usage metadata', async () => {
const messages = [
userMsg('hello'),
assistantMsg('response', { inputTokens: 50, outputTokens: 20, totalTokens: 70 }),
]
const model = mockModel()

const result = await estimateInputTokens(messages, model)

expect(result).toBe(70) // 50 + 20
})

it('adds new message tokens to baseline when messages exist after the assistant message', async () => {
const messages = [
userMsg('hello'),
assistantMsg('response', { inputTokens: 50, outputTokens: 20, totalTokens: 70 }),
userMsg('follow up'),
]
const countTokens = vi.fn().mockResolvedValue(15)
const model = mockModel(countTokens)

const result = await estimateInputTokens(messages, model)

expect(result).toBe(85) // 70 + 15
expect(countTokens).toHaveBeenCalledWith([messages[2]])
})

it('uses the last assistant message with usage (not earlier ones)', async () => {
const messages = [
userMsg('hello'),
assistantMsg('first', { inputTokens: 10, outputTokens: 5, totalTokens: 15 }),
userMsg('second'),
assistantMsg('latest', { inputTokens: 80, outputTokens: 30, totalTokens: 110 }),
]
const model = mockModel()

const result = await estimateInputTokens(messages, model)

expect(result).toBe(110) // 80 + 30
})

it('falls back to model.countTokens when no assistant message has usage metadata', async () => {
const messages = [
userMsg('hello'),
new Message({ role: 'assistant', content: [new TextBlock('no metadata')] }),
userMsg('world'),
]
const countTokens = vi.fn().mockResolvedValue(42)
const model = mockModel(countTokens)

const result = await estimateInputTokens(messages, model)

expect(result).toBe(42)
expect(countTokens).toHaveBeenCalledWith(messages)
})

it('falls back to model.countTokens when there are no assistant messages', async () => {
const messages = [userMsg('hello')]
const countTokens = vi.fn().mockResolvedValue(10)
const model = mockModel(countTokens)

const result = await estimateInputTokens(messages, model)

expect(result).toBe(10)
expect(countTokens).toHaveBeenCalledWith(messages)
})

it('returns undefined on error', async () => {
const messages = [userMsg('hello')]
const countTokens = vi.fn().mockRejectedValue(new Error('API error'))
const model = mockModel(countTokens)

const result = await estimateInputTokens(messages, model)

expect(result).toBeUndefined()
})
})
Loading
Loading