feat(test): add timeout support for agents test command

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,31 +161,36 @@ 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);
           },
           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);
             }

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<T>(
   const headers: Record<string, string> = {
     "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<T>(
 
   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<string, string>,
 ): Promise<Response> {
   const url = `${ctx.baseUrl}${path}`;
-  const headers: Record<string, string> = {
+  const mergedHeaders: Record<string, string> = {
     "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),
   });
 

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<string, string>,
+): 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

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<typeof CLIContextSchema>;
+
+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),
+});