Skip to content

Commit f43172a

Browse files
Re-implement conversation-manager as a hook provider (#222)
* refactor: simplify test helpers and use createMockAgent - Use hooks.addHook() instead of manual registerCallbacks() - Create createMockAgent() helper in fixtures - Simplify test helpers to use invokeCallbacks directly - Update all tests to use createMockAgent helper - Make all tests properly async with await Addresses PR review feedback. refactor: keep conversationManager as first-class config option - Restore conversationManager to AgentConfig and Agent class - Agent automatically registers conversation manager hooks internally - Update tests to use HookRegistryImplementation and focus on behavior - Remove public mentions of hooks from conversation manager docs - Simplify NullConversationManager tests to focus on behavior This maintains conversationManager as a top-level concept while using hooks internally as an implementation detail. refactor!: convert conversation managers to hook providers BREAKING CHANGE: Removed ConversationManager base class and Agent.conversationManager field. Conversation managers are now hook providers that must be registered via the hooks config. Changes: - Added retryModelCall field to AfterModelCallEvent for retry control - Converted NullConversationManager to HookProvider (0 hooks registered) - Converted SlidingWindowConversationManager to HookProvider with AfterInvocationEvent and AfterModelCallEvent callbacks - Removed Agent.conversationManager field and related config - Updated Agent retry logic to use retryModelCall flag instead of catching ContextWindowOverflowError - Added messages field to AgentData interface for hook access - Deleted ConversationManager base class and ConversationContext interface - Updated public API exports to remove ConversationManager Migration: Before: new Agent({ conversationManager: new SlidingWindowConversationManager() }) After: new Agent({ hooks: [new SlidingWindowConversationManager()] }) Resolves: #210 * refactor: address PR review feedback - Create MockAgentData interface with optional messages and state - Update createMockAgent to accept interface-based parameter - Update all test usages to use object syntax - Add retryModelCall tests to agent.hook.test.ts - Update test helpers to use registry.addHook(manager) Addresses review comments on PR #222 --------- Co-authored-by: Strands Agent <[email protected]>
1 parent 76afd8f commit f43172a

File tree

18 files changed

+268
-235
lines changed

18 files changed

+268
-235
lines changed

src/__fixtures__/agent-helpers.ts

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
/**
2+
* Test fixtures and helpers for Agent testing.
3+
* This module provides utilities for testing Agent-related implementations.
4+
*/
5+
6+
import type { Agent } from '../agent/agent.js'
7+
import type { Message } from '../types/messages.js'
8+
import { AgentState } from '../agent/state.js'
9+
import type { JSONValue } from '../types/json.js'
10+
11+
/**
12+
* Data for creating a mock Agent.
13+
*/
14+
export interface MockAgentData {
15+
/**
16+
* Messages for the agent.
17+
*/
18+
messages?: Message[]
19+
/**
20+
* Initial state for the agent.
21+
*/
22+
state?: Record<string, JSONValue>
23+
}
24+
25+
/**
26+
* Helper to create a mock Agent for testing.
27+
* Provides minimal Agent interface with messages and state.
28+
*
29+
* @param data - Optional mock agent data
30+
* @returns Mock Agent object
31+
*/
32+
export function createMockAgent(data?: MockAgentData): Agent {
33+
return {
34+
messages: data?.messages ?? [],
35+
state: new AgentState(data?.state ?? {}),
36+
} as unknown as Agent
37+
}

src/__fixtures__/tool-helpers.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ export function createMockContext(
2323
toolUse,
2424
agent: {
2525
state: new AgentState(agentState),
26+
messages: [],
2627
},
2728
}
2829
}

src/agent/__tests__/agent.hook.test.ts

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import {
99
BeforeToolCallEvent,
1010
MessageAddedEvent,
1111
ModelStreamEventHook,
12+
type HookRegistry,
1213
} from '../../hooks/index.js'
1314
import { MockMessageModel } from '../../__fixtures__/mock-message-model.js'
1415
import { MockHookProvider } from '../../__fixtures__/mock-hook-provider.js'
@@ -301,4 +302,37 @@ describe('Agent Hooks Integration', () => {
301302
)
302303
})
303304
})
305+
306+
describe('AfterModelCallEvent retryModelCall', () => {
307+
it('retries model call when hook sets retryModelCall', async () => {
308+
let callCount = 0
309+
const retryHook = {
310+
registerCallbacks: (registry: HookRegistry) => {
311+
registry.addCallback(AfterModelCallEvent, (event: AfterModelCallEvent) => {
312+
callCount++
313+
if (callCount === 1 && event.error) {
314+
event.retryModelCall = true
315+
}
316+
})
317+
},
318+
}
319+
320+
const model = new MockMessageModel()
321+
.addTurn(new Error('First attempt failed'))
322+
.addTurn({ type: 'textBlock', text: 'Success after retry' })
323+
324+
const agent = new Agent({ model, hooks: [retryHook] })
325+
const result = await agent.invoke('Test')
326+
327+
expect(result.lastMessage.content[0]).toEqual({ type: 'textBlock', text: 'Success after retry' })
328+
expect(callCount).toBe(2)
329+
})
330+
331+
it('does not retry when retryModelCall is not set', async () => {
332+
const model = new MockMessageModel().addTurn(new Error('Failure'))
333+
const agent = new Agent({ model })
334+
335+
await expect(agent.invoke('Test')).rejects.toThrow('Failure')
336+
})
337+
})
304338
})

src/agent/agent.ts

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,15 @@ import {
1212
ToolResultBlock,
1313
type ToolUseBlock,
1414
} from '../index.js'
15-
import { normalizeError, ConcurrentInvocationError, MaxTokensError, ContextWindowOverflowError } from '../errors.js'
15+
import { normalizeError, ConcurrentInvocationError, MaxTokensError } from '../errors.js'
1616
import type { BaseModelConfig, Model, StreamOptions } from '../models/model.js'
1717
import { ToolRegistry } from '../registry/tool-registry.js'
1818
import { AgentState } from './state.js'
1919
import type { AgentData } from '../types/agent.js'
2020
import { AgentPrinter, getDefaultAppender, type Printer } from './printer.js'
21-
import type { ConversationManager } from '../conversation-manager/conversation-manager.js'
21+
import type { HookProvider } from '../hooks/types.js'
2222
import { SlidingWindowConversationManager } from '../conversation-manager/sliding-window-conversation-manager.js'
2323
import { HookRegistryImplementation } from '../hooks/registry.js'
24-
import type { HookProvider } from '../hooks/types.js'
2524
import {
2625
AfterInvocationEvent,
2726
AfterModelCallEvent,
@@ -74,7 +73,7 @@ export type AgentConfig = {
7473
* Conversation manager for handling message history and context overflow.
7574
* Defaults to SlidingWindowConversationManager with windowSize of 40.
7675
*/
77-
conversationManager?: ConversationManager
76+
conversationManager?: HookProvider
7877
/**
7978
* Hook providers to register with the agent.
8079
* Hooks enable observing and extending agent behavior.
@@ -107,7 +106,7 @@ export class Agent implements AgentData {
107106
/**
108107
* Conversation manager for handling message history and context overflow.
109108
*/
110-
public readonly conversationManager: ConversationManager
109+
public readonly conversationManager: HookProvider
111110

112111
private _isInvoking: boolean = false
113112
private _printer?: Printer
@@ -140,10 +139,12 @@ export class Agent implements AgentData {
140139

141140
this.state = new AgentState(config?.state)
142141

142+
// Initialize conversation manager
143143
this.conversationManager = config?.conversationManager ?? new SlidingWindowConversationManager({ windowSize: 40 })
144144

145-
// Initialize hooks
145+
// Initialize hooks and register conversation manager hooks
146146
this.hooks = new HookRegistryImplementation()
147+
this.hooks.addHook(this.conversationManager)
147148
this.hooks.addAllHooks(config?.hooks ?? [])
148149

149150
// Create printer if printer is enabled (default: true)
@@ -283,8 +284,6 @@ export class Agent implements AgentData {
283284
// Continue loop
284285
}
285286
} finally {
286-
this.conversationManager.applyManagement(this)
287-
288287
// Invoke AfterInvocationEvent hook
289288
await this.hooks.invokeCallbacks(new AfterInvocationEvent({ agent: this }))
290289

@@ -362,14 +361,14 @@ export class Agent implements AgentData {
362361
const modelError = normalizeError(error)
363362

364363
// Invoke AfterModelCallEvent hook even on error
365-
await this.hooks.invokeCallbacks(new AfterModelCallEvent({ agent: this, error: modelError }))
364+
const event = await this.hooks.invokeCallbacks(new AfterModelCallEvent({ agent: this, error: modelError }))
366365

367-
if (error instanceof ContextWindowOverflowError) {
368-
// Reduce context and retry
369-
this.conversationManager.reduceContext(this, error)
366+
// Check if hooks request a retry (e.g., after reducing context)
367+
if (event.retryModelCall) {
370368
return yield* this.invokeModel(args)
371369
}
372-
// Re-throw other errors
370+
371+
// Re-throw error
373372
throw error
374373
}
375374
}

src/conversation-manager/__tests__/conversation-manager.test.ts

Lines changed: 0 additions & 11 deletions
This file was deleted.
Lines changed: 19 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,53 +1,42 @@
11
import { describe, it, expect } from 'vitest'
22
import { NullConversationManager } from '../null-conversation-manager.js'
3-
import { ContextWindowOverflowError, Message, TextBlock } from '../../index.js'
4-
import type { Agent } from '../../agent/agent.js'
3+
import { Message, TextBlock } from '../../index.js'
4+
import { HookRegistryImplementation } from '../../hooks/registry.js'
5+
import { AfterInvocationEvent, AfterModelCallEvent } from '../../hooks/events.js'
6+
import { ContextWindowOverflowError } from '../../errors.js'
7+
import { createMockAgent } from '../../__fixtures__/agent-helpers.js'
58

69
describe('NullConversationManager', () => {
7-
describe('applyManagement', () => {
8-
it('does not modify messages array', () => {
10+
describe('behavior', () => {
11+
it('does not modify conversation history', async () => {
912
const manager = new NullConversationManager()
1013
const messages = [
1114
new Message({ role: 'user', content: [new TextBlock('Hello')] }),
1215
new Message({ role: 'assistant', content: [new TextBlock('Hi there')] }),
1316
]
14-
const mockAgent = { messages } as unknown as Agent
17+
const mockAgent = createMockAgent({ messages })
1518

16-
manager.applyManagement(mockAgent)
19+
const registry = new HookRegistryImplementation()
20+
manager.registerCallbacks(registry)
21+
22+
await registry.invokeCallbacks(new AfterInvocationEvent({ agent: mockAgent }))
1723

1824
expect(mockAgent.messages).toHaveLength(2)
1925
expect(mockAgent.messages[0]!.content[0]).toEqual({ type: 'textBlock', text: 'Hello' })
2026
expect(mockAgent.messages[1]!.content[0]).toEqual({ type: 'textBlock', text: 'Hi there' })
2127
})
22-
})
23-
24-
describe('reduceContext', () => {
25-
it('re-throws provided error', () => {
26-
const manager = new NullConversationManager()
27-
const mockAgent = { messages: [] } as unknown as Agent
28-
const testError = new Error('Test error')
29-
30-
expect(() => {
31-
manager.reduceContext(mockAgent, testError)
32-
}).toThrow(testError)
33-
})
3428

35-
it('throws ContextWindowOverflowError when no error provided', () => {
29+
it('does not set retryModelCall on context overflow', async () => {
3630
const manager = new NullConversationManager()
37-
const mockAgent = { messages: [] } as unknown as Agent
31+
const mockAgent = createMockAgent()
32+
const error = new ContextWindowOverflowError('Context overflow')
3833

39-
expect(() => {
40-
manager.reduceContext(mockAgent)
41-
}).toThrow(ContextWindowOverflowError)
42-
})
34+
const registry = new HookRegistryImplementation()
35+
manager.registerCallbacks(registry)
4336

44-
it('throws ContextWindowOverflowError with correct message when no error provided', () => {
45-
const manager = new NullConversationManager()
46-
const mockAgent = { messages: [] } as unknown as Agent
37+
const event = await registry.invokeCallbacks(new AfterModelCallEvent({ agent: mockAgent, error }))
4738

48-
expect(() => {
49-
manager.reduceContext(mockAgent)
50-
}).toThrow('Context window overflowed!')
39+
expect(event.retryModelCall).toBeUndefined()
5140
})
5241
})
5342
})

0 commit comments

Comments
 (0)