|
const DEFAULT_MODEL = "gemini-3.1-pro-preview"; |
|
const MODEL_FALLBACK_CHAIN = [ |
|
"gemini-3.1-pro-preview", |
|
"gemini-3-flash-preview", |
|
"gemini-3.1-flash-lite", |
|
"gemini-2.5-pro", |
|
"gemini-2.5-flash", |
|
"gemini-2.5-flash-lite", |
|
]; |
|
const DEFAULT_TIMEOUT_MS = 180_000; |
|
const DEFAULT_MAX_CHARS = 80_000; |
|
|
|
type GlobalWithProcess = typeof globalThis & { |
|
process?: { |
|
env?: { |
|
HOME?: string; |
|
}; |
|
}; |
|
}; |
|
|
|
const HOME_DIR = (globalThis as GlobalWithProcess).process?.env?.HOME ?? ""; |
|
const SEARCH_ONLY_POLICY = `${HOME_DIR}/.gemini/policies/websearch-only.toml`; |
|
|
|
const websearchParams = { |
|
type: "object", |
|
properties: { |
|
query: { |
|
type: "string", |
|
description: "Self-contained web search or research question. Do not include secrets, private URLs, credentials, or unreleased private details.", |
|
}, |
|
mode: { |
|
type: "string", |
|
enum: ["search", "research"], |
|
description: "Use search for a concise answer; research for deeper source-backed notes. Defaults to research.", |
|
}, |
|
model: { |
|
type: "string", |
|
description: "Preferred Gemini CLI model. Defaults to gemini-3.1-pro-preview. The alias gemini-3-pro is mapped to gemini-3.1-pro-preview; capacity/rate-limit failures devolve through gemini-3-flash-preview, gemini-3.1-flash-lite, gemini-2.5-pro, gemini-2.5-flash, and gemini-2.5-flash-lite.", |
|
}, |
|
timeoutMs: { |
|
type: "number", |
|
description: `Timeout in milliseconds. Defaults to ${DEFAULT_TIMEOUT_MS}.`, |
|
}, |
|
maxChars: { |
|
type: "number", |
|
description: `Maximum characters returned to the model. Defaults to ${DEFAULT_MAX_CHARS}.`, |
|
}, |
|
}, |
|
required: ["query"], |
|
additionalProperties: false, |
|
} as const; |
|
|
|
type WebsearchParams = { |
|
query: string; |
|
mode?: "search" | "research"; |
|
model?: string; |
|
timeoutMs?: number; |
|
maxChars?: number; |
|
}; |
|
|
|
function normalizeModel(model?: string): string { |
|
const trimmed = model?.trim(); |
|
if (!trimmed || trimmed === "gemini-3-pro") return DEFAULT_MODEL; |
|
return trimmed; |
|
} |
|
|
|
export function buildModelChain(model?: string): string[] { |
|
const preferredModel = normalizeModel(model); |
|
const preferredIndex = MODEL_FALLBACK_CHAIN.indexOf(preferredModel); |
|
const lowerModels = preferredIndex >= 0 |
|
? MODEL_FALLBACK_CHAIN.slice(preferredIndex + 1) |
|
: MODEL_FALLBACK_CHAIN; |
|
return [preferredModel, ...lowerModels].filter((candidate, index, allModels) => allModels.indexOf(candidate) === index); |
|
} |
|
|
|
export function isFallbackableGeminiFailure(stdout: string, stderr: string, model: string): boolean { |
|
const combined = `${stdout}\n${stderr}`; |
|
if (/ModelNotFoundError/i.test(combined)) return false; |
|
if (/MODEL_CAPACITY_EXHAUSTED/i.test(combined)) return true; |
|
if (/No capacity available for model/i.test(combined) && combined.includes(model)) return true; |
|
if (/RESOURCE_EXHAUSTED/i.test(combined) && /\b429\b/.test(combined)) return true; |
|
return false; |
|
} |
|
|
|
function buildPrompt(query: string, mode: "search" | "research"): string { |
|
const depth = mode === "search" |
|
? "Answer concisely. Use 2-4 key findings unless more are necessary." |
|
: "Do deeper web research. Compare sources, note conflicts, and include enough detail for another agent to use the result."; |
|
|
|
return `You are a web research subagent being called by the local Pi coding agent. |
|
|
|
Use the google_web_search tool when the answer depends on current, niche, source-backed, or uncertain public information. |
|
|
|
Do not use local filesystem tools, shell tools, code execution, memory tools, MCP tools, web_fetch, or file editing tools. |
|
|
|
${depth} |
|
|
|
Return Markdown only, using this structure: |
|
|
|
## Direct answer |
|
A concise answer to the research question. |
|
|
|
## Key findings |
|
- Important finding with source URL. |
|
- Important finding with source URL. |
|
|
|
## Sources |
|
- Source title -- URL -- what it supports. |
|
|
|
## Caveats |
|
- Note uncertainty, stale data risk, or conflicts between sources. |
|
|
|
Research question: |
|
${query}`; |
|
} |
|
|
|
function extractGeminiResponse(stdout: string): string { |
|
const text = stdout.trim(); |
|
if (!text) return ""; |
|
try { |
|
const parsed = JSON.parse(text); |
|
if (typeof parsed === "string") return parsed; |
|
if (parsed && typeof parsed === "object") { |
|
const object = parsed as Record<string, unknown>; |
|
for (const key of ["response", "text", "output", "content", "message"]) { |
|
if (typeof object[key] === "string") return object[key] as string; |
|
} |
|
} |
|
} catch { |
|
// Fall back to raw stdout below. |
|
} |
|
return text; |
|
} |
|
|
|
function truncate(text: string, maxChars: number): { text: string; truncated: boolean } { |
|
if (text.length <= maxChars) return { text, truncated: false }; |
|
return { |
|
text: `${text.slice(0, maxChars)}\n\n[gemini_websearch output truncated at ${maxChars} characters]`, |
|
truncated: true, |
|
}; |
|
} |
|
|
|
async function fileExists(pi: any, path: string, signal?: AbortSignal): Promise<boolean> { |
|
try { |
|
const result = await pi.exec("test", ["-f", path], { |
|
signal, |
|
timeout: 1_000, |
|
}); |
|
return result.code === 0; |
|
} catch { |
|
return false; |
|
} |
|
} |
|
|
|
export default function geminiWebsearchExtension(pi: any) { |
|
pi.registerTool({ |
|
name: "gemini_websearch", |
|
label: "Gemini Websearch", |
|
description: |
|
"Search or research the public web with the logged-in Gemini CLI using Gemini 3.1 Pro with Gemini 3/2.5 fallback and a search-only policy. Returns Markdown with source URLs.", |
|
promptSnippet: |
|
"Search or research current public web information through Gemini CLI (Gemini 3.1 Pro with Gemini 3/2.5 fallback) with source URLs.", |
|
promptGuidelines: [ |
|
"Use gemini_websearch when current, niche, source-backed, or post-cutoff public information is needed and a normal page extraction tool is insufficient.", |
|
"Do not pass secrets, credentials, private URLs, local file paths, or unreleased private project details to gemini_websearch.", |
|
"Treat gemini_websearch output as research notes, preserve source URLs when relevant, and verify important claims with narrower follow-up searches when needed.", |
|
], |
|
parameters: websearchParams, |
|
async execute(_toolCallId: string, params: WebsearchParams, signal: AbortSignal | undefined, onUpdate: ((value: unknown) => void) | undefined, ctx: any) { |
|
const query = params.query.trim(); |
|
if (!query) { |
|
return { |
|
content: [{ type: "text", text: "gemini_websearch requires a non-empty query." }], |
|
details: { ok: false, reason: "empty_query" }, |
|
isError: true, |
|
}; |
|
} |
|
|
|
const mode = params.mode ?? "research"; |
|
const modelChain = buildModelChain(params.model); |
|
const timeout = Math.max(10_000, Math.floor(params.timeoutMs ?? DEFAULT_TIMEOUT_MS)); |
|
const maxChars = Math.max(1_000, Math.floor(params.maxChars ?? DEFAULT_MAX_CHARS)); |
|
const prompt = buildPrompt(query, mode); |
|
const hasSearchOnlyPolicy = await fileExists(pi, SEARCH_ONLY_POLICY, signal); |
|
if (!hasSearchOnlyPolicy) { |
|
return { |
|
content: [{ type: "text", text: `Gemini websearch policy file not found at ${SEARCH_ONLY_POLICY}. Refusing to run Gemini CLI with --approval-mode=yolo without a search-only policy.` }], |
|
details: { ok: false, query, mode, policy: SEARCH_ONLY_POLICY, reason: "missing_search_only_policy" }, |
|
isError: true, |
|
}; |
|
} |
|
const attemptedModels: Array<{ model: string; code: number; fallbackable: boolean }> = []; |
|
|
|
try { |
|
for (const [index, model] of modelChain.entries()) { |
|
const args = [ |
|
"--prompt", prompt, |
|
"--skip-trust", |
|
"--approval-mode=yolo", |
|
"--output-format", "json", |
|
"--model", model, |
|
]; |
|
|
|
if (hasSearchOnlyPolicy) { |
|
args.push("--policy", SEARCH_ONLY_POLICY); |
|
} |
|
|
|
onUpdate?.({ content: [{ type: "text", text: `Searching the web with Gemini CLI (${model})...` }] }); |
|
|
|
const result = await pi.exec("gemini", args, { |
|
cwd: ctx.cwd, |
|
signal, |
|
timeout, |
|
}); |
|
const fallbackable = result.code !== 0 && index < modelChain.length - 1 && isFallbackableGeminiFailure(result.stdout, result.stderr, model); |
|
attemptedModels.push({ model, code: result.code, fallbackable }); |
|
|
|
if (fallbackable) { |
|
onUpdate?.({ content: [{ type: "text", text: `Gemini CLI capacity/rate-limit failure on ${model}; falling back to ${modelChain[index + 1]}...` }] }); |
|
continue; |
|
} |
|
|
|
const response = extractGeminiResponse(result.stdout); |
|
const rawText = response || result.stderr.trim() || "Gemini CLI returned no output."; |
|
const { text, truncated } = truncate(rawText, maxChars); |
|
const ok = result.code === 0; |
|
|
|
return { |
|
content: [{ type: "text", text }], |
|
details: { |
|
ok, |
|
code: result.code, |
|
query, |
|
mode, |
|
model, |
|
requestedModel: normalizeModel(params.model), |
|
modelChain, |
|
attemptedModels, |
|
fallbackUsed: index > 0, |
|
policy: hasSearchOnlyPolicy ? SEARCH_ONLY_POLICY : undefined, |
|
truncated, |
|
stdoutLength: result.stdout.length, |
|
stderr: result.stderr, |
|
}, |
|
isError: !ok, |
|
}; |
|
} |
|
|
|
return { |
|
content: [{ type: "text", text: "Gemini CLI returned no output." }], |
|
details: { ok: false, query, mode, modelChain, attemptedModels }, |
|
isError: true, |
|
}; |
|
} catch (error) { |
|
const message = error instanceof Error ? error.message : String(error); |
|
const installHint = message.includes("ENOENT") || message.includes("not found") |
|
? "\n\nInstall and log in to Gemini CLI first: npm install -g @google/gemini-cli && gemini" |
|
: ""; |
|
return { |
|
content: [{ type: "text", text: `Gemini websearch failed: ${message}${installHint}` }], |
|
details: { ok: false, query, mode, modelChain, attemptedModels, error: message }, |
|
isError: true, |
|
}; |
|
} |
|
}, |
|
}); |
|
|
|
pi.registerCommand("websearch", { |
|
description: "Ask Gemini CLI to search/research the web: /websearch <query>", |
|
handler: async (args: string, ctx: any) => { |
|
const query = args.trim(); |
|
if (!query) { |
|
ctx.ui.notify("Usage: /websearch <query>", "warning"); |
|
return; |
|
} |
|
const message = `Use gemini_websearch in research mode for this query: ${query}`; |
|
if (ctx.isIdle()) { |
|
pi.sendUserMessage(message); |
|
} else { |
|
pi.sendUserMessage(message, { deliverAs: "followUp" }); |
|
} |
|
}, |
|
}); |
|
} |
