|
| 1 | +import { afterEach, describe, expect, it, vi } from "vitest"; |
| 2 | + |
| 3 | +import { AuthStorage } from "../src/core/auth-storage.ts"; |
| 4 | +import { DEFAULT_COMPACTION_SETTINGS } from "../src/core/compaction/index.ts"; |
| 5 | +import { renderSearchResult } from "../src/core/extensions/builtin/websearch/websearch/renderers.ts"; |
| 6 | +import { createWebSearchTool } from "../src/core/extensions/builtin/websearch/websearch/tool.ts"; |
| 7 | +import type { |
| 8 | + SearchProgressDetails, |
| 9 | + WebsearchConfig, |
| 10 | +} from "../src/core/extensions/builtin/websearch/websearch/types.ts"; |
| 11 | +import type { ExtensionContext } from "../src/core/extensions/types.ts"; |
| 12 | +import { ModelRegistry } from "../src/core/model-registry.ts"; |
| 13 | + |
| 14 | +function minimalToolContext(): ExtensionContext { |
| 15 | + return { |
| 16 | + ui: Object.create(null) as ExtensionContext["ui"], |
| 17 | + mode: "print", |
| 18 | + hasUI: false, |
| 19 | + cwd: process.cwd(), |
| 20 | + sessionManager: Object.create(null) as ExtensionContext["sessionManager"], |
| 21 | + modelRegistry: ModelRegistry.inMemory(AuthStorage.inMemory()), |
| 22 | + model: undefined, |
| 23 | + serviceTier: undefined, |
| 24 | + isIdle: () => true, |
| 25 | + isProjectTrusted: () => true, |
| 26 | + signal: undefined, |
| 27 | + abort: vi.fn(), |
| 28 | + hasPendingMessages: () => false, |
| 29 | + shutdown: vi.fn(), |
| 30 | + getContextUsage: () => undefined, |
| 31 | + getCompactionSettings: () => DEFAULT_COMPACTION_SETTINGS, |
| 32 | + compact: vi.fn(), |
| 33 | + getMessageRevision: () => 0, |
| 34 | + applyCompaction: async () => ({ applied: false, reason: "rejected" }), |
| 35 | + getSystemPrompt: () => "", |
| 36 | + }; |
| 37 | +} |
| 38 | + |
| 39 | +const passthroughTheme = { |
| 40 | + bold: (value: string) => value, |
| 41 | + fg: (_key: string, value: string) => value, |
| 42 | +}; |
| 43 | + |
| 44 | +function fallbackConfig(): WebsearchConfig { |
| 45 | + return { |
| 46 | + strategy: "priority", |
| 47 | + fallback: true, |
| 48 | + auto: false, |
| 49 | + providers: [ |
| 50 | + { id: "primary", provider: "exa", apiKey: "test-key" }, |
| 51 | + { id: "backup", provider: "exa", apiKey: "test-key" }, |
| 52 | + ], |
| 53 | + }; |
| 54 | +} |
| 55 | + |
| 56 | +function exaSuccessResponse(): Response { |
| 57 | + return new Response( |
| 58 | + JSON.stringify({ results: [{ title: "Result", url: "https://example.com/a", text: "snippet" }] }), |
| 59 | + { status: 200, headers: { "Content-Type": "application/json" } }, |
| 60 | + ); |
| 61 | +} |
| 62 | + |
| 63 | +afterEach(() => { |
| 64 | + vi.unstubAllGlobals(); |
| 65 | +}); |
| 66 | + |
| 67 | +describe("websearch per-attempt progress", () => { |
| 68 | + it("#given a failing first provider #when the tool executes with fallback #then emits one progress update per attempt with the current provider", async () => { |
| 69 | + // given |
| 70 | + const responses = [new Response("boom", { status: 500 }), exaSuccessResponse()]; |
| 71 | + const fetchMock = vi.fn<typeof fetch>(async () => { |
| 72 | + const next = responses.shift(); |
| 73 | + if (!next) throw new Error("unexpected fetch call"); |
| 74 | + return next; |
| 75 | + }); |
| 76 | + vi.stubGlobal("fetch", fetchMock); |
| 77 | + const tool = createWebSearchTool(() => ({ ok: true, config: fallbackConfig(), source: "test" })); |
| 78 | + const progress: SearchProgressDetails[] = []; |
| 79 | + |
| 80 | + // when |
| 81 | + const result = await tool.execute( |
| 82 | + "attempt-progress", |
| 83 | + { query: "attempt progress" }, |
| 84 | + undefined, |
| 85 | + (update) => { |
| 86 | + if (update.details && "phase" in update.details && update.details.phase === "searching") { |
| 87 | + progress.push(update.details); |
| 88 | + } |
| 89 | + }, |
| 90 | + minimalToolContext(), |
| 91 | + ); |
| 92 | + |
| 93 | + // then |
| 94 | + expect(progress).toHaveLength(3); |
| 95 | + expect(progress[0]?.currentProvider).toBeUndefined(); |
| 96 | + expect(progress[1]?.currentProvider).toBe("exa/primary"); |
| 97 | + expect(progress[1]?.attempts).toEqual([]); |
| 98 | + expect(progress[2]?.currentProvider).toBe("exa/backup"); |
| 99 | + expect(progress[2]?.attempts).toHaveLength(1); |
| 100 | + expect(progress[2]?.attempts?.[0]?.error).toContain("HTTP 500"); |
| 101 | + expect(result.details && "provider" in result.details ? result.details.entryId : undefined).toBe("backup"); |
| 102 | + }); |
| 103 | + |
| 104 | + it("#given per-attempt progress details #when rendering partial output #then shows only the current provider with its step position", () => { |
| 105 | + // given |
| 106 | + const details: SearchProgressDetails = { |
| 107 | + phase: "searching", |
| 108 | + query: "attempt progress", |
| 109 | + providerLabels: ["exa/primary", "exa/backup"], |
| 110 | + maxResults: 10, |
| 111 | + currentProvider: "exa/backup", |
| 112 | + attempts: [{ provider: "exa", entryId: "primary", durationMs: 12, resultsCount: 0, error: "HTTP 500" }], |
| 113 | + }; |
| 114 | + |
| 115 | + // when |
| 116 | + const collapsed = renderSearchResult( |
| 117 | + { content: [{ type: "text", text: "" }], details }, |
| 118 | + { expanded: false, isPartial: true }, |
| 119 | + passthroughTheme, |
| 120 | + ) |
| 121 | + .render(200) |
| 122 | + .join("\n"); |
| 123 | + const expanded = renderSearchResult( |
| 124 | + { content: [{ type: "text", text: "" }], details }, |
| 125 | + { expanded: true, isPartial: true }, |
| 126 | + passthroughTheme, |
| 127 | + ) |
| 128 | + .render(200) |
| 129 | + .join("\n"); |
| 130 | + |
| 131 | + // then |
| 132 | + expect(collapsed).toContain('Searching "attempt progress" via exa/backup [2/2] (max 10)'); |
| 133 | + expect(collapsed).not.toContain("exa/primary ->"); |
| 134 | + expect(expanded).toContain("route exa/primary:failed"); |
| 135 | + }); |
| 136 | +}); |
0 commit comments