diff --git a/src/agent/executor.ts b/src/agent/executor.ts index 89f50e4a8..7be072d23 100644 --- a/src/agent/executor.ts +++ b/src/agent/executor.ts @@ -127,12 +127,15 @@ export interface ExecutionResult { steps: AgentStep[]; } +type AgentStepCallback = (step: AgentStep) => void | Promise; + export async function executeAgentPlan( plan: AgentPlan, engines: SearchEngine[], router: SmartRouter, budget: ExecutionBudget, prompt = '', + onStep?: AgentStepCallback, ): Promise { const steps: AgentStep[] = []; const allUrls = new Set(); @@ -149,11 +152,13 @@ export async function executeAgentPlan( const searchStart = Date.now(); const searchResults = await executeSearches(plan.searches, engines, budget.deadlineMs, prompt); - steps.push({ + const searchStep: AgentStep = { action: 'search', detail: `Searched ${plan.searches.length} queries, found ${searchResults.length} results`, time_ms: Date.now() - searchStart, - }); + }; + steps.push(searchStep); + await onStep?.(searchStep); for (const result of searchResults) { allUrls.add(result.url); @@ -169,11 +174,13 @@ export async function executeAgentPlan( const fetchStart = Date.now(); const sources = await fetchPages(urlsToFetch, router, budget); - steps.push({ + const fetchStep: AgentStep = { action: 'fetch', detail: `Fetched ${sources.filter((s) => s.fetched).length}/${urlsToFetch.length} pages`, time_ms: Date.now() - fetchStart, - }); + }; + steps.push(fetchStep); + await onStep?.(fetchStep); // Phase 4: Post-fetch relevance scoring // Only filter when a real reranker is configured; the token-overlap diff --git a/src/agent/pipeline.ts b/src/agent/pipeline.ts index 76ef1af97..cab07a13e 100644 --- a/src/agent/pipeline.ts +++ b/src/agent/pipeline.ts @@ -15,6 +15,7 @@ import type { AgentSource, AgentStep, GridConfidence, + ProgressCallback, SearchEngine, } from '../types.js'; import type { SmartRouter } from '../fetch/router.js'; @@ -39,12 +40,23 @@ export async function runAgentPipeline( engines: SearchEngine[], router: SmartRouter, server?: SamplingCapableServer, + onProgress?: ProgressCallback, ): Promise { const start = Date.now(); const maxPages = input.max_pages ?? DEFAULT_MAX_PAGES; const maxTimeMs = input.max_time_ms ?? DEFAULT_MAX_TIME_MS; const deadlineMs = start + maxTimeMs; const steps: AgentStep[] = []; + let completedSteps = 0; + const reportStep = async (step: AgentStep): Promise => { + if (!onProgress) return; + completedSteps += 1; + try { + await onProgress({ progress: completedSteps, message: step.detail }); + } catch (err) { + log.debug('agent progress notification failed', { error: String(err) }); + } + }; try { const planStart = Date.now(); @@ -52,11 +64,13 @@ export async function runAgentPipeline( const plan = await planExecution(input.prompt, input.urls, server); - steps.push({ + const planStep: AgentStep = { action: 'plan', detail: `Generated ${plan.searches.length} searches, ${plan.urls.length} URLs${plan.samplingUsed ? ' (via sampling)' : ' (keyword extraction)'}`, time_ms: Date.now() - planStart, - }); + }; + steps.push(planStep); + await reportStep(planStep); log.info('plan generated', { searches: plan.searches.length, @@ -67,7 +81,7 @@ export async function runAgentPipeline( const execResult = await executeAgentPlan(plan, engines, router, { maxPages, deadlineMs, - }, input.prompt); + }, input.prompt, reportStep); steps.push(...execResult.steps); @@ -87,11 +101,13 @@ export async function runAgentPipeline( const extractStart = Date.now(); const schemaResult = applySchemaExtraction(sources, input.schema as JsonSchema); - steps.push({ + const extractStep: AgentStep = { action: 'extract', detail: `Applied schema extraction to ${fetchedCount} sources`, time_ms: Date.now() - extractStart, - }); + }; + steps.push(extractStep); + await reportStep(extractStep); if (schemaResult && !schemaResult.lowConfidence) { return { @@ -127,11 +143,13 @@ export async function runAgentPipeline( : llmUsed ? ' (via configured LLM)' : ' (evidence fallback)'; - steps.push({ + const synthesizeStep: AgentStep = { action: 'synthesize', detail: `Produced ${resultLen} char result${synthPath}`, time_ms: Date.now() - synthStart, - }); + }; + steps.push(synthesizeStep); + await reportStep(synthesizeStep); // When every fetch failed but the planner produced URLs, // surface that as a partial-fail warning so callers don't see "No data diff --git a/src/server.ts b/src/server.ts index 92a61bc9c..e59135a02 100644 --- a/src/server.ts +++ b/src/server.ts @@ -508,7 +508,7 @@ export function createMcpServer(subsystems: Subsystems): Server { if (name === 'agent') { const input = (args ?? {}) as unknown as AgentInput; const samplingServer = server as unknown as SamplingCapableServer; - const r = await handleAgent(input, searchEngines, router, backendStatus, samplingServer); + const r = await handleAgent(input, searchEngines, router, backendStatus, samplingServer, onProgress); if (!r.ok) { return { content: [{ type: 'text', text: JSON.stringify({ error: r.error, error_reason: r.error_reason, stage: r.stage, ...(r.hint ? { hint: r.hint } : {}) }, null, 2) }], diff --git a/src/tools/agent.ts b/src/tools/agent.ts index d6802ce6b..be9efbfd6 100644 --- a/src/tools/agent.ts +++ b/src/tools/agent.ts @@ -11,6 +11,7 @@ import type { AgentInput, AgentOutput, EvidenceItem, + ProgressCallback, SearchEngine, StageResult, } from '../types.js'; @@ -31,6 +32,7 @@ export async function handleAgent( router: SmartRouter, _backendStatus?: unknown, server?: SamplingCapableServer, + onProgress?: ProgressCallback, ): Promise> { try { if (!input.prompt || typeof input.prompt !== 'string' || input.prompt.trim().length === 0) { @@ -88,6 +90,7 @@ export async function handleAgent( engines, router, server, + input.stream ? onProgress : undefined, ); result.response_time_ms = Date.now() - _start; diff --git a/tests/integration/agent.test.ts b/tests/integration/agent.test.ts index 102df8981..c37b137b9 100644 --- a/tests/integration/agent.test.ts +++ b/tests/integration/agent.test.ts @@ -91,6 +91,61 @@ describe('agent tool integration', () => { expect(synthStep).toBeDefined(); }); + it('emits progress for each completed step when streaming is enabled', async () => { + const onProgress = vi.fn(); + + const response = await handleAgent( + { prompt: 'Find CRM pricing', stream: true }, + [stubEngine], + stubRouter, + undefined, + undefined, + onProgress, + ); + const result = response.ok ? response.data : ({ ...response } as any); + + expect(result.error).toBeUndefined(); + expect(onProgress).toHaveBeenCalledTimes(result.steps.length); + expect(onProgress.mock.calls.map(([update]) => update.progress)).toEqual( + result.steps.map((_, index) => index + 1), + ); + expect(onProgress.mock.calls.map(([update]) => update.message)).toEqual( + result.steps.map((step) => step.detail), + ); + }); + + it('does not emit progress when streaming is disabled', async () => { + const onProgress = vi.fn(); + + const response = await handleAgent( + { prompt: 'Find CRM pricing' }, + [stubEngine], + stubRouter, + undefined, + undefined, + onProgress, + ); + + expect(response.ok).toBe(true); + expect(onProgress).not.toHaveBeenCalled(); + }); + + it('keeps the agent result when a progress callback fails', async () => { + const onProgress = vi.fn().mockRejectedValue(new Error('transport closed')); + + const response = await handleAgent( + { prompt: 'Find CRM pricing', stream: true }, + [stubEngine], + stubRouter, + undefined, + undefined, + onProgress, + ); + + expect(response.ok).toBe(true); + expect(onProgress).toHaveBeenCalled(); + }); + it('full pipeline with explicit URLs', async () => { const input: AgentInput = { prompt: 'Compare pricing', diff --git a/tests/unit/server/schema-registration.test.ts b/tests/unit/server/schema-registration.test.ts index 4275246e0..b43bb6b46 100644 --- a/tests/unit/server/schema-registration.test.ts +++ b/tests/unit/server/schema-registration.test.ts @@ -102,6 +102,7 @@ vi.mock('../../../src/searxng/docker.js', () => ({ vi.mock('../../../src/cache/store.js', () => ({ getCachedContent: vi.fn().mockReturnValue(null), isExpired: vi.fn().mockReturnValue(false), + cacheContent: vi.fn(), })); // Avoid cold ONNX startup on every `connectClient()` — schema registration @@ -118,6 +119,10 @@ vi.mock('../../../src/embedding/embed.js', () => ({ async function connectClient() { const { initSubsystems, createMcpServer } = await import('../../../src/server.js'); const subs = await initSubsystems(); + subs.searchEngines.splice(0, subs.searchEngines.length, { + name: 'stub', + search: vi.fn().mockResolvedValue([]), + }); const server = createMcpServer(subs); const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); @@ -179,6 +184,42 @@ describe('diff + watch tool registration', () => { } }); + it('tools/call agent delivers progress notifications before the final result', async () => { + const { client, teardown } = await connectClient(); + try { + const progress: Array<{ progress: number; total?: number; message?: string }> = []; + const res = await client.callTool( + { + name: 'agent', + arguments: { + prompt: 'Read the example page title', + urls: ['https://example.com/'], + max_pages: 1, + stream: true, + }, + }, + undefined, + { + timeout: 10_000, + onprogress: (update) => progress.push(update), + }, + ); + const block = (res.content as Array<{ type: string; text: string }>)[0]; + const payload = JSON.parse(block.text); + + expect(res.isError).not.toBe(true); + expect(progress).toHaveLength(payload.steps.length); + expect(progress.map((update) => update.progress)).toEqual( + payload.steps.map((_: unknown, index: number) => index + 1), + ); + expect(progress.map((update) => update.message)).toEqual( + payload.steps.map((step: { detail: string }) => step.detail), + ); + } finally { + await teardown(); + } + }); + // The MCP surface accepts the real input shape ({ old, new, output }) and // returns a structured DiffOutput. it('tools/call diff computes a real diff between two markdown bodies', async () => {