diff --git a/packages/shared/tool-sdk/package.json b/packages/shared/tool-sdk/package.json index ad4836a3f..ec9e03c84 100644 --- a/packages/shared/tool-sdk/package.json +++ b/packages/shared/tool-sdk/package.json @@ -19,6 +19,7 @@ }, "dependencies": { "@agentx/agent-platform": "workspace:*", + "@agentx/cache": "workspace:*", "@agentx/core-runtime": "workspace:*", "@agentx/observability": "workspace:*", "yaml": "^2.9.0" diff --git a/packages/shared/tool-sdk/src/pipeline/index.ts b/packages/shared/tool-sdk/src/pipeline/index.ts index a3f26ff3e..fa7757a2d 100644 --- a/packages/shared/tool-sdk/src/pipeline/index.ts +++ b/packages/shared/tool-sdk/src/pipeline/index.ts @@ -6,11 +6,41 @@ import type { ToolExecutionResponse, } from '../interfaces/index.js'; import { Tracer, Metrics } from '@agentx/observability'; +import { CacheManager } from '@agentx/cache'; + +interface CacheKey { + toolName: string; + category: string; + argsHash: string; + workingDirectory: string; +} export class ToolExecutionPipelineImpl implements ToolExecutionPipeline { private hooks: ExecutionHooks[] = []; private tracer = new Tracer('tool-sdk-pipeline'); private metrics = new Metrics(); + private cache: CacheManager; + + constructor(cacheTtlMs: number = 300_000) { + this.cache = new CacheManager(); + this.cacheTtlMs = cacheTtlMs; + } + + private cacheTtlMs: number; + + private generateCacheKey(req: ToolExecutionRequest): string { + const key: CacheKey = { + toolName: req.toolName || 'unknown', + category: req.category || 'unknown', + argsHash: JSON.stringify(req.arguments || {}), + workingDirectory: req.context?.workingDirectory || '', + }; + return `tool:${key.category}:${key.toolName}:${Buffer.from(key.argsHash).toString('base64')}:${key.workingDirectory}`; + } + + private isReadOperation(category: string): boolean { + return category?.endsWith('.read') ?? false; + } public addHook(hook: ExecutionHooks): void { this.hooks.push(hook); @@ -21,6 +51,20 @@ export class ToolExecutionPipelineImpl implements ToolExecutionPipeline { span.setAttribute('tool.name', tool.definition.name); span.setAttribute('tool.category', tool.definition.category); + // Check cache for read operations + const cacheKey = this.generateCacheKey(req); + if (this.isReadOperation(req.category)) { + const cached = await this.cache.get(cacheKey); + if (cached) { + this.metrics.counter('tool_cache_hit', 1, { tool: tool.definition.name }); + span.setAttribute('tool.cache_hit', true); + span.end(); + return cached; + } + this.metrics.counter('tool_cache_miss', 1, { tool: tool.definition.name }); + span.setAttribute('tool.cache_hit', false); + } + // PreExecute hooks for (const hook of this.hooks) { if (hook.preExecute) { @@ -31,6 +75,12 @@ export class ToolExecutionPipelineImpl implements ToolExecutionPipeline { let response: ToolExecutionResponse; try { response = await tool.execute(req); + + // Cache successful read operations + if (this.isReadOperation(req.category) && response.result.success) { + await this.cache.set(cacheKey, response, this.cacheTtlMs); + } + this.metrics.counter('tool_executions_success', 1, { tool: tool.definition.name }); span.setStatus({ code: 0 }); } catch (error: unknown) { diff --git a/packages/shared/tool-sdk/test/tool.test.ts b/packages/shared/tool-sdk/test/tool.test.ts index a7aea5177..b23034ecc 100644 --- a/packages/shared/tool-sdk/test/tool.test.ts +++ b/packages/shared/tool-sdk/test/tool.test.ts @@ -343,4 +343,78 @@ describe('ToolExecutionPipeline', () => { ); expect(errHook).toHaveBeenCalled(); }); + + it('caches read operations and returns cached results', async () => { + const pipeline = new ToolExecutionPipelineImpl(300_000); + const req = { + toolName: 'ReadTool', + category: 'fs.read', + arguments: { path: '/test' }, + context: { workingDirectory: '/workspace', taskId: 't1', traceId: 'tr1', agentRole: 'coder' }, + } as unknown as ToolExecutionRequest; + const res = { + result: { success: true, output: 'cached content' }, + } as unknown as ToolExecutionResponse; + const executeFn = vi.fn().mockResolvedValue(res); + const tool = { + execute: executeFn, + definition: { name: 'ReadTool', category: 'fs.read' }, + } as unknown as ITool; + + // First call - cache miss + const result1 = await pipeline.execute(req, tool); + expect(result1).toBe(res); + expect(executeFn).toHaveBeenCalledTimes(1); + + // Second call - cache hit + const result2 = await pipeline.execute(req, tool); + expect(result2).toBe(res); + expect(executeFn).toHaveBeenCalledTimes(1); // Not called again + }); + + it('does not cache write operations', async () => { + const pipeline = new ToolExecutionPipelineImpl(300_000); + const res = { + result: { success: true, output: 'written' }, + } as unknown as ToolExecutionResponse; + const executeFn = vi.fn().mockResolvedValue(res); + const req = { + toolName: 'WriteTool', + category: 'fs.write', + arguments: { path: '/test', content: 'data' }, + context: { workingDirectory: '/workspace', taskId: 't1', traceId: 'tr1', agentRole: 'coder' }, + } as unknown as ToolExecutionRequest; + const tool = { + execute: executeFn, + definition: { name: 'WriteTool', category: 'fs.write' }, + } as unknown as ITool; + + const result1 = await pipeline.execute(req, tool); + const result2 = await pipeline.execute(req, tool); + expect(executeFn).toHaveBeenCalledTimes(2); // Called twice + expect(result1).toBe(res); + expect(result2).toBe(res); + }); + + it('does not cache failed operations', async () => { + const pipeline = new ToolExecutionPipelineImpl(300_000); + const failRes = { + result: { success: false, error: 'File not found' }, + } as unknown as ToolExecutionResponse; + const executeFn = vi.fn().mockResolvedValue(failRes); + const req = { + toolName: 'ReadTool', + category: 'fs.read', + arguments: { path: '/missing' }, + context: { workingDirectory: '/workspace', taskId: 't1', traceId: 'tr1', agentRole: 'coder' }, + } as unknown as ToolExecutionRequest; + const tool = { + execute: executeFn, + definition: { name: 'ReadTool', category: 'fs.read' }, + } as unknown as ITool; + + await pipeline.execute(req, tool); + await pipeline.execute(req, tool); + expect(executeFn).toHaveBeenCalledTimes(2); // Called twice (not cached) + }); }); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 54eb18f4e..554ea7f81 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -651,6 +651,9 @@ importers: '@agentx/agent-platform': specifier: workspace:* version: link:../../agent/agent-platform + '@agentx/cache': + specifier: workspace:* + version: link:../cache '@agentx/core-runtime': specifier: workspace:* version: link:../core-runtime