From d512cf823e245eb4c5d1e7087e60584c49902753 Mon Sep 17 00:00:00 2001 From: Ivan Leo Date: Fri, 5 Jun 2026 12:34:32 -0700 Subject: [PATCH 1/4] feat(test): add timeout support for agents test command --- src/commands/agents/test.ts | 52 +++++++++++++++++++++++++------------ src/lib/api.ts | 10 ++++--- src/lib/output.ts | 25 +++++++++++------- src/lib/schemas.ts | 14 ++++++++++ 4 files changed, 72 insertions(+), 29 deletions(-) diff --git a/src/commands/agents/test.ts b/src/commands/agents/test.ts index 9cd5549..7247e93 100644 --- a/src/commands/agents/test.ts +++ b/src/commands/agents/test.ts @@ -20,11 +20,13 @@ import { 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,20 +69,33 @@ Examples: type: "string", description: "Use existing environment", }, + timeout: { + type: "string", + description: "Override timeout in seconds", + }, }, async run({ args }) { try { - const agentDir = args.path as string; - const prompt = args.prompt as string; + 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 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); const inlineFiles = await collectInlineFiles(agentDir); @@ -92,8 +107,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) { @@ -140,23 +155,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 ?? 600; + 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); @@ -164,7 +184,7 @@ Examples: onComplete: () => {}, }); } else { - const verbose = args.verbose as boolean; + const verbose = parsedArgs.verbose; const renderer = new HumanStreamRenderer(process.stdout, verbose); await processStream(response, { @@ -176,7 +196,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..3ca527b 100644 --- a/src/lib/api.ts +++ b/src/lib/api.ts @@ -85,7 +85,7 @@ export async function apiRequest( const response = await fetchWithTimeout(url, { method, - headers, + headers: headers, body: body ? JSON.stringify(body) : undefined, }); @@ -187,21 +187,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", + ...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..1e41117 100644 --- a/src/lib/output.ts +++ b/src/lib/output.ts @@ -14,15 +14,22 @@ 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"] ?? "30000"; + 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..b7e711d 100644 --- a/src/lib/schemas.ts +++ b/src/lib/schemas.ts @@ -82,6 +82,7 @@ export const AgentConfigSchema = z sources: z.array(SourceSchema).optional(), environment: EnvironmentSchema.optional(), examples: z.array(ExampleSchema).optional(), + timeout: z.number().optional(), }) .strict(); @@ -93,3 +94,16 @@ 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: z.coerce.number().int().positive().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), +}); From b493f0f38f7350795b3a080c16eef1bc6a6eb4ca Mon Sep 17 00:00:00 2001 From: Ivan Leo Date: Fri, 5 Jun 2026 12:50:05 -0700 Subject: [PATCH 2/4] fix(test): validate agent timeout config --- src/commands/agents/test.ts | 2 +- src/lib/schemas.ts | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/commands/agents/test.ts b/src/commands/agents/test.ts index 7247e93..72558bd 100644 --- a/src/commands/agents/test.ts +++ b/src/commands/agents/test.ts @@ -162,7 +162,7 @@ Examples: const body = buildInteractionRequest(runOpts); - const timeoutSeconds = parsedArgs.timeout ?? config.timeout ?? 600; + const timeoutSeconds = parsedArgs.timeout ?? config.timeout ?? 300000; const headers = { "x-server-timeout": timeoutSeconds.toString(), }; diff --git a/src/lib/schemas.ts b/src/lib/schemas.ts index b7e711d..af7d3a4 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,7 +84,7 @@ export const AgentConfigSchema = z sources: z.array(SourceSchema).optional(), environment: EnvironmentSchema.optional(), examples: z.array(ExampleSchema).optional(), - timeout: z.number().optional(), + timeout: TimeoutSecondsSchema.optional(), }) .strict(); @@ -100,7 +102,7 @@ export const TestArgsSchema = z.object({ path: z.string().default("."), "previous-interaction-id": z.string().optional(), environment: z.string().optional(), - timeout: z.coerce.number().int().positive().optional(), + timeout: TimeoutSecondsSchema.optional(), "api-key": z.string().optional(), "base-url": z.string().optional(), json: z.boolean().default(false), From 27bd50a865810c9a8b36da28aaf610a85ab58692 Mon Sep 17 00:00:00 2001 From: Ivan Leo Date: Fri, 5 Jun 2026 12:51:36 -0700 Subject: [PATCH 3/4] fix(api): centralize server timeout default --- src/commands/agents/test.ts | 3 ++- src/lib/api.ts | 5 +++-- src/lib/output.ts | 3 ++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/src/commands/agents/test.ts b/src/commands/agents/test.ts index 72558bd..a1e9b14 100644 --- a/src/commands/agents/test.ts +++ b/src/commands/agents/test.ts @@ -16,6 +16,7 @@ import { defineCommand } from "citty"; import { apiStreamRequest, buildInteractionRequest, + DEFAULT_SERVER_TIMEOUT_SECONDS, isAgentName, normalizeSources, type RunOptions, @@ -162,7 +163,7 @@ Examples: const body = buildInteractionRequest(runOpts); - const timeoutSeconds = parsedArgs.timeout ?? config.timeout ?? 300000; + const timeoutSeconds = parsedArgs.timeout ?? config.timeout ?? DEFAULT_SERVER_TIMEOUT_SECONDS; const headers = { "x-server-timeout": timeoutSeconds.toString(), }; diff --git a/src/lib/api.ts b/src/lib/api.ts index 3ca527b..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")) { @@ -193,7 +194,7 @@ export async function apiStreamRequest( 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, }; diff --git a/src/lib/output.ts b/src/lib/output.ts index 1e41117..53ef9ca 100644 --- a/src/lib/output.ts +++ b/src/lib/output.ts @@ -12,6 +12,7 @@ // 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( @@ -21,7 +22,7 @@ export function printCurl( body?: unknown, headers?: Record, ): void { - const serverTimeout = headers?.["x-server-timeout"] ?? "30000"; + const serverTimeout = headers?.["x-server-timeout"] ?? DEFAULT_SERVER_TIMEOUT_SECONDS.toString(); const apiRevision = url.includes("/interactions") ? ` \\\n -H "Api-Revision: 2026-05-20"` : ""; From fbb94026221538adeed7e013301278b6e36467eb Mon Sep 17 00:00:00 2001 From: Ivan Leo Date: Sun, 7 Jun 2026 13:37:48 -0700 Subject: [PATCH 4/4] fix(test): resolve timeout env flag merge --- src/commands/agents/test.ts | 6 ++---- src/lib/schemas.ts | 1 + 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/commands/agents/test.ts b/src/commands/agents/test.ts index c3af35d..7b3dd07 100644 --- a/src/commands/agents/test.ts +++ b/src/commands/agents/test.ts @@ -73,6 +73,7 @@ Examples: timeout: { type: "string", description: "Override timeout in seconds", + }, env: { type: "string", alias: "e", @@ -92,10 +93,7 @@ Examples: const agentDir = parsedArgs.path; const prompt = parsedArgs.prompt; - - const agentDir = args.path as string; - const prompt = args.prompt as string; - const envFile = args.env as string | undefined; + const envFile = parsedArgs.env; const sharedFlags = { apiKey: parsedArgs["api-key"], baseUrl: parsedArgs["base-url"], diff --git a/src/lib/schemas.ts b/src/lib/schemas.ts index af7d3a4..9b116c2 100644 --- a/src/lib/schemas.ts +++ b/src/lib/schemas.ts @@ -103,6 +103,7 @@ export const TestArgsSchema = z.object({ "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),