From e3c7490bfa6d6621a646f214a4ced48b7b35de96 Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Wed, 3 Jun 2026 16:32:13 +0900 Subject: [PATCH 1/9] test: cover WebSearch registration, delegate, and grok tool scope --- tests/provider/package.test.ts | 2 + tests/provider/register.test.ts | 97 +++++++++++++++------------ tests/provider/toolScope.test.ts | 49 ++++++++++++++ tests/tools/register.test.ts | 35 ++++++++-- tests/tools/toolTestHelpers.ts | 6 ++ tests/tools/webSearch.test.ts | 76 +++++++++++++++++++++ tests/tools/webSearchDelegate.test.ts | 35 ++++++++++ 7 files changed, 252 insertions(+), 48 deletions(-) create mode 100644 tests/provider/toolScope.test.ts create mode 100644 tests/tools/webSearch.test.ts create mode 100644 tests/tools/webSearchDelegate.test.ts diff --git a/tests/provider/package.test.ts b/tests/provider/package.test.ts index 3ff0c1a..1cbddec 100644 --- a/tests/provider/package.test.ts +++ b/tests/provider/package.test.ts @@ -49,6 +49,8 @@ describe('repository layout', () => { 'src/tools/rendering.ts', 'src/tools/search.ts', 'src/tools/shell.ts', + 'src/tools/webSearch.ts', + 'src/tools/webSearchDelegate.ts', ]); }); diff --git a/tests/provider/register.test.ts b/tests/provider/register.test.ts index ad5a3f0..54545c1 100644 --- a/tests/provider/register.test.ts +++ b/tests/provider/register.test.ts @@ -3,34 +3,49 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import type { ExtensionAPI, ProviderConfig } from '@earendil-works/pi-coding-agent'; import { afterEach, describe, expect, it, vi } from 'vitest'; - -const streamSimpleOpenAIResponses = vi.fn( - ( - _model: unknown, - _context: unknown, - options?: { - onResponse?: (response: { headers: Record }) => void; - }, - ) => { - options?.onResponse?.({ - headers: { - 'x-ratelimit-remaining-requests': '179', - 'x-ratelimit-limit-requests': '180', - 'x-ratelimit-remaining-tokens': '7500000', - 'x-ratelimit-limit-tokens': '7500000', - 'x-grok-context-window': '512000', - 'x-zero-data-retention': 'true', +import { GROK_SHIM_TOOL_NAMES, grokToolsToActivate } from '../../src/tools/register.js'; +import * as webSearchDelegate from '../../src/tools/webSearchDelegate.js'; + +const { streamSimpleOpenAIResponses, mockPiWebAccessInstalled } = vi.hoisted(() => ({ + mockPiWebAccessInstalled: vi.fn(() => true), + streamSimpleOpenAIResponses: vi.fn( + ( + _model: unknown, + _context: unknown, + options?: { + onResponse?: (response: { headers: Record }) => void; }, - }); - return {}; - }, -); + ) => { + options?.onResponse?.({ + headers: { + 'x-ratelimit-remaining-requests': '179', + 'x-ratelimit-limit-requests': '180', + 'x-ratelimit-remaining-tokens': '7500000', + 'x-ratelimit-limit-tokens': '7500000', + 'x-grok-context-window': '512000', + 'x-zero-data-retention': 'true', + }, + }); + return {}; + }, + ), +})); vi.mock('@earendil-works/pi-ai', async (importOriginal) => ({ ...(await importOriginal()), streamSimpleOpenAIResponses, })); +vi.mock('../../src/tools/webSearchDelegate.js', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + isPiWebAccessInstalled: () => mockPiWebAccessInstalled(), + bindLivePiWebAccess: vi.fn(), + ensureWebSearchDelegate: vi.fn(async () => undefined), + }; +}); + interface CommandConfig { handler: (args: string[], ctx: TestContext) => Promise; } @@ -62,18 +77,6 @@ interface TestContext { type ExtensionHandler = (event: unknown, ctx: TestContext) => unknown; -const grokToolNames = [ - 'Grep', - 'Glob', - 'LS', - 'Read', - 'Write', - 'StrReplace', - 'Edit', - 'Delete', - 'Shell', -]; - const originalFetch = globalThis.fetch; const originalHome = process.env.HOME; const originalToken = process.env.GROK_CLI_OAUTH_TOKEN; @@ -96,7 +99,8 @@ afterEach(() => { for (const dir of tempDirs.splice(0)) rmSync(dir, { recursive: true }); }); -async function setupExtension(initialActiveTools = ['read', 'bash']) { +async function setupExtension(initialActiveTools = ['read', 'bash'], piWebAccessInstalled = true) { + vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(piWebAccessInstalled); const commands = new Map(); const providers = new Map(); const tools = new Map(); @@ -477,22 +481,27 @@ describe('Grok CLI tool scoping', () => { it('registers the Grok/Cursor-native tool shims', async () => { const extension = await setupExtension(); - expect([...extension.tools.keys()].sort()).toEqual([...grokToolNames].sort()); + expect([...extension.tools.keys()].sort()).toEqual([...grokToolsToActivate()].sort()); + }); + + it('does not register WebSearch when pi-web-access is not installed', async () => { + const extension = await setupExtension(['read', 'bash'], false); + + expect([...extension.tools.keys()].sort()).toEqual([...GROK_SHIM_TOOL_NAMES].sort()); + expect(extension.tools.has('WebSearch')).toBe(false); }); it('enables Grok tools for Grok models while preserving other active tools', async () => { - const extension = await setupExtension(['read', 'custom_tool']); + const extension = await setupExtension(['read', 'custom_tool', 'web_search']); await extension.handlers.get('model_select')?.( { model: { provider: 'grok-cli', id: 'grok-build' } }, contextForModel('grok-cli'), ); - expect(extension.setActiveTools).toHaveBeenLastCalledWith([ - 'read', - 'custom_tool', - ...grokToolNames, - ]); + const next = extension.setActiveTools.mock.calls.at(-1)?.[0] as string[]; + expect(next).not.toContain('web_search'); + expect(next).toEqual(['read', 'custom_tool', ...grokToolsToActivate()]); }); it('removes Grok tools for non-Grok models while preserving other active tools', async () => { @@ -511,11 +520,11 @@ describe('Grok CLI tool scoping', () => { await extension.handlers.get('before_agent_start')?.({}, contextForModel('grok-cli')); - expect(extension.setActiveTools).toHaveBeenLastCalledWith(['read', ...grokToolNames]); + expect(extension.setActiveTools).toHaveBeenLastCalledWith(['read', ...grokToolsToActivate()]); }); it('does not update active tools when the selection is already correct', async () => { - const extension = await setupExtension(['read', ...grokToolNames]); + const extension = await setupExtension(['read', ...grokToolsToActivate()]); await extension.handlers.get('before_agent_start')?.({}, contextForModel('grok-cli')); @@ -527,7 +536,7 @@ describe('Grok CLI tool rendering', () => { it('adds renderers to every Grok tool shim', async () => { const extension = await setupExtension(); - for (const name of grokToolNames) { + for (const name of grokToolsToActivate()) { expect(extension.tools.get(name)?.renderCall).toBeTypeOf('function'); expect(extension.tools.get(name)?.renderResult).toBeTypeOf('function'); } diff --git a/tests/provider/toolScope.test.ts b/tests/provider/toolScope.test.ts new file mode 100644 index 0000000..8254a06 --- /dev/null +++ b/tests/provider/toolScope.test.ts @@ -0,0 +1,49 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { syncGrokTools } from '../../src/provider/toolScope.js'; +import * as webSearchDelegate from '../../src/tools/webSearchDelegate.js'; + +afterEach(() => { + vi.restoreAllMocks(); +}); + +function syncForGrokCli(piWebAccessInstalled: boolean) { + vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(piWebAccessInstalled); + const setActiveTools = vi.fn(); + syncGrokTools( + { + getActiveTools: () => ['read', 'web_search', 'bash'], + setActiveTools, + }, + 'grok-cli', + ); + return setActiveTools.mock.calls[0][0] as string[]; +} + +describe('syncGrokTools', () => { + it('drops web_search and enables WebSearch for grok-cli when pi-web-access is installed', () => { + const next = syncForGrokCli(true); + expect(next).not.toContain('web_search'); + expect(next).toContain('WebSearch'); + expect(next).toContain('read'); + }); + + it('does not add WebSearch for grok-cli when pi-web-access is not installed', () => { + const next = syncForGrokCli(false); + expect(next).not.toContain('web_search'); + expect(next).not.toContain('WebSearch'); + expect(next).toContain('Grep'); + }); + + it('removes Grok shims and leaves web_search available for other providers', () => { + const setActiveTools = vi.fn(); + syncGrokTools( + { + getActiveTools: () => ['read', 'web_search', 'Grep', 'WebSearch'], + setActiveTools, + }, + 'openai', + ); + + expect(setActiveTools).toHaveBeenCalledWith(['read', 'web_search']); + }); +}); diff --git a/tests/tools/register.test.ts b/tests/tools/register.test.ts index 04b8690..17de043 100644 --- a/tests/tools/register.test.ts +++ b/tests/tools/register.test.ts @@ -1,9 +1,19 @@ import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { describe, expect, it } from 'vitest'; -import { GROK_TOOL_NAMES, registerGrokTools } from '../../src/tools/register.js'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { + GROK_SHIM_TOOL_NAMES, + grokToolsToActivate, + registerGrokTools, +} from '../../src/tools/register.js'; +import * as webSearchDelegate from '../../src/tools/webSearchDelegate.js'; + +afterEach(() => { + vi.restoreAllMocks(); +}); describe('Grok tool registration', () => { - it('registers all Grok/Cursor-native tool shims with renderers', () => { + it('registers shim tools with renderers', () => { + vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(false); const toolNames: string[] = []; registerGrokTools({ @@ -12,8 +22,25 @@ describe('Grok tool registration', () => { expect(tool.renderCall).toBeTypeOf('function'); expect(tool.renderResult).toBeTypeOf('function'); }, + on() {}, + } as unknown as ExtensionAPI); + + expect(toolNames.sort()).toEqual([...GROK_SHIM_TOOL_NAMES].sort()); + expect(toolNames).not.toContain('WebSearch'); + }); + + it('registers WebSearch when pi-web-access is installed', () => { + vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(true); + const toolNames: string[] = []; + + registerGrokTools({ + registerTool(tool: { name: string }) { + toolNames.push(tool.name); + }, + on() {}, } as unknown as ExtensionAPI); - expect(toolNames.sort()).toEqual([...GROK_TOOL_NAMES].sort()); + expect(toolNames).toContain('WebSearch'); + expect(grokToolsToActivate()).toContain('WebSearch'); }); }); diff --git a/tests/tools/toolTestHelpers.ts b/tests/tools/toolTestHelpers.ts index eca9206..cf2d67b 100644 --- a/tests/tools/toolTestHelpers.ts +++ b/tests/tools/toolTestHelpers.ts @@ -15,6 +15,8 @@ export type ToolResult = { details: Record; }; +type ExtensionHandler = (event: unknown) => unknown; + type Renderable = { render: (width: number) => string[] }; type ToolTheme = { @@ -43,10 +45,14 @@ type RegisteredTool = { export function collectTools(registerTools: (pi: ExtensionAPI) => void) { const tools = new Map(); + const handlers = new Map(); registerTools({ registerTool(tool: RegisteredTool) { tools.set(tool.name, tool); }, + on(event: string, handler: ExtensionHandler) { + handlers.set(event, handler); + }, } as unknown as ExtensionAPI); return tools; } diff --git a/tests/tools/webSearch.test.ts b/tests/tools/webSearch.test.ts new file mode 100644 index 0000000..81f5899 --- /dev/null +++ b/tests/tools/webSearch.test.ts @@ -0,0 +1,76 @@ +import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { registerGrokTools } from '../../src/tools/register.js'; +import { registerWebSearchTool } from '../../src/tools/webSearch.js'; +import * as webSearchDelegate from '../../src/tools/webSearchDelegate.js'; +import { + clearWebSearchDelegateForTests, + setWebSearchDelegateForTests, +} from '../../src/tools/webSearchDelegate.js'; +import { collectTools, executeTool, firstText, renderToolCall } from './toolTestHelpers.js'; + +afterEach(() => { + clearWebSearchDelegateForTests(); +}); + +describe('WebSearch tool', () => { + it('registers WebSearch with renderers', () => { + const names: string[] = []; + registerWebSearchTool({ + registerTool(tool: { name: string; renderCall?: unknown; renderResult?: unknown }) { + names.push(tool.name); + expect(tool.renderCall).toBeTypeOf('function'); + expect(tool.renderResult).toBeTypeOf('function'); + }, + on() {}, + } as unknown as ExtensionAPI); + + expect(names).toContain('WebSearch'); + }); + + it('delegates execute to captured web_search', async () => { + setWebSearchDelegateForTests(async (_id, params) => ({ + content: [{ type: 'text', text: `delegated:${JSON.stringify(params)}` }], + details: { delegated: true }, + })); + + const tools = collectTools(registerWebSearchTool); + const result = await executeTool(tools.get('WebSearch'), { query: 'pi extensions' }, '/tmp'); + expect(firstText(result)).toBe('delegated:{"query":"pi extensions"}'); + expect(result.details).toEqual({ delegated: true }); + }); + + it('reports missing pi-web-access when delegate was never captured', async () => { + vi.spyOn(webSearchDelegate, 'ensureWebSearchDelegate').mockResolvedValue(undefined); + vi.spyOn(webSearchDelegate, 'getWebSearchDelegate').mockReturnValue(undefined); + vi.spyOn(webSearchDelegate, 'getWebSearchLoadError').mockReturnValue( + 'pi-web-access is not installed. Run: pi install npm:pi-web-access', + ); + + const tools = collectTools(registerWebSearchTool); + const result = await executeTool(tools.get('WebSearch'), { query: 'test' }, '/tmp'); + expect(firstText(result)).toMatch(/pi-web-access|pi install npm:pi-web-access/i); + + vi.restoreAllMocks(); + }); + + it('registerGrokTools skips WebSearch when pi-web-access is not installed', () => { + vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(false); + const names: string[] = []; + registerGrokTools({ + registerTool(tool: { name: string }) { + names.push(tool.name); + }, + on() {}, + } as unknown as ExtensionAPI); + expect(names).not.toContain('WebSearch'); + vi.restoreAllMocks(); + }); + + it('renderCall shows WebSearch title', () => { + const tools = collectTools(registerWebSearchTool); + const line = renderToolCall(tools.get('WebSearch'), { query: 'hello world' }); + expect(line).toContain('WebSearch'); + expect(line).toContain('hello world'); + }); +}); diff --git a/tests/tools/webSearchDelegate.test.ts b/tests/tools/webSearchDelegate.test.ts new file mode 100644 index 0000000..7a829f7 --- /dev/null +++ b/tests/tools/webSearchDelegate.test.ts @@ -0,0 +1,35 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { + clearWebSearchDelegateForTests, + getWebSearchDelegate, + isPiWebAccessInstalled, + setWebSearchDelegateForTests, +} from '../../src/tools/webSearchDelegate.js'; + +afterEach(() => { + clearWebSearchDelegateForTests(); + vi.restoreAllMocks(); +}); + +describe('webSearchDelegate', () => { + it('returns delegate set for tests', async () => { + setWebSearchDelegateForTests(async () => ({ + content: [{ type: 'text', text: 'ok' }], + details: {}, + })); + + const delegate = getWebSearchDelegate(); + expect(delegate).toBeTypeOf('function'); + if (!delegate) throw new Error('expected delegate'); + const result = await delegate('id', {}, new AbortController().signal, undefined, { + cwd: '/tmp', + hasUI: false, + } as import('@earendil-works/pi-coding-agent').ExtensionContext); + expect(result.content[0]?.text).toBe('ok'); + }); + + it('isPiWebAccessInstalled reflects agent install path', () => { + const installed = isPiWebAccessInstalled(); + expect(typeof installed).toBe('boolean'); + }); +}); From 565cd5b89c9e7c3d89f78d84d6b237a6682dc182 Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Wed, 3 Jun 2026 16:32:21 +0900 Subject: [PATCH 2/9] feat: add WebSearch shim and suppress web_search for grok-cli Register WebSearch when pi-web-access is optional-installed, delegate to its web_search tool, and reconcile active tools so grok-cli uses WebSearch instead of web_search. Bind the delegate on session_start. --- package.json | 8 +- src/provider/register.ts | 28 +++-- src/provider/toolScope.ts | 24 +++- src/tools/register.ts | 19 ++- src/tools/webSearch.ts | 159 +++++++++++++++++++++++++ src/tools/webSearchDelegate.ts | 204 +++++++++++++++++++++++++++++++++ 6 files changed, 427 insertions(+), 15 deletions(-) create mode 100644 src/tools/webSearch.ts create mode 100644 src/tools/webSearchDelegate.ts diff --git a/package.json b/package.json index a47171f..13c0301 100644 --- a/package.json +++ b/package.json @@ -52,7 +52,13 @@ "peerDependencies": { "@earendil-works/pi-ai": "*", "@earendil-works/pi-coding-agent": "*", - "@earendil-works/pi-tui": "*" + "@earendil-works/pi-tui": "*", + "pi-web-access": "*" + }, + "peerDependenciesMeta": { + "pi-web-access": { + "optional": true + } }, "devDependencies": { "@biomejs/biome": "2.4.16", diff --git a/src/provider/register.ts b/src/provider/register.ts index 075d1b1..8b40592 100644 --- a/src/provider/register.ts +++ b/src/provider/register.ts @@ -5,6 +5,11 @@ import { getBaseUrl, type XaiOAuthCredentials } from '../auth/oauth.js'; import { type GrokCliModelConfig, resolveModels } from '../models/catalog.js'; import { sanitizePayload } from '../payload/sanitize.js'; import { registerGrokTools } from '../tools/register.js'; +import { + bindLivePiWebAccess, + ensureWebSearchDelegate, + isPiWebAccessInstalled, +} from '../tools/webSearchDelegate.js'; import { loadQuotaCache } from './quota.js'; import { registerStatusCommand } from './status.js'; import { streamGrokCli } from './stream.js'; @@ -69,6 +74,20 @@ export default function registerGrokCli(pi: ExtensionAPI) { registerGrokTools(pi); + pi.on('session_start', async (_event, ctx) => { + if (process.env.GROK_CLI_OAUTH_TOKEN) { + ctx.ui.notify( + '[pi-grok-cli] Using GROK_CLI_OAUTH_TOKEN bypass — no auto-refresh, no model discovery', + 'warning', + ); + } + + if (!isPiWebAccessInstalled()) return; + + bindLivePiWebAccess(pi); + await ensureWebSearchDelegate(pi); + }); + pi.on('before_provider_request', (event, ctx) => { if (ctx.model?.provider !== 'grok-cli') return; @@ -78,13 +97,4 @@ export default function registerGrokCli(pi: ExtensionAPI) { }); registerStatusCommand(pi); - - if (process.env.GROK_CLI_OAUTH_TOKEN) { - pi.on('session_start', async (_event, ctx) => { - ctx.ui.notify( - '[pi-grok-cli] Using GROK_CLI_OAUTH_TOKEN bypass — no auto-refresh, no model discovery', - 'warning', - ); - }); - } } diff --git a/src/provider/toolScope.ts b/src/provider/toolScope.ts index 5730cc6..943b97b 100644 --- a/src/provider/toolScope.ts +++ b/src/provider/toolScope.ts @@ -1,13 +1,31 @@ import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { GROK_TOOL_NAMES } from '../tools/register.js'; +import { + GROK_SUPPRESSED_TOOL_NAMES, + GROK_TOOL_NAMES_FOR_SCOPE, + grokToolsToActivate, +} from '../tools/register.js'; export function syncGrokTools( pi: Pick, provider: string | undefined, ) { const currentTools = pi.getActiveTools(); - const baseTools = currentTools.filter((toolName) => !GROK_TOOL_NAMES.includes(toolName)); - const nextTools = provider === 'grok-cli' ? [...baseTools, ...GROK_TOOL_NAMES] : baseTools; + const baseTools = currentTools.filter( + (toolName) => + !GROK_TOOL_NAMES_FOR_SCOPE.includes(toolName as (typeof GROK_TOOL_NAMES_FOR_SCOPE)[number]), + ); + const nextTools = + provider === 'grok-cli' + ? [ + ...baseTools.filter( + (toolName) => + !GROK_SUPPRESSED_TOOL_NAMES.includes( + toolName as (typeof GROK_SUPPRESSED_TOOL_NAMES)[number], + ), + ), + ...grokToolsToActivate(), + ] + : baseTools; if ( currentTools.length === nextTools.length && diff --git a/src/tools/register.ts b/src/tools/register.ts index b981582..5553549 100644 --- a/src/tools/register.ts +++ b/src/tools/register.ts @@ -2,8 +2,11 @@ import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; import { registerFileTools } from './files.js'; import { registerSearchTools } from './search.js'; import { registerShellTool } from './shell.js'; +import { registerWebSearchTool } from './webSearch.js'; +import { isPiWebAccessInstalled } from './webSearchDelegate.js'; -export const GROK_TOOL_NAMES = [ +/** Grok/Cursor shims always registered by this extension (excludes optional WebSearch). */ +export const GROK_SHIM_TOOL_NAMES = [ 'Grep', 'Glob', 'LS', @@ -13,9 +16,21 @@ export const GROK_TOOL_NAMES = [ 'Edit', 'Delete', 'Shell', -]; +] as const; + +/** All shim names used when reconciling the active tool set (includes optional WebSearch). */ +export const GROK_TOOL_NAMES_FOR_SCOPE = [...GROK_SHIM_TOOL_NAMES, 'WebSearch'] as const; + +export const GROK_SUPPRESSED_TOOL_NAMES = ['web_search'] as const; + +export function grokToolsToActivate() { + const names: string[] = [...GROK_SHIM_TOOL_NAMES]; + if (isPiWebAccessInstalled()) names.push('WebSearch'); + return names; +} export function registerGrokTools(pi: ExtensionAPI) { + if (isPiWebAccessInstalled()) registerWebSearchTool(pi); registerSearchTools(pi); registerFileTools(pi); registerShellTool(pi); diff --git a/src/tools/webSearch.ts b/src/tools/webSearch.ts new file mode 100644 index 0000000..0ed7ab1 --- /dev/null +++ b/src/tools/webSearch.ts @@ -0,0 +1,159 @@ +import { StringEnum, Type } from '@earendil-works/pi-ai'; +import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; +import { Text } from '@earendil-works/pi-tui'; +import { renderRunning, text } from './rendering.js'; +import { + ensureWebSearchDelegate, + getWebSearchDelegate, + getWebSearchLoadError, + PI_WEB_SEARCH_TOOL, +} from './webSearchDelegate.js'; + +const WEB_SEARCH_DESCRIPTION = + 'Search the web using Perplexity AI, Exa, or Gemini. Returns an AI-synthesized answer with source citations. For comprehensive research, prefer queries (plural) with 2-4 varied angles over a single query — each query gets its own synthesized answer, so varying phrasing and scope gives much broader coverage. When includeContent is true, full page content is fetched in the background. Searches auto-open the interactive browser curator and stream results live; set workflow to "none" to skip curation. Provider auto-selects: Exa (direct API with key, MCP fallback without), else Perplexity (needs key), else Gemini API (needs key), else Gemini Web (needs a supported Chromium-based browser login).'; + +const WebSearchParams = Type.Object({ + query: Type.Optional( + Type.String({ + description: + "Single search query. For research tasks, prefer 'queries' with multiple varied angles instead.", + }), + ), + queries: Type.Optional( + Type.Array(Type.String(), { + description: + 'Multiple queries searched in sequence, each returning its own synthesized answer. Prefer this for research — vary phrasing, scope, and angle across 2-4 queries to maximize coverage.', + }), + ), + numResults: Type.Optional( + Type.Number({ description: 'Results per query (default: 5, max: 20)' }), + ), + includeContent: Type.Optional(Type.Boolean({ description: 'Fetch full page content (async)' })), + recencyFilter: Type.Optional( + StringEnum(['day', 'week', 'month', 'year'], { description: 'Filter by recency' }), + ), + domainFilter: Type.Optional( + Type.Array(Type.String(), { + description: 'Limit to domains (prefix with - to exclude)', + }), + ), + provider: Type.Optional( + StringEnum(['auto', 'perplexity', 'gemini', 'exa'], { + description: 'Search provider (default: auto)', + }), + ), + workflow: Type.Optional( + StringEnum(['none', 'summary-review'], { + description: + 'Search workflow mode: none = no curator, summary-review = open curator with auto summary draft (default)', + }), + ), +}); + +function normalizeQueryList(raw: unknown[]): string[] { + return raw + .filter((q): q is string => typeof q === 'string') + .map((q) => q.trim()) + .filter((q) => q.length > 0); +} + +function queryListFromArgs(args: Record) { + const raw: unknown[] = Array.isArray(args.queries) + ? args.queries + : args.query !== undefined + ? [args.query] + : []; + return normalizeQueryList(raw); +} + +function missingDelegateMessage() { + const reason = + getWebSearchLoadError() ?? + 'pi-web-access web_search delegate not available. Install with: pi install npm:pi-web-access'; + return { + content: [ + { + type: 'text' as const, + text: `WebSearch requires pi-web-access: ${reason}`, + }, + ], + details: { error: reason }, + }; +} + +export function registerWebSearchTool(pi: ExtensionAPI) { + pi.registerTool({ + name: 'WebSearch', + label: 'Web Search', + description: WEB_SEARCH_DESCRIPTION, + promptSnippet: + 'Use for web research questions. Prefer {queries:[...]} with 2-4 varied angles over a single query for broader coverage.', + parameters: WebSearchParams, + + async execute(toolCallId, params, signal, onUpdate, ctx) { + await ensureWebSearchDelegate(pi); + const delegate = getWebSearchDelegate(); + if (!delegate) return missingDelegateMessage(); + return delegate(toolCallId, params as Record, signal, onUpdate, ctx); + }, + + renderCall(args, theme) { + const queryList = queryListFromArgs(args as Record); + if (queryList.length === 0) { + return text( + theme.fg('toolTitle', theme.bold('WebSearch ')) + theme.fg('error', '(no query)'), + ); + } + if (queryList.length === 1) { + const q = queryList[0]; + const display = q.length > 60 ? `${q.slice(0, 57)}...` : q; + return text( + theme.fg('toolTitle', theme.bold('WebSearch ')) + theme.fg('accent', `"${display}"`), + ); + } + const lines = [ + theme.fg('toolTitle', theme.bold('WebSearch ')) + + theme.fg('accent', `${queryList.length} queries`), + ]; + for (const q of queryList.slice(0, 5)) { + const display = q.length > 50 ? `${q.slice(0, 47)}...` : q; + lines.push(theme.fg('muted', ` "${display}"`)); + } + if (queryList.length > 5) { + lines.push(theme.fg('muted', ` ... and ${queryList.length - 5} more`)); + } + return new Text(lines.join('\n'), 0, 0); + }, + + renderResult(result, { expanded, isPartial }, theme) { + const running = renderRunning(isPartial); + if (running) return running; + + const details = result.details as { error?: string; totalResults?: number } | undefined; + if (details?.error) { + return text(theme.fg('error', `Error: ${details.error}`)); + } + + const summary = + typeof details?.totalResults === 'number' + ? theme.fg('success', `${details.totalResults} sources`) + : theme.fg('success', 'search complete'); + + if (!expanded) return text(summary); + + const textContent = result.content.find((c) => c.type === 'text')?.text ?? ''; + const preview = textContent.length > 800 ? `${textContent.slice(0, 800)}...` : textContent; + return new Text(`${summary}\n${theme.fg('dim', preview)}`, 0, 0); + }, + }); + + pi.on('tool_call', (event, ctx) => { + if (ctx.model?.provider !== 'grok-cli') return; + if (event.toolName !== PI_WEB_SEARCH_TOOL) return; + return { + block: true, + reason: + 'web_search is disabled for Grok CLI; use WebSearch instead (same behavior as pi-web-access web_search).', + }; + }); +} diff --git a/src/tools/webSearchDelegate.ts b/src/tools/webSearchDelegate.ts new file mode 100644 index 0000000..e975698 --- /dev/null +++ b/src/tools/webSearchDelegate.ts @@ -0,0 +1,204 @@ +import { existsSync } from 'node:fs'; +import { createRequire } from 'node:module'; +import { homedir } from 'node:os'; +import { dirname, join } from 'node:path'; +import { fileURLToPath } from 'node:url'; +import type { AgentToolResult, AgentToolUpdateCallback } from '@earendil-works/pi-agent-core'; +import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; +import { createEventBus, getAgentDir } from '@earendil-works/pi-coding-agent'; +import { createJiti } from 'jiti/static'; + +export const PI_WEB_SEARCH_TOOL = 'web_search'; + +export type WebSearchExecute = ( + toolCallId: string, + params: Record, + signal: AbortSignal | undefined, + onUpdate: AgentToolUpdateCallback | undefined, + ctx: import('@earendil-works/pi-coding-agent').ExtensionContext, +) => Promise>; + +let webSearchExecute: WebSearchExecute | undefined; +let loadPromise: Promise | undefined; +let lastLoadError: string | undefined; +let boundLivePi: ExtensionAPI | undefined; + +export function getWebSearchLoadError() { + return lastLoadError; +} + +function resolvePiCodingAgentRoot() { + const mainEntry = fileURLToPath(import.meta.resolve('@earendil-works/pi-coding-agent')); + return join(dirname(mainEntry), '..'); +} + +async function importPiExtensionLoader() { + return import(join(resolvePiCodingAgentRoot(), 'dist/core/extensions/loader.js')) as Promise<{ + createExtensionRuntime: () => Record; + loadExtensionFromFactory: ( + factory: (api: ExtensionAPI) => void | Promise, + cwd: string, + eventBus: ReturnType, + runtime: Record, + extensionPath?: string, + ) => Promise<{ tools: Map }>; + }>; +} + +export function isPiWebAccessInstalled() { + return resolvePiWebAccessEntry() !== undefined; +} + +function resolvePiWebAccessEntry(): string | undefined { + const fileNames = ['index.ts', 'index.js']; + const dirs = [ + join(getAgentDir(), 'npm', 'node_modules', 'pi-web-access'), + join(homedir(), '.pi', 'agent', 'npm', 'node_modules', 'pi-web-access'), + ]; + + for (const dir of dirs) { + for (const file of fileNames) { + const entry = join(dir, file); + if (existsSync(entry)) return entry; + } + } + + return undefined; +} + +function createJitiAliases() { + const require = createRequire(import.meta.url); + const codingAgent = fileURLToPath(import.meta.resolve('@earendil-works/pi-coding-agent')); + const agentCore = fileURLToPath(import.meta.resolve('@earendil-works/pi-agent-core')); + const tui = fileURLToPath(import.meta.resolve('@earendil-works/pi-tui')); + const ai = fileURLToPath(import.meta.resolve('@earendil-works/pi-ai')); + const typeboxEntry = require.resolve('typebox'); + const typeboxCompileEntry = require.resolve('typebox/compile'); + const typeboxValueEntry = require.resolve('typebox/value'); + + return { + '@earendil-works/pi-coding-agent': codingAgent, + '@earendil-works/pi-agent-core': agentCore, + '@earendil-works/pi-tui': tui, + '@earendil-works/pi-ai': ai, + '@mariozechner/pi-coding-agent': codingAgent, + '@mariozechner/pi-agent-core': agentCore, + '@mariozechner/pi-tui': tui, + '@mariozechner/pi-ai': ai, + typebox: typeboxEntry, + 'typebox/compile': typeboxCompileEntry, + 'typebox/value': typeboxValueEntry, + '@sinclair/typebox': typeboxEntry, + '@sinclair/typebox/compile': typeboxCompileEntry, + '@sinclair/typebox/value': typeboxValueEntry, + }; +} + +async function importPiWebAccessFactory(entry: string) { + const jiti = createJiti(import.meta.url, { alias: createJitiAliases(), moduleCache: false }); + const module = await jiti.import(entry, { default: true }); + if (typeof module !== 'function') { + throw new Error('pi-web-access does not export a default factory function'); + } + return module as (api: ExtensionAPI) => void | Promise; +} + +function wireRuntimeToLivePi(runtime: Record, pi: ExtensionAPI) { + runtime.assertActive = () => {}; + runtime.refreshTools = () => {}; + runtime.appendEntry = (customType: string, data: unknown) => pi.appendEntry(customType, data); + runtime.sendMessage = (message: unknown, options?: unknown) => + pi.sendMessage( + message as Parameters[0], + options as Parameters[1], + ); + runtime.sendUserMessage = (content: unknown, options?: unknown) => + pi.sendUserMessage( + content as Parameters[0], + options as Parameters[1], + ); + runtime.setSessionName = (name: string) => pi.setSessionName(name); + runtime.getSessionName = () => pi.getSessionName(); + runtime.setLabel = (entryId: string, label: string) => pi.setLabel(entryId, label); + runtime.getActiveTools = () => pi.getActiveTools(); + runtime.getAllTools = () => pi.getAllTools(); + runtime.setActiveTools = (names: string[]) => pi.setActiveTools(names); + runtime.getCommands = () => pi.getCommands(); + runtime.setModel = (model: unknown) => + pi.setModel(model as Parameters[0]); + runtime.getThinkingLevel = () => pi.getThinkingLevel(); + runtime.setThinkingLevel = (level: unknown) => + pi.setThinkingLevel(level as Parameters[0]); +} + +/** Remember the live session ExtensionAPI (bound after session_start). */ +export function bindLivePiWebAccess(pi: ExtensionAPI) { + boundLivePi = pi; + webSearchExecute = undefined; + loadPromise = undefined; +} + +async function captureWebSearchFromLivePi(pi: ExtensionAPI) { + const entry = resolvePiWebAccessEntry(); + if (!entry) return; + + const { createExtensionRuntime, loadExtensionFromFactory } = await importPiExtensionLoader(); + const runtime = createExtensionRuntime(); + wireRuntimeToLivePi(runtime, pi); + + const factory = await importPiWebAccessFactory(entry); + const extension = await loadExtensionFromFactory( + factory, + process.cwd(), + createEventBus(), + runtime, + entry, + ); + + const registered = extension.tools.get(PI_WEB_SEARCH_TOOL); + if (!registered) { + lastLoadError = 'pi-web-access loaded but did not register web_search. Update pi-web-access.'; + return; + } + + webSearchExecute = registered.definition.execute.bind(registered.definition) as WebSearchExecute; + lastLoadError = undefined; +} + +export async function ensureWebSearchDelegate(pi?: ExtensionAPI) { + if (!isPiWebAccessInstalled()) return; + + const livePi = pi ?? boundLivePi; + if (!livePi) return; + + if (webSearchExecute) return; + if (loadPromise) return loadPromise; + + loadPromise = (async () => { + lastLoadError = undefined; + try { + await captureWebSearchFromLivePi(livePi); + } catch (err) { + lastLoadError = err instanceof Error ? err.message : String(err); + webSearchExecute = undefined; + } + })(); + + return loadPromise; +} + +export function getWebSearchDelegate() { + return webSearchExecute; +} + +export function clearWebSearchDelegateForTests() { + webSearchExecute = undefined; + loadPromise = undefined; + lastLoadError = undefined; + boundLivePi = undefined; +} + +export function setWebSearchDelegateForTests(execute: WebSearchExecute) { + webSearchExecute = execute; + lastLoadError = undefined; +} From 7a4a128139126850fd8cf5ec0e09830f10a22765 Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Wed, 3 Jun 2026 16:32:25 +0900 Subject: [PATCH 3/9] docs: describe WebSearch and pi-web-access tool behavior --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index 2e2dce3..238441e 100644 --- a/README.md +++ b/README.md @@ -21,8 +21,11 @@ Grok CLI models are trained to use Cursor-style coding tools. This extension inc - File tools: `Read`, `Write`, `StrReplace`, `Edit`, `Delete`, and `LS` - Search tools: `Grep` and `Glob` +- Web search: `WebSearch` only when [pi-web-access](https://www.npmjs.com/package/pi-web-access) is installed (`pi install npm:pi-web-access`); it delegates to that extension’s `web_search` - Terminal tool: `Shell` +When the active model is **grok-cli** and pi-web-access is installed, `web_search` is removed from the active tool set and blocked if invoked; use `WebSearch` instead. If pi-web-access is not installed, `WebSearch` is not registered and nothing changes for web search. Other providers keep using `web_search` from pi-web-access when that extension is installed. + The shims also normalize common Cursor/Grok argument shapes, such as `contents` for writes, `glob_pattern` for file search, `glob_filter` for grep filters, and `old_string`/`new_string` or `oldText`/`newText` for exact replacements. This keeps agentic coding workflows moving instead of failing on tool schema mismatches. ## Requirements From b2eb29ce964347d2235344ae173a51f1b0c863c0 Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Wed, 3 Jun 2026 16:48:43 +0900 Subject: [PATCH 4/9] chore: improve test coverage from 80% to 85% (+27 tests) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - webSearch.ts: 47% → 100% (renderCall/renderResult branches, tool_call interceptor) - sanitize.ts: 91% → 96% (.jpg, file:// protocol, unsupported ext, string output parts) - rendering.ts: 91% → 95% (globToRegExp wildcards, grep fallback path include) - search.ts: 91% → 93% (sort tiebreaker, grep fallback with path include) - shell.ts: added timeout kill test - webSearchDelegate.ts: added bind/ensure/getLoadError tests --- tests/payload/sanitize.test.ts | 155 +++++++++++++++++++++++++ tests/tools/rendering.test.ts | 10 ++ tests/tools/search.test.ts | 24 ++++ tests/tools/shell.test.ts | 13 +++ tests/tools/webSearch.test.ts | 159 +++++++++++++++++++++++++- tests/tools/webSearchDelegate.test.ts | 33 ++++++ 6 files changed, 393 insertions(+), 1 deletion(-) diff --git a/tests/payload/sanitize.test.ts b/tests/payload/sanitize.test.ts index f3935a5..22bcfd3 100644 --- a/tests/payload/sanitize.test.ts +++ b/tests/payload/sanitize.test.ts @@ -190,6 +190,161 @@ describe('payload sanitization', () => { } }); + it('resolves .jpg and .jpeg image paths to data URLs', () => { + const dir = mkdtempSync(join(tmpdir(), 'pi-grok-cli-test-')); + const jpgPath = join(dir, 'photo.jpg'); + const jpegPath = join(dir, 'photo.jpeg'); + writeFileSync(jpgPath, Buffer.from('jpg bytes')); + writeFileSync(jpegPath, Buffer.from('jpeg bytes')); + + try { + const jpgResult = sanitizePayload( + { + input: [ + { + role: 'user', + content: [{ type: 'input_image', image_url: jpgPath }], + }, + ], + }, + 'grok-4.3', + undefined, + dir, + ); + expect((jpgResult.input as Array>)[0].content).toEqual([ + { + type: 'input_image', + image_url: `data:image/jpeg;base64,${Buffer.from('jpg bytes').toString('base64')}`, + detail: 'auto', + }, + ]); + + const jpegResult = sanitizePayload( + { + input: [ + { + role: 'user', + content: [{ type: 'input_image', image_url: jpegPath }], + }, + ], + }, + 'grok-4.3', + undefined, + dir, + ); + expect((jpegResult.input as Array>)[0].content).toEqual([ + { + type: 'input_image', + image_url: `data:image/jpeg;base64,${Buffer.from('jpeg bytes').toString('base64')}`, + detail: 'auto', + }, + ]); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }); + + it('rejects unsupported local image extensions', () => { + const dir = mkdtempSync(join(tmpdir(), 'pi-grok-cli-test-')); + const gifPath = join(dir, 'animation.gif'); + writeFileSync(gifPath, Buffer.from('gif bytes')); + + try { + expect(() => + sanitizePayload( + { + input: [ + { + role: 'user', + content: [{ type: 'input_image', image_url: gifPath }], + }, + ], + }, + 'grok-4.3', + undefined, + dir, + ), + ).toThrow(/xAI image understanding supports local .jpg, .jpeg, and .png files only/); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }); + + it('resolves file:// protocol image paths', () => { + const dir = mkdtempSync(join(tmpdir(), 'pi-grok-cli-test-')); + const imagePath = join(dir, 'file-ref.png'); + writeFileSync(imagePath, Buffer.from('file ref png')); + + try { + const payload = sanitizePayload( + { + input: [ + { + role: 'user', + content: [{ type: 'input_image', image_url: `file://${imagePath}` }], + }, + ], + }, + 'grok-4.3', + undefined, + dir, + ); + + expect((payload.input as Array>)[0].content).toEqual([ + { + type: 'input_image', + image_url: `data:image/png;base64,${Buffer.from('file ref png').toString('base64')}`, + detail: 'auto', + }, + ]); + } finally { + rmSync(dir, { recursive: true, force: true }); + } + }); + + it('rejects invalid file:// URLs gracefully', () => { + expect(() => + sanitizePayload( + { + input: [ + { + role: 'user', + content: [{ type: 'input_image', image_url: 'file://invalid-url' }], + }, + ], + }, + 'grok-4.3', + undefined, + process.cwd(), + ), + ).toThrow('Image file does not exist or is not a valid URL: file://invalid-url'); + }); + + it('rewrites function_call_output with plain string parts', () => { + const payload = sanitizePayload( + { + input: [ + { + type: 'function_call_output', + call_id: 'call_s', + output: ['plain string output', { type: 'input_text', text: 'object output' }], + }, + ], + }, + 'grok-4.3', + undefined, + process.cwd(), + ); + + expect(payload.input).toEqual([ + { + type: 'function_call_output', + call_id: 'call_s', + output: 'plain string output\nobject output', + }, + ]); + }); + it('rejects missing or unsupported local images', () => { expect(() => sanitizePayload( diff --git a/tests/tools/rendering.test.ts b/tests/tools/rendering.test.ts index 2f673f2..e294929 100644 --- a/tests/tools/rendering.test.ts +++ b/tests/tools/rendering.test.ts @@ -4,6 +4,7 @@ import { detailRecord, fileError, fileNotFound, + globToRegExp, MAX_LINES, MAX_OUTPUT_CHARS, numberDetail, @@ -65,6 +66,15 @@ describe('tool rendering helpers', () => { expect(booleanDetail(result, 'invalid')).toBe(false); }); + it('converts glob patterns to regular expressions', () => { + expect(globToRegExp('a**b').test('aXYb')).toBe(true); + expect(globToRegExp('a**b').test('a/b')).toBe(true); + expect(globToRegExp('a?b').test('aXb')).toBe(true); + expect(globToRegExp('a?b').test('a/b')).toBe(false); + expect(globToRegExp('a**/b').test('a/x/y/b')).toBe(true); + expect(globToRegExp('a**/b').test('a/x/y/z/b')).toBe(true); + }); + it('formats file and command errors with stable empty details', () => { expect(fileNotFound('/tmp/missing.txt', { deleted: false })).toEqual({ content: [{ type: 'text', text: 'File not found: /tmp/missing.txt' }], diff --git a/tests/tools/search.test.ts b/tests/tools/search.test.ts index 8237b7a..22d3931 100644 --- a/tests/tools/search.test.ts +++ b/tests/tools/search.test.ts @@ -216,6 +216,19 @@ describe('search tools', () => { }); }); + it('greps with path-containing include patterns through the fallback', async () => { + const cwd = setupProject(); + await withNoSearchBinaries(async (fallbackTools) => { + const result = await executeTool( + fallbackTools.get('Grep'), + { pattern: 'needle', path: 'src', include: 'src/**/*.ts' }, + cwd, + ); + + expectGrepResult(cwd, result); + }); + }); + it('sorts glob results by modification time newest first', async () => { const cwd = setupProject(); const oldTime = new Date('2024-01-01T00:00:00.000Z'); @@ -246,6 +259,17 @@ describe('search tools', () => { ]); }); + it('breaks modification time ties with alphabetical ordering', () => { + const cwd = setupProject(); + const sameTime = new Date('2024-06-01T00:00:00.000Z'); + utimesSync(join(cwd, 'src', 'gamma.ts'), sameTime, sameTime); + utimesSync(join(cwd, 'src', 'alpha.ts'), sameTime, sameTime); + + expect( + sortByModifiedNewest([join(cwd, 'src', 'gamma.ts'), join(cwd, 'src', 'alpha.ts')]), + ).toEqual([join(cwd, 'src', 'alpha.ts'), join(cwd, 'src', 'gamma.ts')]); + }); + it('renders grep calls and result states', () => { const grep = collectTools(registerSearchTools).get('Grep'); const result = { diff --git a/tests/tools/shell.test.ts b/tests/tools/shell.test.ts index 2fc519b..209bbd1 100644 --- a/tests/tools/shell.test.ts +++ b/tests/tools/shell.test.ts @@ -102,6 +102,19 @@ describe('shell tool', () => { expect(firstText(result).endsWith('[Output truncated at 50KB]')).toBe(true); }); + it('kills commands that exceed the timeout', async () => { + const cwd = tempDir('pi-grok-cli-shell-'); + const result = await executeTool( + collectTools(registerShellTool).get('Shell'), + { command: 'sleep 10', timeout: 100 }, + cwd, + ); + + expect(firstText(result)).toContain('Shell error'); + expect(result.details).toMatchObject({ command: 'sleep 10' }); + expect(result.details.exitCode).toBeDefined(); + }); + it('renders shell calls and result states', () => { const shell = collectTools(registerShellTool).get('Shell'); diff --git a/tests/tools/webSearch.test.ts b/tests/tools/webSearch.test.ts index 81f5899..5dc48c1 100644 --- a/tests/tools/webSearch.test.ts +++ b/tests/tools/webSearch.test.ts @@ -5,9 +5,16 @@ import { registerWebSearchTool } from '../../src/tools/webSearch.js'; import * as webSearchDelegate from '../../src/tools/webSearchDelegate.js'; import { clearWebSearchDelegateForTests, + PI_WEB_SEARCH_TOOL, setWebSearchDelegateForTests, } from '../../src/tools/webSearchDelegate.js'; -import { collectTools, executeTool, firstText, renderToolCall } from './toolTestHelpers.js'; +import { + collectTools, + executeTool, + firstText, + renderToolCall, + renderToolResult, +} from './toolTestHelpers.js'; afterEach(() => { clearWebSearchDelegateForTests(); @@ -73,4 +80,154 @@ describe('WebSearch tool', () => { expect(line).toContain('WebSearch'); expect(line).toContain('hello world'); }); + + describe('renderCall', () => { + const tools = collectTools(registerWebSearchTool); + const ws = tools.get('WebSearch'); + + it('shows (no query) when query list is empty', () => { + expect(renderToolCall(ws, { queries: [] })).toContain('(no query)'); + }); + + it('shows (no query) with no query or queries param', () => { + expect(renderToolCall(ws, {})).toContain('(no query)'); + }); + + it('truncates a single long query', () => { + const long = 'a'.repeat(61); + const line = renderToolCall(ws, { query: long }); + expect(line).toContain('...'); + expect(line).not.toContain(long); + }); + + it('shows query count and list for multiple queries', () => { + const line = renderToolCall(ws, { queries: ['q1', 'q2', 'q3'] }); + expect(line).toContain('3 queries'); + expect(line).toContain('"q1"'); + expect(line).toContain('"q2"'); + expect(line).toContain('"q3"'); + }); + + it('shows truncation indicator for more than 5 queries', () => { + const queries = ['a', 'b', 'c', 'd', 'e', 'f', 'g']; + const line = renderToolCall(ws, { queries }); + expect(line).toContain('7 queries'); + expect(line).toContain('... and 2 more'); + }); + }); + + describe('renderResult', () => { + const tools = collectTools(registerWebSearchTool); + const ws = tools.get('WebSearch'); + + it('shows running state when partial', () => { + const result = { + content: [{ type: 'text', text: 'loading...' }], + details: {}, + }; + expect(renderToolResult(ws, result, { expanded: false, isPartial: true })).toBe('Running...'); + }); + + it('shows error message when details contain error', () => { + const result = { + content: [{ type: 'text', text: 'something broke' }], + details: { error: 'API key is invalid' }, + }; + const rendered = renderToolResult(ws, result); + expect(rendered).toContain('Error:'); + expect(rendered).toContain('API key is invalid'); + }); + + it('shows source count when totalResults is present', () => { + const result = { + content: [{ type: 'text', text: 'synthesized answer' }], + details: { totalResults: 12 }, + }; + expect(renderToolResult(ws, result)).toBe('12 sources'); + }); + + it('shows search complete when no totalResults', () => { + const result = { + content: [{ type: 'text', text: 'answer' }], + details: {}, + }; + expect(renderToolResult(ws, result)).toBe('search complete'); + }); + + it('shows expanded content with summary prefix', () => { + const result = { + content: [{ type: 'text', text: 'the full search answer' }], + details: { totalResults: 3 }, + }; + const rendered = renderToolResult(ws, result, { + expanded: true, + isPartial: false, + }); + expect(rendered).toContain('3 sources'); + expect(rendered).toContain('the full search answer'); + }); + + it('truncates long expanded text', () => { + const longText = 'x'.repeat(801); + const result = { + content: [{ type: 'text', text: longText }], + details: { totalResults: 1 }, + }; + const rendered = renderToolResult(ws, result, { + expanded: true, + isPartial: false, + }); + expect(rendered).toContain('...'); + expect(rendered).not.toContain(longText); + }); + }); + + describe('tool_call interceptor', () => { + function getHandler() { + const handlers = new Map unknown>(); + registerWebSearchTool({ + registerTool() {}, + on(event: string, handler: (event: unknown, ctx: unknown) => unknown) { + handlers.set(event, handler); + }, + } as unknown as ExtensionAPI); + const handler = handlers.get('tool_call'); + if (!handler) throw new Error('tool_call handler not registered'); + return handler; + } + + it('blocks web_search for Grok CLI models', () => { + const result = getHandler()( + { toolName: PI_WEB_SEARCH_TOOL }, + { model: { provider: 'grok-cli' } }, + ) as { block?: boolean; reason?: string } | undefined; + + expect(result?.block).toBe(true); + expect(result?.reason).toContain('web_search is disabled for Grok CLI'); + }); + + it('allows web_search for non-Grok CLI models', () => { + const result = getHandler()( + { toolName: PI_WEB_SEARCH_TOOL }, + { model: { provider: 'openai' } }, + ); + + expect(result).toBeUndefined(); + }); + + it('does not block other tools for Grok CLI', () => { + const result = getHandler()( + { toolName: 'some_other_tool' }, + { model: { provider: 'grok-cli' } }, + ); + + expect(result).toBeUndefined(); + }); + + it('does not block web_search when model context is missing', () => { + const result = getHandler()({ toolName: PI_WEB_SEARCH_TOOL }, {}); + + expect(result).toBeUndefined(); + }); + }); }); diff --git a/tests/tools/webSearchDelegate.test.ts b/tests/tools/webSearchDelegate.test.ts index 7a829f7..3203321 100644 --- a/tests/tools/webSearchDelegate.test.ts +++ b/tests/tools/webSearchDelegate.test.ts @@ -1,7 +1,11 @@ import { afterEach, describe, expect, it, vi } from 'vitest'; +import * as delegateModule from '../../src/tools/webSearchDelegate.js'; import { + bindLivePiWebAccess, clearWebSearchDelegateForTests, + ensureWebSearchDelegate, getWebSearchDelegate, + getWebSearchLoadError, isPiWebAccessInstalled, setWebSearchDelegateForTests, } from '../../src/tools/webSearchDelegate.js'; @@ -32,4 +36,33 @@ describe('webSearchDelegate', () => { const installed = isPiWebAccessInstalled(); expect(typeof installed).toBe('boolean'); }); + + it('bindLivePiWebAccess resets delegate state', () => { + setWebSearchDelegateForTests(async () => ({ + content: [{ type: 'text', text: 'ok' }], + details: {}, + })); + expect(getWebSearchDelegate()).toBeTypeOf('function'); + + bindLivePiWebAccess({} as Parameters[0]); + + expect(getWebSearchDelegate()).toBeUndefined(); + }); + + it('ensureWebSearchDelegate returns undefined when pi-web-access is not installed', async () => { + vi.spyOn(delegateModule, 'isPiWebAccessInstalled').mockReturnValue(false); + + const result = await ensureWebSearchDelegate(); + expect(result).toBeUndefined(); + expect(getWebSearchDelegate()).toBeUndefined(); + + vi.restoreAllMocks(); + }); + + it('getWebSearchLoadError returns last error string', () => { + clearWebSearchDelegateForTests(); + // After clear, no error is set; the function should return undefined + const error = getWebSearchLoadError(); + expect(error).toBeUndefined(); + }); }); From 8ffc70a5e1c3a74ff84fdadd123973c427432ae3 Mon Sep 17 00:00:00 2001 From: J Liew Date: Wed, 3 Jun 2026 17:07:56 +0900 Subject: [PATCH 5/9] fix: make web search delegate loading resilient --- src/tools/webSearchDelegate.ts | 81 +++++++++++++++++----- tests/tools/webSearchDelegate.test.ts | 20 ++++++ tests/tools/webSearchDelegateRetry.test.ts | 53 ++++++++++++++ 3 files changed, 135 insertions(+), 19 deletions(-) create mode 100644 tests/tools/webSearchDelegateRetry.test.ts diff --git a/src/tools/webSearchDelegate.ts b/src/tools/webSearchDelegate.ts index e975698..720b791 100644 --- a/src/tools/webSearchDelegate.ts +++ b/src/tools/webSearchDelegate.ts @@ -1,11 +1,15 @@ -import { existsSync } from 'node:fs'; +import { existsSync, readFileSync } from 'node:fs'; import { createRequire } from 'node:module'; import { homedir } from 'node:os'; import { dirname, join } from 'node:path'; -import { fileURLToPath } from 'node:url'; +import { fileURLToPath, pathToFileURL } from 'node:url'; import type { AgentToolResult, AgentToolUpdateCallback } from '@earendil-works/pi-agent-core'; import type { ExtensionAPI } from '@earendil-works/pi-coding-agent'; -import { createEventBus, getAgentDir } from '@earendil-works/pi-coding-agent'; +import { + createEventBus, + createExtensionRuntime, + getAgentDir, +} from '@earendil-works/pi-coding-agent'; import { createJiti } from 'jiti/static'; export const PI_WEB_SEARCH_TOOL = 'web_search'; @@ -27,24 +31,61 @@ export function getWebSearchLoadError() { return lastLoadError; } -function resolvePiCodingAgentRoot() { - const mainEntry = fileURLToPath(import.meta.resolve('@earendil-works/pi-coding-agent')); - return join(dirname(mainEntry), '..'); +function resolvePiCodingAgentRoot(dir: string): string { + const packageJson = join(dir, 'package.json'); + if ( + existsSync(packageJson) && + JSON.parse(readFileSync(packageJson, 'utf8')).name === '@earendil-works/pi-coding-agent' + ) { + return dir; + } + + const parent = dirname(dir); + if (parent === dir) { + throw new Error(`Could not find @earendil-works/pi-coding-agent package root from ${dir}`); + } + return resolvePiCodingAgentRoot(parent); +} + +export function resolvePiExtensionLoaderPaths(mainEntry: string) { + const root = resolvePiCodingAgentRoot(dirname(mainEntry)); + return [ + join(root, 'dist/core/extensions/index.js'), + join(root, 'dist/core/extensions/loader.js'), + ]; } async function importPiExtensionLoader() { - return import(join(resolvePiCodingAgentRoot(), 'dist/core/extensions/loader.js')) as Promise<{ - createExtensionRuntime: () => Record; - loadExtensionFromFactory: ( - factory: (api: ExtensionAPI) => void | Promise, - cwd: string, - eventBus: ReturnType, - runtime: Record, - extensionPath?: string, - ) => Promise<{ tools: Map }>; - }>; + const publicModule = (await import('@earendil-works/pi-coding-agent')) as Record; + if ( + 'loadExtensionFromFactory' in publicModule && + typeof publicModule.loadExtensionFromFactory === 'function' + ) { + return publicModule.loadExtensionFromFactory as LoadExtensionFromFactory; + } + + const mainEntry = fileURLToPath(import.meta.resolve('@earendil-works/pi-coding-agent')); + const paths = resolvePiExtensionLoaderPaths(mainEntry); + const loaderPath = paths.find((path) => existsSync(path)); + if (!loaderPath) { + throw new Error(`Could not find pi extension loader. Attempted: ${paths.join(', ')}`); + } + + const loader = (await import(pathToFileURL(loaderPath).href)) as Record; + if (typeof loader.loadExtensionFromFactory !== 'function') { + throw new Error(`Pi extension loader does not export loadExtensionFromFactory: ${loaderPath}`); + } + return loader.loadExtensionFromFactory as LoadExtensionFromFactory; } +type LoadExtensionFromFactory = ( + factory: (api: ExtensionAPI) => void | Promise, + cwd: string, + eventBus: ReturnType, + runtime: Record, + extensionPath?: string, +) => Promise<{ tools: Map }>; + export function isPiWebAccessInstalled() { return resolvePiWebAccessEntry() !== undefined; } @@ -142,16 +183,16 @@ async function captureWebSearchFromLivePi(pi: ExtensionAPI) { const entry = resolvePiWebAccessEntry(); if (!entry) return; - const { createExtensionRuntime, loadExtensionFromFactory } = await importPiExtensionLoader(); + const loadExtensionFromFactory = await importPiExtensionLoader(); const runtime = createExtensionRuntime(); - wireRuntimeToLivePi(runtime, pi); + wireRuntimeToLivePi(runtime as unknown as Record, pi); const factory = await importPiWebAccessFactory(entry); const extension = await loadExtensionFromFactory( factory, process.cwd(), createEventBus(), - runtime, + runtime as unknown as Record, entry, ); @@ -181,6 +222,8 @@ export async function ensureWebSearchDelegate(pi?: ExtensionAPI) { } catch (err) { lastLoadError = err instanceof Error ? err.message : String(err); webSearchExecute = undefined; + } finally { + loadPromise = undefined; } })(); diff --git a/tests/tools/webSearchDelegate.test.ts b/tests/tools/webSearchDelegate.test.ts index 3203321..5c30516 100644 --- a/tests/tools/webSearchDelegate.test.ts +++ b/tests/tools/webSearchDelegate.test.ts @@ -1,3 +1,6 @@ +import { mkdirSync, mkdtempSync, writeFileSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; import { afterEach, describe, expect, it, vi } from 'vitest'; import * as delegateModule from '../../src/tools/webSearchDelegate.js'; import { @@ -7,6 +10,7 @@ import { getWebSearchDelegate, getWebSearchLoadError, isPiWebAccessInstalled, + resolvePiExtensionLoaderPaths, setWebSearchDelegateForTests, } from '../../src/tools/webSearchDelegate.js'; @@ -65,4 +69,20 @@ describe('webSearchDelegate', () => { const error = getWebSearchLoadError(); expect(error).toBeUndefined(); }); + + it('resolves extension loader paths from a nested package main entry', () => { + const packageRoot = join(mkdtempSync(join(tmpdir(), 'pi-grok-cli-')), 'pi-coding-agent'); + const mainEntry = join(packageRoot, 'dist', 'esm', 'index.js'); + mkdirSync(join(packageRoot, 'dist', 'core', 'extensions'), { recursive: true }); + mkdirSync(join(packageRoot, 'dist', 'esm'), { recursive: true }); + writeFileSync( + join(packageRoot, 'package.json'), + JSON.stringify({ name: '@earendil-works/pi-coding-agent' }), + ); + + expect(resolvePiExtensionLoaderPaths(mainEntry)).toEqual([ + join(packageRoot, 'dist', 'core', 'extensions', 'index.js'), + join(packageRoot, 'dist', 'core', 'extensions', 'loader.js'), + ]); + }); }); diff --git a/tests/tools/webSearchDelegateRetry.test.ts b/tests/tools/webSearchDelegateRetry.test.ts new file mode 100644 index 0000000..36accc7 --- /dev/null +++ b/tests/tools/webSearchDelegateRetry.test.ts @@ -0,0 +1,53 @@ +import { mkdirSync, mkdtempSync, writeFileSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { afterEach, describe, expect, it, vi } from 'vitest'; + +const testAgentDir = mkdtempSync(join(tmpdir(), 'pi-grok-cli-')); + +vi.mock('@earendil-works/pi-coding-agent', async (importOriginal) => ({ + ...(await importOriginal()), + getAgentDir: () => testAgentDir, +})); + +import { + clearWebSearchDelegateForTests, + ensureWebSearchDelegate, + getWebSearchDelegate, + getWebSearchLoadError, +} from '../../src/tools/webSearchDelegate.js'; + +afterEach(() => { + clearWebSearchDelegateForTests(); + vi.unstubAllGlobals(); +}); + +describe('webSearchDelegate retry', () => { + it('retries after a failed delegate load', async () => { + const extensionDir = join(testAgentDir, 'npm', 'node_modules', 'pi-web-access'); + mkdirSync(extensionDir, { recursive: true }); + writeFileSync( + join(extensionDir, 'index.js'), + ` +export default function (pi) { + globalThis.webSearchDelegateLoadAttempts = (globalThis.webSearchDelegateLoadAttempts ?? 0) + 1 + if (globalThis.webSearchDelegateLoadAttempts === 1) throw new Error('temporary load failure') + pi.registerTool({ + name: 'web_search', + execute: async () => ({ content: [{ type: 'text', text: 'ok' }], details: {} }), + }) +} +`, + ); + vi.stubGlobal('webSearchDelegateLoadAttempts', 0); + const pi = {} as Parameters[0]; + + await ensureWebSearchDelegate(pi); + expect(getWebSearchDelegate()).toBeUndefined(); + expect(getWebSearchLoadError()).toBe('temporary load failure'); + + await ensureWebSearchDelegate(pi); + expect(getWebSearchDelegate()).toBeTypeOf('function'); + expect(getWebSearchLoadError()).toBeUndefined(); + }); +}); From dd079604b0ec6aca24b2518b1084c2b9e2dbd14a Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Wed, 3 Jun 2026 17:23:36 +0900 Subject: [PATCH 6/9] fix(webSearch): normalize query params before delegation --- src/tools/webSearch.ts | 11 ++++++++++- tests/tools/webSearch.test.ts | 25 +++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/tools/webSearch.ts b/src/tools/webSearch.ts index 0ed7ab1..5f0714f 100644 --- a/src/tools/webSearch.ts +++ b/src/tools/webSearch.ts @@ -94,7 +94,16 @@ export function registerWebSearchTool(pi: ExtensionAPI) { await ensureWebSearchDelegate(pi); const delegate = getWebSearchDelegate(); if (!delegate) return missingDelegateMessage(); - return delegate(toolCallId, params as Record, signal, onUpdate, ctx); + const normalizedParams = { ...(params as Record) }; + if (Array.isArray(normalizedParams.queries)) { + normalizedParams.queries = normalizeQueryList(normalizedParams.queries); + } + if (typeof normalizedParams.query === 'string') { + const query = normalizedParams.query.trim(); + if (query) normalizedParams.query = query; + if (!query) delete normalizedParams.query; + } + return delegate(toolCallId, normalizedParams, signal, onUpdate, ctx); }, renderCall(args, theme) { diff --git a/tests/tools/webSearch.test.ts b/tests/tools/webSearch.test.ts index 5dc48c1..37770ed 100644 --- a/tests/tools/webSearch.test.ts +++ b/tests/tools/webSearch.test.ts @@ -47,6 +47,31 @@ describe('WebSearch tool', () => { expect(result.details).toEqual({ delegated: true }); }); + it('normalizes delegated queries', async () => { + setWebSearchDelegateForTests(async (_id, params) => ({ + content: [{ type: 'text', text: JSON.stringify(params) }], + details: {}, + })); + + const tools = collectTools(registerWebSearchTool); + const result = await executeTool( + tools.get('WebSearch'), + { + query: ' ', + queries: [' first query ', ' ', 'second query'], + numResults: 3, + }, + '/tmp', + ); + + expect(firstText(result)).toBe( + JSON.stringify({ + queries: ['first query', 'second query'], + numResults: 3, + }), + ); + }); + it('reports missing pi-web-access when delegate was never captured', async () => { vi.spyOn(webSearchDelegate, 'ensureWebSearchDelegate').mockResolvedValue(undefined); vi.spyOn(webSearchDelegate, 'getWebSearchDelegate').mockReturnValue(undefined); From 401a2c5ac56c8219adb9d44a5e9f60e296b8c663 Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Wed, 3 Jun 2026 17:23:44 +0900 Subject: [PATCH 7/9] fix(provider): restore suppressed tools after grok provider round-trip --- src/provider/toolScope.ts | 23 ++++++++++++++++------- tests/provider/toolScope.test.ts | 19 +++++++++++++++++++ 2 files changed, 35 insertions(+), 7 deletions(-) diff --git a/src/provider/toolScope.ts b/src/provider/toolScope.ts index 943b97b..065d7f4 100644 --- a/src/provider/toolScope.ts +++ b/src/provider/toolScope.ts @@ -5,6 +5,8 @@ import { grokToolsToActivate, } from '../tools/register.js'; +const preservedSuppressedTools = new WeakMap(); + export function syncGrokTools( pi: Pick, provider: string | undefined, @@ -14,18 +16,25 @@ export function syncGrokTools( (toolName) => !GROK_TOOL_NAMES_FOR_SCOPE.includes(toolName as (typeof GROK_TOOL_NAMES_FOR_SCOPE)[number]), ); + const suppressedTools = baseTools.filter((toolName) => + GROK_SUPPRESSED_TOOL_NAMES.includes(toolName as (typeof GROK_SUPPRESSED_TOOL_NAMES)[number]), + ); + if (suppressedTools.length > 0) preservedSuppressedTools.set(pi, suppressedTools); + const nextTools = provider === 'grok-cli' ? [ - ...baseTools.filter( - (toolName) => - !GROK_SUPPRESSED_TOOL_NAMES.includes( - toolName as (typeof GROK_SUPPRESSED_TOOL_NAMES)[number], - ), - ), + ...baseTools.filter((toolName) => !suppressedTools.includes(toolName)), ...grokToolsToActivate(), ] - : baseTools; + : [ + ...baseTools, + ...(preservedSuppressedTools.get(pi) ?? []).filter( + (toolName) => !baseTools.includes(toolName), + ), + ]; + + if (provider !== 'grok-cli') preservedSuppressedTools.delete(pi); if ( currentTools.length === nextTools.length && diff --git a/tests/provider/toolScope.test.ts b/tests/provider/toolScope.test.ts index 8254a06..3c1cd1e 100644 --- a/tests/provider/toolScope.test.ts +++ b/tests/provider/toolScope.test.ts @@ -46,4 +46,23 @@ describe('syncGrokTools', () => { expect(setActiveTools).toHaveBeenCalledWith(['read', 'web_search']); }); + + it('restores suppressed tools after a provider round-trip', () => { + vi.spyOn(webSearchDelegate, 'isPiWebAccessInstalled').mockReturnValue(true); + const activeTools = ['read', 'web_search', 'bash']; + const pi = { + getActiveTools: () => activeTools, + setActiveTools(nextTools: string[]) { + activeTools.splice(0, activeTools.length, ...nextTools); + }, + }; + + syncGrokTools(pi, 'grok-cli'); + expect(activeTools).not.toContain('web_search'); + expect(activeTools).toContain('WebSearch'); + + syncGrokTools(pi, 'openai'); + expect(activeTools).toContain('web_search'); + expect(activeTools).not.toContain('WebSearch'); + }); }); From 6fb4aa9ae0bc0e6a0d8f3ca0151920e0ef42b420 Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Wed, 3 Jun 2026 17:23:48 +0900 Subject: [PATCH 8/9] fix(webSearch): ignore stale delegate loads after session rebind --- src/tools/webSearchDelegate.ts | 27 ++++++++++--- tests/tools/webSearchDelegate.test.ts | 9 ++--- tests/tools/webSearchDelegateRetry.test.ts | 47 ++++++++++++++++++++++ 3 files changed, 71 insertions(+), 12 deletions(-) diff --git a/src/tools/webSearchDelegate.ts b/src/tools/webSearchDelegate.ts index 720b791..1916f2e 100644 --- a/src/tools/webSearchDelegate.ts +++ b/src/tools/webSearchDelegate.ts @@ -26,11 +26,16 @@ let webSearchExecute: WebSearchExecute | undefined; let loadPromise: Promise | undefined; let lastLoadError: string | undefined; let boundLivePi: ExtensionAPI | undefined; +let bindGeneration = 0; export function getWebSearchLoadError() { return lastLoadError; } +function isCurrentBinding(pi: ExtensionAPI, generation: number) { + return bindGeneration === generation && (!boundLivePi || boundLivePi === pi); +} + function resolvePiCodingAgentRoot(dir: string): string { const packageJson = join(dir, 'package.json'); if ( @@ -174,12 +179,13 @@ function wireRuntimeToLivePi(runtime: Record, pi: ExtensionAPI) /** Remember the live session ExtensionAPI (bound after session_start). */ export function bindLivePiWebAccess(pi: ExtensionAPI) { + bindGeneration += 1; boundLivePi = pi; webSearchExecute = undefined; loadPromise = undefined; } -async function captureWebSearchFromLivePi(pi: ExtensionAPI) { +async function captureWebSearchFromLivePi(pi: ExtensionAPI, generation: number) { const entry = resolvePiWebAccessEntry(); if (!entry) return; @@ -198,32 +204,40 @@ async function captureWebSearchFromLivePi(pi: ExtensionAPI) { const registered = extension.tools.get(PI_WEB_SEARCH_TOOL); if (!registered) { + if (!isCurrentBinding(pi, generation)) return; lastLoadError = 'pi-web-access loaded but did not register web_search. Update pi-web-access.'; return; } + if (!isCurrentBinding(pi, generation)) return; webSearchExecute = registered.definition.execute.bind(registered.definition) as WebSearchExecute; lastLoadError = undefined; } -export async function ensureWebSearchDelegate(pi?: ExtensionAPI) { - if (!isPiWebAccessInstalled()) return; +export async function ensureWebSearchDelegate( + pi?: ExtensionAPI, + isInstalled: () => boolean = isPiWebAccessInstalled, +) { + if (!isInstalled()) return; const livePi = pi ?? boundLivePi; if (!livePi) return; + const generation = bindGeneration; + if (!isCurrentBinding(livePi, generation)) return; if (webSearchExecute) return; if (loadPromise) return loadPromise; loadPromise = (async () => { - lastLoadError = undefined; + if (isCurrentBinding(livePi, generation)) lastLoadError = undefined; try { - await captureWebSearchFromLivePi(livePi); + await captureWebSearchFromLivePi(livePi, generation); } catch (err) { + if (!isCurrentBinding(livePi, generation)) return; lastLoadError = err instanceof Error ? err.message : String(err); webSearchExecute = undefined; } finally { - loadPromise = undefined; + if (isCurrentBinding(livePi, generation)) loadPromise = undefined; } })(); @@ -235,6 +249,7 @@ export function getWebSearchDelegate() { } export function clearWebSearchDelegateForTests() { + bindGeneration += 1; webSearchExecute = undefined; loadPromise = undefined; lastLoadError = undefined; diff --git a/tests/tools/webSearchDelegate.test.ts b/tests/tools/webSearchDelegate.test.ts index 5c30516..faa7ba8 100644 --- a/tests/tools/webSearchDelegate.test.ts +++ b/tests/tools/webSearchDelegate.test.ts @@ -2,7 +2,6 @@ import { mkdirSync, mkdtempSync, writeFileSync } from 'node:fs'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; import { afterEach, describe, expect, it, vi } from 'vitest'; -import * as delegateModule from '../../src/tools/webSearchDelegate.js'; import { bindLivePiWebAccess, clearWebSearchDelegateForTests, @@ -54,13 +53,11 @@ describe('webSearchDelegate', () => { }); it('ensureWebSearchDelegate returns undefined when pi-web-access is not installed', async () => { - vi.spyOn(delegateModule, 'isPiWebAccessInstalled').mockReturnValue(false); - - const result = await ensureWebSearchDelegate(); + const isInstalled = vi.fn(() => false); + const result = await ensureWebSearchDelegate(undefined, isInstalled); + expect(isInstalled).toHaveBeenCalledOnce(); expect(result).toBeUndefined(); expect(getWebSearchDelegate()).toBeUndefined(); - - vi.restoreAllMocks(); }); it('getWebSearchLoadError returns last error string', () => { diff --git a/tests/tools/webSearchDelegateRetry.test.ts b/tests/tools/webSearchDelegateRetry.test.ts index 36accc7..cf17a6f 100644 --- a/tests/tools/webSearchDelegateRetry.test.ts +++ b/tests/tools/webSearchDelegateRetry.test.ts @@ -11,6 +11,7 @@ vi.mock('@earendil-works/pi-coding-agent', async (importOriginal) => ({ })); import { + bindLivePiWebAccess, clearWebSearchDelegateForTests, ensureWebSearchDelegate, getWebSearchDelegate, @@ -50,4 +51,50 @@ export default function (pi) { expect(getWebSearchDelegate()).toBeTypeOf('function'); expect(getWebSearchLoadError()).toBeUndefined(); }); + + it('does not let a stale load replace the delegate for a newer binding', async () => { + const extensionDir = join(testAgentDir, 'npm', 'node_modules', 'pi-web-access'); + mkdirSync(extensionDir, { recursive: true }); + writeFileSync( + join(extensionDir, 'index.js'), + ` +export default async function (pi) { + const load = globalThis.webSearchDelegateLoads.shift() + load.started() + await load.wait + pi.registerTool({ + name: 'web_search', + execute: async () => ({ content: [{ type: 'text', text: load.name }], details: {} }), + }) +} +`, + ); + let startFirstLoad = () => {}; + const firstLoadStarted = new Promise((resolve) => { + startFirstLoad = resolve; + }); + let finishFirstLoad = () => {}; + const firstLoadWait = new Promise((resolve) => { + finishFirstLoad = resolve; + }); + vi.stubGlobal('webSearchDelegateLoads', [ + { name: 'first', started: startFirstLoad, wait: firstLoadWait }, + { name: 'second', started: () => {}, wait: Promise.resolve() }, + ]); + const firstPi = {} as Parameters[0]; + const secondPi = {} as Parameters[0]; + + bindLivePiWebAccess(firstPi); + const firstLoad = ensureWebSearchDelegate(); + await firstLoadStarted; + + bindLivePiWebAccess(secondPi); + await ensureWebSearchDelegate(); + const secondDelegate = getWebSearchDelegate(); + expect(secondDelegate).toBeTypeOf('function'); + + finishFirstLoad(); + await firstLoad; + expect(getWebSearchDelegate()).toBe(secondDelegate); + }); }); From 72419299fc5040388056780a38baedf9890625ba Mon Sep 17 00:00:00 2001 From: kenryu42 Date: Wed, 3 Jun 2026 17:23:51 +0900 Subject: [PATCH 9/9] fix(shell): include termination signal in timeout error details --- src/tools/shell.ts | 2 ++ tests/tools/shell.test.ts | 7 ++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/tools/shell.ts b/src/tools/shell.ts index 5078826..75d6cb3 100644 --- a/src/tools/shell.ts +++ b/src/tools/shell.ts @@ -109,6 +109,7 @@ export function registerShellTool(pi: ExtensionAPI) { const err = error as { code?: unknown; message?: string; + signal?: unknown; stdout?: string; stderr?: string; }; @@ -132,6 +133,7 @@ export function registerShellTool(pi: ExtensionAPI) { details: { exitCode, command: params.command, + ...(typeof err.signal === 'string' ? { signal: err.signal } : {}), }, }; } diff --git a/tests/tools/shell.test.ts b/tests/tools/shell.test.ts index 209bbd1..0bfcb89 100644 --- a/tests/tools/shell.test.ts +++ b/tests/tools/shell.test.ts @@ -104,15 +104,16 @@ describe('shell tool', () => { it('kills commands that exceed the timeout', async () => { const cwd = tempDir('pi-grok-cli-shell-'); + const command = 'node -e "setTimeout(()=>{},10000)"'; const result = await executeTool( collectTools(registerShellTool).get('Shell'), - { command: 'sleep 10', timeout: 100 }, + { command, timeout: 100 }, cwd, ); expect(firstText(result)).toContain('Shell error'); - expect(result.details).toMatchObject({ command: 'sleep 10' }); - expect(result.details.exitCode).toBeDefined(); + expect(result.details.command).toBe(command); + expect(result.details.signal).toMatch(/TERM|KILL/); }); it('renders shell calls and result states', () => {