diff --git a/src/commands/agents/test.ts b/src/commands/agents/test.ts index 225a942..7b3dd07 100644 --- a/src/commands/agents/test.ts +++ b/src/commands/agents/test.ts @@ -16,15 +16,18 @@ import { defineCommand } from "citty"; import { apiStreamRequest, buildInteractionRequest, + DEFAULT_SERVER_TIMEOUT_SECONDS, isAgentName, normalizeSources, type RunOptions, resolveContext, + SharedFlags, type Source, validateSources, } from "../../lib/api"; import { loadAgent } from "../../lib/config"; import { CLIError, ConfigError } from "../../lib/errors"; +import { TestArgsSchema } from "../../lib/schemas"; import { collectInlineFiles } from "../../lib/files"; import { logRequest, logResponse } from "../../lib/logger"; import { @@ -67,6 +70,10 @@ Examples: type: "string", description: "Use existing environment", }, + timeout: { + type: "string", + description: "Override timeout in seconds", + }, env: { type: "string", alias: "e", @@ -75,17 +82,26 @@ Examples: }, async run({ args }) { try { - const agentDir = args.path as string; - const prompt = args.prompt as string; - const envFile = args.env as string | undefined; + const parseResult = TestArgsSchema.safeParse(args); + if (!parseResult.success) { + const errors = parseResult.error.issues + .map((i) => ` - ${i.path.join(".")}: ${i.message}`) + .join("\n"); + throw new CLIError(`Invalid arguments:\n${errors}`); + } + const parsedArgs = parseResult.data; + + const agentDir = parsedArgs.path; + const prompt = parsedArgs.prompt; + const envFile = parsedArgs.env; const sharedFlags = { - apiKey: (args["api-key"] || args.apiKey) as string | undefined, - baseUrl: (args["base-url"] || args.baseUrl) as string | undefined, - json: args.json as boolean, - dryRun: (args["dry-run"] || args.dryRun) as boolean, + apiKey: parsedArgs["api-key"], + baseUrl: parsedArgs["base-url"], + json: parsedArgs.json, + dryRun: parsedArgs["dry-run"], }; - const ctx = resolveContext(sharedFlags); + const ctx = resolveContext(sharedFlags as SharedFlags); const { config } = await loadAgent(agentDir, { envFile }); const inlineFiles = await collectInlineFiles(agentDir); @@ -97,8 +113,8 @@ Examples: // Build environment config let environment: any; - if (args.environment) { - environment = args.environment; + if (parsedArgs.environment) { + environment = parsedArgs.environment; } else { const sources: Source[] = [...inlineFiles] as unknown as Source[]; if (config.sources) { @@ -145,23 +161,28 @@ Examples: input: prompt, systemInstruction: systemInstruction, tools: config.tools as any, - previousInteractionId: args["previous-interaction-id"] as string | undefined, + previousInteractionId: parsedArgs["previous-interaction-id"], stream: true, environment: environment || undefined, }; const body = buildInteractionRequest(runOpts); - if (args["dry-run"]) { - printCurl("POST", `${ctx.baseUrl}/interactions`, ctx.apiKey, body); + const timeoutSeconds = parsedArgs.timeout ?? config.timeout ?? DEFAULT_SERVER_TIMEOUT_SECONDS; + const headers = { + "x-server-timeout": timeoutSeconds.toString(), + }; + + if (parsedArgs["dry-run"]) { + printCurl("POST", `${ctx.baseUrl}/interactions`, ctx.apiKey, body, headers); return; } const startTime = performance.now(); - const response = await apiStreamRequest(ctx, "/interactions", body); + const response = await apiStreamRequest(ctx, "/interactions", body, headers); - if (args.json) { + if (parsedArgs.json) { await processStream(response, { onEvent: (event) => { console.log(event.raw); @@ -169,7 +190,7 @@ Examples: onComplete: () => {}, }); } else { - const verbose = args.verbose as boolean; + const verbose = parsedArgs.verbose; const renderer = new HumanStreamRenderer(process.stdout, verbose); await processStream(response, { @@ -181,7 +202,7 @@ Examples: renderer.finish(); const latencySeconds = (performance.now() - startTime) / 1000; printCompletionSummary(result, latencySeconds, verbose); - if (!args.json) { + if (!parsedArgs.json) { logRequest(result.interactionId, body); logResponse(result.interactionId, result); } diff --git a/src/lib/api.ts b/src/lib/api.ts index a83b4b7..0876e51 100644 --- a/src/lib/api.ts +++ b/src/lib/api.ts @@ -15,6 +15,7 @@ import { CLIError } from "./errors"; export const DEFAULT_BASE_URL = "https://generativelanguage.googleapis.com/v1beta"; +export const DEFAULT_SERVER_TIMEOUT_SECONDS = 300000; export async function fetchWithTimeout( url: string, @@ -76,7 +77,7 @@ export async function apiRequest( const headers: Record = { "Content-Type": "application/json", "x-goog-api-key": ctx.apiKey, - "x-server-timeout": "30000", + "x-server-timeout": DEFAULT_SERVER_TIMEOUT_SECONDS.toString(), }; if (path.includes("/interactions")) { @@ -85,7 +86,7 @@ export async function apiRequest( const response = await fetchWithTimeout(url, { method, - headers, + headers: headers, body: body ? JSON.stringify(body) : undefined, }); @@ -187,21 +188,23 @@ export async function apiStreamRequest( ctx: CLIContext, path: string, body: unknown, + headerOverrides?: Record, ): Promise { const url = `${ctx.baseUrl}${path}`; - const headers: Record = { + const mergedHeaders: Record = { "Content-Type": "application/json", "x-goog-api-key": ctx.apiKey, - "x-server-timeout": "30000", + "x-server-timeout": DEFAULT_SERVER_TIMEOUT_SECONDS.toString(), + ...headerOverrides, }; if (path.includes("/interactions")) { - headers["Api-Revision"] = "2026-05-20"; + mergedHeaders["Api-Revision"] = "2026-05-20"; } const response = await fetchWithTimeout(url, { method: "POST", - headers, + headers: mergedHeaders, body: JSON.stringify(body), }); diff --git a/src/lib/output.ts b/src/lib/output.ts index bc81337..53ef9ca 100644 --- a/src/lib/output.ts +++ b/src/lib/output.ts @@ -12,17 +12,25 @@ // See the License for the specific language governing permissions and // limitations under the License. +import { DEFAULT_SERVER_TIMEOUT_SECONDS } from "./api"; import type { ContentBlock, StreamEvent, StreamResult } from "./stream"; -export function printCurl(method: string, url: string, apiKey: string, body?: unknown): void { - let curl = `curl -X ${method} "${url}" \\\n`; - curl += ` -H "Content-Type: application/json" \\\n`; - curl += ` -H "x-goog-api-key: ${apiKey}" \\\n`; - curl += ` -H "x-server-timeout: 30000"`; - - if (url.includes("/interactions")) { - curl += ` \\\n -H "Api-Revision: 2026-05-20"`; - } +export function printCurl( + method: string, + url: string, + apiKey: string, + body?: unknown, + headers?: Record, +): void { + const serverTimeout = headers?.["x-server-timeout"] ?? DEFAULT_SERVER_TIMEOUT_SECONDS.toString(); + const apiRevision = url.includes("/interactions") + ? ` \\\n -H "Api-Revision: 2026-05-20"` + : ""; + + let curl = `curl -X ${method} "${url}" \\ + -H "Content-Type: application/json" \\ + -H "x-goog-api-key: ${apiKey}" \\ + -H "x-server-timeout: ${serverTimeout}"${apiRevision}`; if (body) { // Escape single quotes in body for bash diff --git a/src/lib/schemas.ts b/src/lib/schemas.ts index 2172406..9b116c2 100644 --- a/src/lib/schemas.ts +++ b/src/lib/schemas.ts @@ -71,6 +71,8 @@ const ExampleSchema = z.object({ prompt: z.string(), }); +const TimeoutSecondsSchema = z.coerce.number().int().positive(); + export const AgentConfigSchema = z .object({ id: z.string(), @@ -82,6 +84,7 @@ export const AgentConfigSchema = z sources: z.array(SourceSchema).optional(), environment: EnvironmentSchema.optional(), examples: z.array(ExampleSchema).optional(), + timeout: TimeoutSecondsSchema.optional(), }) .strict(); @@ -93,3 +96,17 @@ export const CLIContextSchema = z.object({ }); export type CLIContext = z.infer; + +export const TestArgsSchema = z.object({ + prompt: z.string(), + path: z.string().default("."), + "previous-interaction-id": z.string().optional(), + environment: z.string().optional(), + timeout: TimeoutSecondsSchema.optional(), + env: z.string().optional(), + "api-key": z.string().optional(), + "base-url": z.string().optional(), + json: z.boolean().default(false), + "dry-run": z.boolean().default(false), + verbose: z.boolean().default(false), +});