diff --git a/AGENTS.md b/AGENTS.md index b69ed30..236be48 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,23 +8,23 @@ Grok CLI (`@vibe-kit/grok-cli`) is a single-package TypeScript CLI tool — no d ### Quick reference -| Action | Command | -|--------|---------| -| Install deps | `bun install` (installs Husky; pre-commit runs Biome on staged files) | -| Typecheck | `bun run typecheck` | -| Build | `bun run build` | -| Run built CLI | `node dist/index.js` | -| Headless mode | `node dist/index.js --prompt "..." --max-tool-rounds N` | -| CLI help | `node dist/index.js --help` | + +| Action | Command | +| ------------- | --------------------------------------------------------------------- | +| Install deps | `bun install` (installs Husky; pre-commit runs Biome on staged files) | +| Typecheck | `bun run typecheck` | +| Build | `bun run build` | +| Run built CLI | `node dist/index.js` | +| Headless mode | `node dist/index.js --prompt "..." --max-tool-rounds N` | +| CLI help | `node dist/index.js --help` | + ### Known issues - **ESLint config is broken**: The repo has `.eslintrc.js` (legacy format) but uses ESLint 9 (`^9.31.0`) + `@typescript-eslint` v8, which require flat config (`eslint.config.js`). Additionally, `.eslintrc.js` uses `module.exports` (CJS) but `package.json` has `"type": "module"` (ESM). Running `bun run lint` will fail. Use `bun run typecheck` as the primary code quality check (this is also what CI enforces). - **Dev mode (`bun run dev` / `bun run dev:node`) fails at runtime**: `src/utils/model-config.ts` imports TypeScript interfaces (`UserSettings`, `ProjectSettings`) as value imports from `settings-manager.ts`. These type-only exports are erased at runtime by Bun and tsx, causing `SyntaxError: export '...' not found`. The fix is to use `import type` syntax, but this is a pre-existing repo issue. **Workaround**: build first (`bun run build`), then run the compiled version (`node dist/index.js`). -- **Grok API 410 error**: The client sends `search_parameters` in every API request, but xAI has deprecated "Live search". This causes a 410 error on headless prompts. This is an API compatibility issue in the codebase, not an environment problem. ### Environment - **Bun** must be installed (not pre-installed on Cloud VMs). The update script handles this. -- `GROK_API_KEY` environment variable is required for API calls. Set it as a secret. -- `MORPH_API_KEY` is optional (enables Morph Fast Apply editing). +- `GROK_API_KEY` environment variable is required for API calls. Set it as a secret. \ No newline at end of file diff --git a/README.md b/README.md index 9f59ba8..8a099f2 100644 --- a/README.md +++ b/README.md @@ -60,9 +60,14 @@ grok --prompt "run the test suite and summarize failures" grok -p "show me package.json" --directory /path/to/project grok --prompt "refactor X" --max-tool-rounds 30 grok --prompt "summarize the repo state" --format json +grok --prompt "review the repo overnight" --batch-api grok --verify ``` +`--batch-api` uses xAI's Batch API for lower-cost unattended runs. It is a good +fit for scripts, CI, schedules, and other non-interactive workflows where a +delayed result is fine. + **Continue a saved session:** ```bash diff --git a/bun.lock b/bun.lock index 5d8ec4d..6a15a62 100644 --- a/bun.lock +++ b/bun.lock @@ -6,6 +6,7 @@ "name": "@vibe-kit/grok-cli", "dependencies": { "@ai-sdk/mcp": "^1.0.25", + "@ai-sdk/provider-utils": "^4.0.21", "@ai-sdk/xai": "^3.0.67", "@modelcontextprotocol/sdk": "^1.27.1", "@opentui/core": "^0.1.88", @@ -41,7 +42,7 @@ "@ai-sdk/provider": ["@ai-sdk/provider@3.0.8", "", { "dependencies": { "json-schema": "^0.4.0" } }, "sha512-oGMAgGoQdBXbZqNG0Ze56CHjDZ1IDYOwGYxYjO5KLSlz5HiNQ9udIXsPZ61VWaHGZ5XW/jyjmr6t2xz2jGVwbQ=="], - "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.19", "", { "dependencies": { "@ai-sdk/provider": "3.0.8", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-3eG55CrSWCu2SXlqq2QCsFjo3+E7+Gmg7i/oRVoSZzIodTuDSfLb3MRje67xE9RFea73Zao7Lm4mADIfUETKGg=="], + "@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.21", "", { "dependencies": { "@ai-sdk/provider": "3.0.8", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-MtFUYI1/8mgDvRmaBDjbLJPFFrMG777AvSgyIFQtZHIMzm88R/12vYBBpnk7pfiWLFE1DSZzY4WDYzGbKAcmiw=="], "@ai-sdk/xai": ["@ai-sdk/xai@3.0.67", "", { "dependencies": { "@ai-sdk/openai-compatible": "2.0.35", "@ai-sdk/provider": "3.0.8", "@ai-sdk/provider-utils": "4.0.19" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-KQQIDc91dUA5IGFMnXBuvPBeraYNTdpDC1qUS+JG8vE+/299//5sZFafI1kKYUu3f3p7LaZrKXYgZ1Ni7QIRbw=="], @@ -679,6 +680,14 @@ "zod-to-json-schema": ["zod-to-json-schema@3.25.1", "", { "peerDependencies": { "zod": "^3.25 || ^4" } }, "sha512-pM/SU9d3YAggzi6MtR4h7ruuQlqKtad8e9S0fmxcMi+ueAK5Korys/aWcV9LIIHTVbj01NdzxcnXSN+O74ZIVA=="], + "@ai-sdk/gateway/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.19", "", { "dependencies": { "@ai-sdk/provider": "3.0.8", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-3eG55CrSWCu2SXlqq2QCsFjo3+E7+Gmg7i/oRVoSZzIodTuDSfLb3MRje67xE9RFea73Zao7Lm4mADIfUETKGg=="], + + "@ai-sdk/mcp/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.19", "", { "dependencies": { "@ai-sdk/provider": "3.0.8", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-3eG55CrSWCu2SXlqq2QCsFjo3+E7+Gmg7i/oRVoSZzIodTuDSfLb3MRje67xE9RFea73Zao7Lm4mADIfUETKGg=="], + + "@ai-sdk/openai-compatible/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.19", "", { "dependencies": { "@ai-sdk/provider": "3.0.8", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-3eG55CrSWCu2SXlqq2QCsFjo3+E7+Gmg7i/oRVoSZzIodTuDSfLb3MRje67xE9RFea73Zao7Lm4mADIfUETKGg=="], + + "@ai-sdk/xai/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.19", "", { "dependencies": { "@ai-sdk/provider": "3.0.8", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-3eG55CrSWCu2SXlqq2QCsFjo3+E7+Gmg7i/oRVoSZzIodTuDSfLb3MRje67xE9RFea73Zao7Lm4mADIfUETKGg=="], + "@jimp/plugin-blit/zod": ["zod@3.25.76", "", {}, "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ=="], "@jimp/plugin-circle/zod": ["zod@3.25.76", "", {}, "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ=="], @@ -713,6 +722,8 @@ "@opentui/core/diff": ["diff@8.0.2", "", {}, "sha512-sSuxWU5j5SR9QQji/o2qMvqRNYRDOcBTgsJ/DeCf4iSN4gW+gNMXM7wFIP+fdXZxoNiAnHUTGjCr+TSWXdRDKg=="], + "ai/@ai-sdk/provider-utils": ["@ai-sdk/provider-utils@4.0.19", "", { "dependencies": { "@ai-sdk/provider": "3.0.8", "@standard-schema/spec": "^1.1.0", "eventsource-parser": "^3.0.6" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-3eG55CrSWCu2SXlqq2QCsFjo3+E7+Gmg7i/oRVoSZzIodTuDSfLb3MRje67xE9RFea73Zao7Lm4mADIfUETKGg=="], + "image-q/@types/node": ["@types/node@16.9.1", "", {}, "sha512-QpLcX9ZSsq3YYUUnD3nFDY8H7wctAhQj/TFKL8Ya8v5fMm3CFXxo8zStsLAl780ltoYoo1WvKUVGBQK+1ifr7g=="], "lint-staged/commander": ["commander@14.0.3", "", {}, "sha512-H+y0Jo/T1RZ9qPP4Eh1pkcQcLRglraJaSLoyOtHxu6AapkjWVCy2Sit1QQ4x3Dng8qDlSsZEet7g5Pq06MvTgw=="], diff --git a/package.json b/package.json index 50c560d..c860a74 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "grok-dev", - "version": "1.1.0", + "version": "1.1.1", "description": "An open-source AI coding agent powered by Grok, built with Bun and OpenTUI.", "type": "module", "main": "dist/index.js", @@ -43,6 +43,7 @@ "license": "MIT", "dependencies": { "@ai-sdk/mcp": "^1.0.25", + "@ai-sdk/provider-utils": "^4.0.21", "@ai-sdk/xai": "^3.0.67", "@modelcontextprotocol/sdk": "^1.27.1", "@opentui/core": "^0.1.88", diff --git a/src/agent/agent.ts b/src/agent/agent.ts index 02b6fd9..b25a1b6 100644 --- a/src/agent/agent.ts +++ b/src/agent/agent.ts @@ -1,6 +1,20 @@ +import { convertToBase64 } from "@ai-sdk/provider-utils"; import { type ModelMessage, stepCountIs, streamText, type ToolSet } from "ai"; +import { + addBatchRequests, + type BatchChatCompletionRequest, + type BatchChatCompletionResponse, + type BatchChatMessage, + type BatchClientOptions, + type BatchFunctionTool, + type BatchToolCall, + createBatch, + getBatchChatCompletion, + pollBatchRequestResult, +} from "../grok/batch"; import { createProvider, generateTitle as genTitle, resolveModelRuntime, type XaiProvider } from "../grok/client"; import { DEFAULT_MODEL, getModelInfo, normalizeModelId } from "../grok/models"; +import { toolSetToBatchTools } from "../grok/tool-schemas"; import { createTools } from "../grok/tools"; import { buildMcpToolSet } from "../mcp/runtime"; import { @@ -65,6 +79,7 @@ interface AgentOptions { session?: string; sandboxMode?: SandboxMode; sandboxSettings?: SandboxSettings; + batchApi?: boolean; } type ProcessMessageFinishReason = "stop" | "length" | "content-filter" | "tool-calls" | "error" | "other"; @@ -73,6 +88,7 @@ export interface ProcessMessageUsage { inputTokens?: number; outputTokens?: number; totalTokens?: number; + costUsdTicks?: number; } export interface ProcessMessageStepStart { @@ -418,6 +434,7 @@ export class Agent { private planContext: string | null = null; private subagentStatusListeners = new Set<(status: SubagentStatus | null) => void>(); private sendTelegramFile: ((filePath: string) => Promise) | null = null; + private batchApi = false; constructor( apiKey: string | undefined, @@ -443,6 +460,7 @@ export class Agent { this.maxToolRounds = maxToolRounds || MAX_TOOL_ROUNDS; const envMax = Number(process.env.GROK_MAX_TOKENS); this.maxTokens = Number.isFinite(envMax) && envMax > 0 ? envMax : 16_384; + this.batchApi = options.batchApi ?? false; if (options.persistSession !== false) { this.sessionStore = new SessionStore(this.bash.getCwd()); @@ -679,6 +697,219 @@ export class Agent { } } + private getBatchClientOptions(signal?: AbortSignal): BatchClientOptions { + if (!this.apiKey) { + throw new Error("API key required. Add an API key to continue."); + } + + return { + apiKey: this.apiKey, + baseURL: this.baseURL ?? undefined, + signal, + }; + } + + private async executeBatchToolCall( + tools: ToolSet, + toolCall: ToolCall, + messages: ModelMessage[], + signal?: AbortSignal, + ): Promise<{ input: unknown; result: ToolResult }> { + const tool = tools[toolCall.function.name]; + if (!tool || tool.type === "provider" || typeof tool.execute !== "function") { + return { + input: parseToolArgumentsOrRaw(toolCall.function.arguments), + result: { + success: false, + output: `Tool "${toolCall.function.name}" is unavailable in batch mode.`, + }, + }; + } + + let parsedInput: unknown; + try { + parsedInput = toolCall.function.arguments.trim() ? JSON.parse(toolCall.function.arguments) : {}; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + return { + input: toolCall.function.arguments, + result: { + success: false, + output: `Tool "${toolCall.function.name}" received invalid JSON arguments: ${message}`, + }, + }; + } + + try { + const output = await tool.execute(parsedInput as never, { + toolCallId: toolCall.id, + messages, + abortSignal: signal, + }); + return { + input: parsedInput, + result: toToolResult(output), + }; + } catch (error) { + if (signal?.aborted) { + throw error; + } + const message = error instanceof Error ? error.message : String(error); + return { + input: parsedInput, + result: { + success: false, + output: `Tool "${toolCall.function.name}" failed: ${message}`, + }, + }; + } + } + + private async runTaskRequestBatch(args: { + request: TaskRequest; + childMessages: ModelMessage[]; + childSystem: string; + childRuntime: ReturnType; + childTools: ToolSet; + maxSteps: number; + initialDetail: string; + onActivity?: (detail: string) => void; + signal?: AbortSignal; + }): Promise { + const { + request, + childMessages, + childSystem, + childRuntime, + childTools, + maxSteps, + initialDetail, + onActivity, + signal, + } = args; + + if (childRuntime.modelInfo?.responsesOnly) { + throw new Error("Batch mode currently supports chat-completions models only."); + } + + const batchTools = + childRuntime.modelInfo?.supportsClientTools === false ? [] : await toolSetToBatchTools(childTools); + const batch = await createBatch({ + ...this.getBatchClientOptions(signal), + name: buildBatchName(`task-${request.agent}`, request.description), + }); + + const turnMessages: ModelMessage[] = []; + const totalUsage: ProcessMessageUsage = {}; + let assistantText = ""; + let lastActivity = initialDetail; + + for (let round = 0; round < maxSteps; round++) { + const batchRequestId = `task-${Date.now()}-${round + 1}`; + await addBatchRequests({ + ...this.getBatchClientOptions(signal), + batchId: batch.batch_id, + batchRequests: [ + { + batch_request_id: batchRequestId, + batch_request: { + chat_get_completion: buildBatchChatCompletionRequest({ + modelId: childRuntime.modelId, + system: childSystem, + messages: [...childMessages, ...turnMessages], + temperature: request.agent === "explore" ? 0.2 : 0.5, + maxOutputTokens: + childRuntime.modelInfo?.supportsMaxOutputTokens === false + ? undefined + : Math.min(this.maxTokens, 8_192), + reasoningEffort: childRuntime.providerOptions?.xai.reasoningEffort, + tools: batchTools, + }), + }, + }, + ], + }); + + const result = await pollBatchRequestResult({ + ...this.getBatchClientOptions(signal), + batchId: batch.batch_id, + batchRequestId, + }); + const response = getBatchChatCompletion(result); + accumulateUsage(totalUsage, getBatchUsage(response)); + + const choice = response.choices[0]; + if (!choice) { + throw new Error("Batch response did not contain any choices."); + } + const content = choice?.message.content ?? ""; + if (content) { + assistantText += content; + } + + const requestMessages = [...childMessages, ...turnMessages]; + const toolCalls = (choice?.message.tool_calls ?? []).map(toLocalToolCall); + const assistantMessage = buildAssistantBatchMessage(content, toolCalls); + if (assistantMessage) { + turnMessages.push(assistantMessage); + } + + if (toolCalls.length === 0) { + if (hasUsage(totalUsage)) { + this.recordUsage(totalUsage, "task", childRuntime.modelId); + } + const output = assistantText.trim() || `Task completed. Last action: ${lastActivity}`; + return { + success: true, + output, + task: { + agent: request.agent, + description: request.description, + summary: firstLine(output), + activity: lastActivity, + }, + }; + } + + const toolParts: ExecutedBatchTool[] = []; + for (const toolCall of toolCalls) { + const nextActivity = formatSubagentActivity( + toolCall.function.name, + parseToolArgumentsOrRaw(toolCall.function.arguments), + ); + lastActivity = nextActivity; + onActivity?.(nextActivity); + + const executed = await this.executeBatchToolCall(childTools, toolCall, requestMessages, signal); + toolParts.push({ + toolCall, + input: executed.input, + toolResult: executed.result, + }); + } + + const toolMessage = buildToolBatchMessage(toolParts); + if (toolMessage) { + turnMessages.push(toolMessage); + } + } + + if (hasUsage(totalUsage)) { + this.recordUsage(totalUsage, "task", childRuntime.modelId); + } + const output = assistantText.trim() || `Task stopped after ${maxSteps} batch rounds. Last action: ${lastActivity}`; + return { + success: false, + output, + task: { + agent: request.agent, + description: request.description, + summary: output, + activity: lastActivity, + }, + }; + } + async runTaskRequest( request: TaskRequest, onActivity?: (detail: string) => void, @@ -768,6 +999,20 @@ export class Agent { ? await buildVisionUserMessages(request.prompt, childBash.getCwd(), signal) : [{ role: "user" as const, content: request.prompt }]; + if (this.batchApi) { + return await this.runTaskRequestBatch({ + request, + childMessages, + childSystem, + childRuntime, + childTools, + maxSteps: Math.min(this.maxToolRounds, isExplore ? 60 : 120), + initialDetail, + onActivity, + signal, + }); + } + const result = streamText({ model: childRuntime.model, system: childSystem, @@ -869,6 +1114,7 @@ export class Agent { sandboxSettings: this.bash.getSandboxSettings(), maxToolRounds: this.maxToolRounds, maxTokens: this.maxTokens, + batchApi: this.batchApi, }); } catch (err: unknown) { if (abortSignal?.aborted) throw err; @@ -956,6 +1202,218 @@ export class Agent { return true; } + private async *processMessageBatchTurn(args: { + userModelMessage: ModelMessage; + observer?: ProcessMessageObserver; + provider: XaiProvider; + subagents: CustomSubagentConfig[]; + system: string; + runtime: ReturnType; + modelInfo: ReturnType; + signal: AbortSignal; + }): AsyncGenerator { + const { userModelMessage, observer, provider, subagents, system, runtime, modelInfo, signal } = args; + let attemptedOverflowRecovery = false; + + while (true) { + let closeMcp: (() => Promise) | undefined; + const turnMessages: ModelMessage[] = []; + const totalUsage: ProcessMessageUsage = {}; + + try { + const settings = attemptedOverflowRecovery + ? relaxCompactionSettings(this.getCompactionSettings()) + : this.getCompactionSettings(); + if (modelInfo) { + await this.compactForContext( + provider, + system, + modelInfo.contextWindow, + signal, + settings, + attemptedOverflowRecovery, + ); + } + + if (runtime.modelInfo?.responsesOnly) { + throw new Error("Batch mode currently supports chat-completions models only."); + } + + const baseTools = createTools(this.bash, provider, this.mode, { + runTask: (request, abortSignal) => this.runTask(request, combineAbortSignals(signal, abortSignal)), + runDelegation: (request, abortSignal) => + this.runDelegation(request, combineAbortSignals(signal, abortSignal)), + readDelegation: (id) => this.readDelegation(id), + listDelegations: () => this.listDelegations(), + scheduleManager: this.schedules, + subagents, + sendTelegramFile: this.sendTelegramFile ?? undefined, + }); + let tools: ToolSet = runtime.modelInfo?.supportsClientTools === false ? {} : baseTools; + if (this.mode === "agent" && runtime.modelInfo?.supportsClientTools !== false) { + const mcpBundle = await buildMcpToolSet(loadMcpServers()); + closeMcp = mcpBundle.close; + tools = { ...baseTools, ...mcpBundle.tools }; + if (mcpBundle.errors.length > 0) { + yield { type: "content", content: `MCP unavailable: ${mcpBundle.errors.join(" | ")}\n\n` }; + } + } + + const batchTools = runtime.modelInfo?.supportsClientTools === false ? [] : await toolSetToBatchTools(tools); + const batch = await createBatch({ + ...this.getBatchClientOptions(signal), + name: buildBatchName("session", this.getSessionId() || runtime.modelId), + }); + + for (let round = 0; round < this.maxToolRounds; round++) { + const stepNumber = round + 1; + notifyObserver(observer?.onStepStart, { + stepNumber, + timestamp: Date.now(), + }); + + const batchRequestId = `turn-${Date.now()}-${stepNumber}`; + await addBatchRequests({ + ...this.getBatchClientOptions(signal), + batchId: batch.batch_id, + batchRequests: [ + { + batch_request_id: batchRequestId, + batch_request: { + chat_get_completion: buildBatchChatCompletionRequest({ + modelId: runtime.modelId, + system, + messages: [...this.messages, ...turnMessages], + temperature: 0.7, + maxOutputTokens: runtime.modelInfo?.supportsMaxOutputTokens === false ? undefined : this.maxTokens, + reasoningEffort: runtime.providerOptions?.xai.reasoningEffort, + tools: batchTools, + }), + }, + }, + ], + }); + + const result = await pollBatchRequestResult({ + ...this.getBatchClientOptions(signal), + batchId: batch.batch_id, + batchRequestId, + }); + const response = getBatchChatCompletion(result); + const choice = response.choices[0]; + if (!choice) { + throw new Error("Batch response did not contain any choices."); + } + + const usage = getBatchUsage(response); + accumulateUsage(totalUsage, usage); + const finishReason = getBatchFinishReason(choice.finish_reason); + + const content = choice.message.content ?? ""; + if (content) { + yield { type: "content", content }; + } + + const requestMessages = [...this.messages, ...turnMessages]; + const toolCalls = (choice.message.tool_calls ?? []).map(toLocalToolCall); + const assistantMessage = buildAssistantBatchMessage(content, toolCalls); + if (assistantMessage) { + turnMessages.push(assistantMessage); + } + + if (toolCalls.length === 0) { + notifyObserver(observer?.onStepFinish, { + stepNumber, + timestamp: Date.now(), + finishReason, + usage, + }); + if (hasUsage(totalUsage)) { + this.recordUsage(totalUsage, "message", runtime.modelId); + } + this.appendCompletedTurn(userModelMessage, turnMessages); + yield { type: "done" }; + return; + } + + yield { type: "tool_calls", toolCalls }; + + const toolParts: ExecutedBatchTool[] = []; + for (const toolCall of toolCalls) { + notifyObserver(observer?.onToolStart, { + toolCall, + timestamp: Date.now(), + }); + + const executed = await this.executeBatchToolCall(tools, toolCall, requestMessages, signal); + notifyObserver(observer?.onToolFinish, { + toolCall, + toolResult: executed.result, + timestamp: Date.now(), + }); + yield { type: "tool_result", toolCall, toolResult: executed.result }; + toolParts.push({ + toolCall, + input: executed.input, + toolResult: executed.result, + }); + } + + const toolMessage = buildToolBatchMessage(toolParts); + if (toolMessage) { + turnMessages.push(toolMessage); + } + notifyObserver(observer?.onStepFinish, { + stepNumber, + timestamp: Date.now(), + finishReason, + usage, + }); + } + + const message = `Error: Reached max tool rounds (${this.maxToolRounds}) in batch mode.`; + notifyObserver(observer?.onError, { + message, + timestamp: Date.now(), + }); + if (hasUsage(totalUsage)) { + this.recordUsage(totalUsage, "message", runtime.modelId); + } + this.appendCompletedTurn(userModelMessage, turnMessages); + yield { type: "error", content: message }; + yield { type: "done" }; + return; + } catch (err: unknown) { + if (signal.aborted) { + this.discardAbortedTurn(userModelMessage); + yield { type: "content", content: "\n\n[Cancelled]" }; + yield { type: "done" }; + return; + } + + if (!attemptedOverflowRecovery && turnMessages.length === 0 && modelInfo && isContextLimitError(err)) { + attemptedOverflowRecovery = true; + continue; + } + + const msg = err instanceof Error ? err.message : String(err); + notifyObserver(observer?.onError, { + message: `Error: ${msg}`, + timestamp: Date.now(), + }); + if (hasUsage(totalUsage)) { + this.recordUsage(totalUsage, "message", runtime.modelId); + } + this.appendCompletedTurn(userModelMessage, turnMessages); + yield { type: "error", content: `Error: ${msg}` }; + yield { type: "done" }; + return; + } finally { + await closeMcp?.().catch(() => {}); + } + } + } + private appendCompletedTurn(userMessage: ModelMessage, newMessages: ModelMessage[]): void { if (newMessages.length === 0) return; @@ -1010,6 +1468,26 @@ export class Agent { this.planContext = null; let attemptedOverflowRecovery = false; + if (this.batchApi) { + try { + yield* this.processMessageBatchTurn({ + userModelMessage, + observer, + provider, + subagents, + system, + runtime, + modelInfo, + signal, + }); + } finally { + if (this.abortController?.signal === signal) { + this.abortController = null; + } + } + return; + } + try { while (true) { let assistantText = ""; @@ -1253,6 +1731,14 @@ export class Agent { } } +interface ExecutedBatchTool { + toolCall: ToolCall; + input: unknown; + toolResult: ToolResult; +} + +type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue }; + function extractJsonObject(text: string): string | null { const start = text.indexOf("{"); const end = text.lastIndexOf("}"); @@ -1260,6 +1746,296 @@ function extractJsonObject(text: string): string | null { return text.slice(start, end + 1); } +function buildBatchName(prefix: string, label: string): string { + const compact = + label + .replace(/\s+/g, "-") + .replace(/[^a-zA-Z0-9._-]+/g, "") + .slice(0, 48) || "run"; + return `grok-cli-${prefix}-${compact}`; +} + +function buildBatchChatCompletionRequest(args: { + modelId: string; + system: string; + messages: ModelMessage[]; + temperature: number; + maxOutputTokens?: number; + reasoningEffort?: BatchChatCompletionRequest["reasoning_effort"]; + tools: BatchFunctionTool[]; +}): BatchChatCompletionRequest { + return { + model: args.modelId, + messages: toBatchChatMessages(args.system, args.messages), + temperature: args.temperature, + ...(args.maxOutputTokens != null ? { max_completion_tokens: args.maxOutputTokens } : {}), + ...(args.reasoningEffort ? { reasoning_effort: args.reasoningEffort } : {}), + ...(args.tools.length > 0 ? { tools: args.tools } : {}), + }; +} + +function toBatchChatMessages(system: string, messages: ModelMessage[]): BatchChatMessage[] { + const batchMessages: BatchChatMessage[] = [{ role: "system", content: system }]; + + for (const message of messages) { + const { role, content } = message; + + switch (role) { + case "system": + batchMessages.push({ role: "system", content }); + break; + + case "user": { + if (typeof content === "string") { + batchMessages.push({ role: "user", content }); + break; + } + + if (!Array.isArray(content)) { + break; + } + + if (content.length === 1 && content[0]?.type === "text") { + batchMessages.push({ role: "user", content: content[0].text }); + break; + } + + const userContent: Array<{ type: "text"; text: string } | { type: "image_url"; image_url: { url: string } }> = + []; + for (const part of content) { + switch (part.type) { + case "text": + userContent.push({ type: "text", text: part.text }); + break; + + case "image": { + const mediaType = part.mediaType === "image/*" || !part.mediaType ? "image/jpeg" : part.mediaType; + const data = + part.image instanceof URL + ? part.image.toString() + : `data:${mediaType};base64,${toBase64DataContent(part.image)}`; + userContent.push({ type: "image_url", image_url: { url: data } }); + break; + } + + case "file": { + if (!part.mediaType.startsWith("image/")) { + break; + } + const mediaType = part.mediaType === "image/*" ? "image/jpeg" : part.mediaType; + const data = + part.data instanceof URL + ? part.data.toString() + : `data:${mediaType};base64,${toBase64DataContent(part.data)}`; + userContent.push({ type: "image_url", image_url: { url: data } }); + break; + } + } + } + batchMessages.push({ + role: "user", + content: userContent, + }); + break; + } + + case "assistant": { + if (typeof content === "string") { + batchMessages.push({ role: "assistant", content }); + break; + } + + if (!Array.isArray(content)) { + break; + } + + let assistantText = ""; + const toolCalls: BatchToolCall[] = []; + for (const part of content) { + if (part.type === "text") { + assistantText += part.text; + } else if (part.type === "tool-call") { + toolCalls.push({ + id: part.toolCallId, + type: "function", + function: { + name: part.toolName, + arguments: JSON.stringify(part.input), + }, + }); + } + } + + if (assistantText || toolCalls.length > 0) { + batchMessages.push({ + role: "assistant", + content: assistantText, + ...(toolCalls.length > 0 ? { tool_calls: toolCalls } : {}), + }); + } + break; + } + + case "tool": + for (const part of content) { + if (part.type === "tool-approval-response") { + continue; + } + batchMessages.push({ + role: "tool", + tool_call_id: part.toolCallId, + content: toolOutputToText(part.output), + }); + } + break; + } + } + + return batchMessages; +} + +function toBase64DataContent(value: string | Uint8Array | ArrayBuffer): string { + return convertToBase64(value instanceof ArrayBuffer ? new Uint8Array(value) : value); +} + +function toolOutputToText(output: { + type: "text" | "json" | "execution-denied" | "error-text" | "error-json" | "content"; + value?: unknown; + reason?: string; +}): string { + switch (output.type) { + case "text": + case "error-text": + return String(output.value ?? ""); + case "execution-denied": + return output.reason ?? "Tool execution denied."; + case "json": + case "error-json": + case "content": + return JSON.stringify(output.value ?? null); + } +} + +function getBatchUsage(response: BatchChatCompletionResponse): ProcessMessageUsage { + const usage = response.usage ?? {}; + const inputTokens = asNumber(usage.input_tokens) ?? asNumber(usage.prompt_tokens); + const outputTokens = asNumber(usage.output_tokens) ?? asNumber(usage.completion_tokens); + const totalTokens = asNumber(usage.total_tokens) ?? sumDefined(inputTokens, outputTokens); + return { + inputTokens, + outputTokens, + totalTokens, + costUsdTicks: asNumber(usage.cost_in_usd_ticks), + }; +} + +function accumulateUsage(target: ProcessMessageUsage, usage: ProcessMessageUsage): void { + target.inputTokens = (target.inputTokens ?? 0) + (usage.inputTokens ?? 0); + target.outputTokens = (target.outputTokens ?? 0) + (usage.outputTokens ?? 0); + target.totalTokens = (target.totalTokens ?? 0) + (usage.totalTokens ?? 0); + target.costUsdTicks = (target.costUsdTicks ?? 0) + (usage.costUsdTicks ?? 0); +} + +function hasUsage(usage: ProcessMessageUsage): boolean { + return Boolean( + (usage.inputTokens ?? 0) || (usage.outputTokens ?? 0) || (usage.totalTokens ?? 0) || (usage.costUsdTicks ?? 0), + ); +} + +function getBatchFinishReason(finishReason: string | null | undefined): ProcessMessageFinishReason { + switch (finishReason) { + case "stop": + case "length": + case "content-filter": + case "tool-calls": + case "error": + case "other": + return finishReason; + case "tool_calls": + return "tool-calls"; + default: + return "other"; + } +} + +function toLocalToolCall(toolCall: BatchToolCall): ToolCall { + return { + id: toolCall.id, + type: "function", + function: { + name: toolCall.function.name, + arguments: toolCall.function.arguments, + }, + }; +} + +function buildAssistantBatchMessage(content: string, toolCalls: ToolCall[]): ModelMessage | null { + if (toolCalls.length === 0) { + return content ? { role: "assistant", content } : null; + } + + const parts: Array< + { type: "text"; text: string } | { type: "tool-call"; toolCallId: string; toolName: string; input: unknown } + > = []; + if (content) { + parts.push({ type: "text", text: content }); + } + for (const toolCall of toolCalls) { + parts.push({ + type: "tool-call", + toolCallId: toolCall.id, + toolName: toolCall.function.name, + input: parseToolArgumentsOrRaw(toolCall.function.arguments), + }); + } + return { role: "assistant", content: parts }; +} + +function buildToolBatchMessage(toolParts: ExecutedBatchTool[]): ModelMessage | null { + if (toolParts.length === 0) { + return null; + } + + return { + role: "tool", + content: toolParts.map((part) => ({ + type: "tool-result" as const, + toolCallId: part.toolCall.id, + toolName: part.toolCall.function.name, + output: part.toolResult.success + ? ({ type: "json", value: toSerializableValue(part.toolResult) } as const) + : ({ type: "error-json", value: toSerializableValue(part.toolResult) } as const), + })), + }; +} + +function parseToolArgumentsOrRaw(raw: string): unknown { + try { + return raw.trim() ? JSON.parse(raw) : {}; + } catch { + return raw; + } +} + +function toSerializableValue(value: unknown): JsonValue { + try { + return JSON.parse(JSON.stringify(value ?? null)) as JsonValue; + } catch { + return String(value); + } +} + +function asNumber(value: unknown): number | undefined { + return typeof value === "number" ? value : undefined; +} + +function sumDefined(left?: number, right?: number): number | undefined { + if (left == null && right == null) { + return undefined; + } + return (left ?? 0) + (right ?? 0); +} + function toToolCall(part: { toolCallId: string; toolName: string; args?: unknown; input?: unknown }): ToolCall { return { id: part.toolCallId, diff --git a/src/agent/batch-mode.test.ts b/src/agent/batch-mode.test.ts new file mode 100644 index 0000000..3c34c3b --- /dev/null +++ b/src/agent/batch-mode.test.ts @@ -0,0 +1,130 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +async function importAgentModuleWithBatchMocks() { + vi.resetModules(); + vi.doMock("../storage/index", () => ({ + appendCompaction: vi.fn(), + appendMessages: vi.fn(() => []), + appendSystemMessage: vi.fn(() => 0), + buildChatEntries: vi.fn(() => []), + getNextMessageSequence: vi.fn(() => 0), + getSessionTotalTokens: vi.fn(() => 0), + loadTranscript: vi.fn(() => []), + loadTranscriptState: vi.fn(() => ({ messages: [], seqs: [] })), + recordUsageEvent: vi.fn(), + SessionStore: class { + getWorkspace() { + return null; + } + openSession() { + return null; + } + createSession() { + return null; + } + setModel() {} + getRequiredSession() { + return null; + } + setMode() {} + touchSession() {} + }, + })); + + const createBatch = vi.fn(async () => ({ + batch_id: "batch-1", + state: { + num_requests: 0, + num_pending: 0, + num_success: 0, + num_error: 0, + num_cancelled: 0, + }, + })); + const addBatchRequests = vi.fn(async () => {}); + const pollBatchRequestResult = vi.fn(async () => ({ batch_request_id: "req-1" })); + const getBatchChatCompletion = vi.fn(() => ({ + choices: [], + usage: {}, + })); + + vi.doMock("../grok/batch", async () => { + const actual = await vi.importActual("../grok/batch"); + return { + ...actual, + createBatch, + addBatchRequests, + pollBatchRequestResult, + getBatchChatCompletion, + }; + }); + + const mod = await import("./agent"); + return { + ...mod, + mocks: { + createBatch, + addBatchRequests, + pollBatchRequestResult, + getBatchChatCompletion, + }, + }; +} + +afterEach(() => { + vi.restoreAllMocks(); + vi.resetModules(); + vi.doUnmock("../storage/index"); + vi.doUnmock("../grok/batch"); +}); + +describe("Agent batch mode", () => { + it("throws when a child batch response has no choices", async () => { + const { Agent, mocks } = await importAgentModuleWithBatchMocks(); + const agent = new Agent("test-key", "https://api.x.ai/v1", undefined, undefined, { + persistSession: false, + }); + + const runTaskRequestBatch = ( + agent as unknown as { + runTaskRequestBatch: (args: { + request: { agent: string; description: string; prompt: string }; + childMessages: Array<{ role: "user"; content: string }>; + childSystem: string; + childRuntime: unknown; + childTools: Record; + maxSteps: number; + initialDetail: string; + }) => Promise; + } + ).runTaskRequestBatch.bind(agent); + + await expect( + runTaskRequestBatch({ + request: { + agent: "general", + description: "Child task", + prompt: "Do the thing", + }, + childMessages: [{ role: "user", content: "Do the thing" }], + childSystem: "system", + childRuntime: { + modelId: "grok-4-1-fast-reasoning", + modelInfo: { + supportsClientTools: false, + supportsMaxOutputTokens: true, + }, + providerOptions: undefined, + }, + childTools: {}, + maxSteps: 1, + initialDetail: "Starting child task", + }), + ).rejects.toThrow("Batch response did not contain any choices."); + + expect(mocks.createBatch).toHaveBeenCalled(); + expect(mocks.addBatchRequests).toHaveBeenCalled(); + expect(mocks.pollBatchRequestResult).toHaveBeenCalled(); + expect(mocks.getBatchChatCompletion).toHaveBeenCalled(); + }); +}); diff --git a/src/agent/delegations.test.ts b/src/agent/delegations.test.ts index 046f115..613c882 100644 --- a/src/agent/delegations.test.ts +++ b/src/agent/delegations.test.ts @@ -4,6 +4,7 @@ import path from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; const tempDirs: string[] = []; +const originalHome = process.env.HOME; function makeTempDir(prefix: string): string { const dir = fs.mkdtempSync(path.join(os.tmpdir(), prefix)); @@ -11,19 +12,22 @@ function makeTempDir(prefix: string): string { return dir; } +function readStoredDelegationRecord(home: string): Record { + const delegationsRoot = path.join(home, ".grok", "delegations"); + const projectDirs = fs.readdirSync(delegationsRoot); + expect(projectDirs).toHaveLength(1); + const projectDir = path.join(delegationsRoot, projectDirs[0] as string); + const files = fs.readdirSync(projectDir).filter((file) => file.endsWith(".json")); + expect(files).toHaveLength(1); + return JSON.parse(fs.readFileSync(path.join(projectDir, files[0] as string), "utf8")) as Record; +} + async function importDelegationsModule(options: { home?: string; spawnMock?: ReturnType } = {}) { vi.resetModules(); - vi.doUnmock("os"); vi.doUnmock("child_process"); if (options.home) { - vi.doMock("os", async () => { - const actual = await vi.importActual("os"); - return { - ...actual, - homedir: () => options.home!, - }; - }); + process.env.HOME = options.home; } if (options.spawnMock) { @@ -43,9 +47,9 @@ afterEach(() => { for (const dir of tempDirs.splice(0)) { fs.rmSync(dir, { recursive: true, force: true }); } + process.env.HOME = originalHome; vi.restoreAllMocks(); vi.resetModules(); - vi.doUnmock("os"); vi.doUnmock("child_process"); }); @@ -81,9 +85,7 @@ describe("DelegationManager sandbox propagation", () => { expect.objectContaining({ cwd }), ); - const cliArgs = (spawnMock.mock.calls[0] as unknown[])?.[1] as string[]; - const jobPath = cliArgs[cliArgs.indexOf("--background-task-file") + 1] as string; - const record = JSON.parse(fs.readFileSync(jobPath, "utf8")) as { + const record = readStoredDelegationRecord(home) as { sandboxMode: string; pid: number; }; @@ -124,9 +126,7 @@ describe("DelegationManager sandbox propagation", () => { }, ); - const cliArgs = (spawnMock.mock.calls[0] as unknown[])?.[1] as string[]; - const jobPath = cliArgs[cliArgs.indexOf("--background-task-file") + 1] as string; - const record = JSON.parse(fs.readFileSync(jobPath, "utf8")) as { + const record = readStoredDelegationRecord(home) as { sandboxMode: string; sandboxSettings?: { allowNet: boolean; diff --git a/src/agent/delegations.ts b/src/agent/delegations.ts index 20c5745..e245b04 100644 --- a/src/agent/delegations.ts +++ b/src/agent/delegations.ts @@ -23,6 +23,7 @@ export interface StoredDelegation { sandboxSettings?: SandboxSettings; maxToolRounds: number; maxTokens: number; + batchApi?: boolean; status: DelegationStatus; startedAt: string; completedAt?: string; @@ -45,6 +46,7 @@ interface StartDelegationOptions { sandboxSettings?: SandboxSettings; maxToolRounds: number; maxTokens: number; + batchApi?: boolean; } export class DelegationManager { @@ -83,6 +85,7 @@ export class DelegationManager { sandboxSettings: options.sandboxSettings, maxToolRounds: options.maxToolRounds, maxTokens: options.maxTokens, + batchApi: options.batchApi, status: "running", startedAt: new Date().toISOString(), outputPath, @@ -92,7 +95,14 @@ export class DelegationManager { const child = spawn( process.execPath, - [...resolveCliArgs(), "--directory", cwd, "--background-task-file", jobPath], + [ + ...resolveCliArgs(), + "--directory", + cwd, + "--background-task-file", + jobPath, + ...(options.batchApi ? ["--batch-api"] : []), + ], { cwd, detached: true, diff --git a/src/grok/batch.test.ts b/src/grok/batch.test.ts new file mode 100644 index 0000000..3e2654e --- /dev/null +++ b/src/grok/batch.test.ts @@ -0,0 +1,166 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { getBatchChatCompletion, pollBatchRequestResult } from "./batch"; + +function jsonResponse(body: unknown, init?: ResponseInit): Response { + return new Response(JSON.stringify(body), { + status: 200, + headers: { "Content-Type": "application/json" }, + ...init, + }); +} + +describe("getBatchChatCompletion", () => { + afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); + }); + + it("extracts the chat completion payload from a batch result", () => { + const response = getBatchChatCompletion({ + batch_request_id: "req-1", + batch_result: { + response: { + chat_get_completion: { + id: "chatcmpl-1", + choices: [ + { + finish_reason: "stop", + message: { + content: "done", + }, + }, + ], + usage: { + total_tokens: 12, + }, + }, + }, + }, + }); + + expect(response.choices[0]?.message.content).toBe("done"); + expect(response.usage?.total_tokens).toBe(12); + }); + + it("throws when the batch result is an error", () => { + expect(() => + getBatchChatCompletion({ + batch_request_id: "req-2", + error_message: "boom", + }), + ).toThrow('Batch request "req-2" failed: boom'); + }); + + it("waits through transient ClickHouse visibility errors before returning a result", async () => { + const fetchMock = vi + .fn() + .mockResolvedValueOnce( + jsonResponse({ + results: [ + { + batch_request_id: "req-3", + batch_result: { + error: "Batch request not yet visible in ClickHouse", + }, + }, + ], + pagination_token: null, + }), + ) + .mockResolvedValueOnce( + jsonResponse({ + batch_id: "batch-1", + state: { + num_requests: 1, + num_pending: 1, + num_success: 0, + num_error: 0, + num_cancelled: 0, + }, + }), + ) + .mockResolvedValueOnce( + jsonResponse({ + results: [ + { + batch_request_id: "req-3", + batch_result: { + response: { + chat_get_completion: { + choices: [{ message: { content: "done" } }], + }, + }, + }, + }, + ], + pagination_token: null, + }), + ); + vi.stubGlobal("fetch", fetchMock); + + const result = await pollBatchRequestResult({ + apiKey: "test-key", + batchId: "batch-1", + batchRequestId: "req-3", + initialPollMs: 0, + maxPollMs: 0, + timeoutMs: 1000, + }); + + expect(result.batch_result?.response?.chat_get_completion?.choices[0]?.message.content).toBe("done"); + expect(fetchMock).toHaveBeenCalledTimes(3); + }); + + it("keeps polling after batch completion until the result row becomes visible", async () => { + const fetchMock = vi + .fn() + .mockResolvedValueOnce( + jsonResponse({ + results: [], + pagination_token: null, + }), + ) + .mockResolvedValueOnce( + jsonResponse({ + batch_id: "batch-2", + state: { + num_requests: 1, + num_pending: 0, + num_success: 1, + num_error: 0, + num_cancelled: 0, + }, + }), + ) + .mockResolvedValueOnce( + jsonResponse({ + results: [ + { + batch_request_id: "req-4", + batch_result: { + response: { + chat_get_completion: { + choices: [{ message: { content: "visible now" } }], + }, + }, + }, + }, + ], + pagination_token: null, + }), + ); + vi.stubGlobal("fetch", fetchMock); + + const result = await pollBatchRequestResult({ + apiKey: "test-key", + batchId: "batch-2", + batchRequestId: "req-4", + initialPollMs: 0, + maxPollMs: 0, + timeoutMs: 1000, + }); + + expect(result.batch_result?.response?.chat_get_completion?.choices[0]?.message.content).toBe("visible now"); + expect(fetchMock).toHaveBeenCalledTimes(3); + }); +}); diff --git a/src/grok/batch.ts b/src/grok/batch.ts new file mode 100644 index 0000000..36fccca --- /dev/null +++ b/src/grok/batch.ts @@ -0,0 +1,405 @@ +import process from "node:process"; + +const DEFAULT_BASE_URL = "https://api.x.ai/v1"; +const DEFAULT_INITIAL_POLL_MS = 2_000; +const DEFAULT_MAX_POLL_MS = 30_000; +const DEFAULT_TIMEOUT_MS = 10 * 60_000; +const DEFAULT_BATCH_PAGE_SIZE = 100; +const DEFAULT_RATE_LIMIT_RETRIES = 5; + +export interface BatchClientOptions { + apiKey: string; + baseURL?: string; + signal?: AbortSignal; +} + +export interface BatchState { + num_requests: number; + num_pending: number; + num_success: number; + num_error: number; + num_cancelled: number; +} + +export interface BatchCostBreakdown { + total_cost_usd_ticks?: number; +} + +export interface BatchInfo { + batch_id: string; + name?: string; + expires_at?: string; + state: BatchState; + cost_breakdown?: BatchCostBreakdown; +} + +export interface BatchFunctionTool { + type: "function"; + function: { + name: string; + description?: string; + parameters: unknown; + strict?: boolean; + }; +} + +export interface BatchToolCall { + id: string; + type: "function"; + function: { + name: string; + arguments: string; + }; +} + +export interface BatchChatMessage { + role: "system" | "user" | "assistant" | "tool"; + content: string | Array<{ type: "text"; text: string } | { type: "image_url"; image_url: { url: string } }>; + tool_calls?: BatchToolCall[]; + tool_call_id?: string; +} + +export interface BatchChatCompletionRequest { + model: string; + messages: BatchChatMessage[]; + temperature?: number; + max_completion_tokens?: number; + reasoning_effort?: "low" | "medium" | "high" | "xhigh"; + tools?: BatchFunctionTool[]; + tool_choice?: + | "auto" + | "none" + | "required" + | { + type: "function"; + function: { + name: string; + }; + }; +} + +export interface BatchRequestEntry { + batch_request_id: string; + batch_request: { + chat_get_completion: BatchChatCompletionRequest; + }; +} + +export interface BatchChatCompletionChoice { + index?: number; + finish_reason?: string | null; + message: { + role?: "assistant"; + content?: string | null; + tool_calls?: BatchToolCall[]; + }; +} + +export interface BatchChatCompletionUsage { + prompt_tokens?: number; + completion_tokens?: number; + total_tokens?: number; + input_tokens?: number; + output_tokens?: number; + cost_in_usd_ticks?: number; +} + +export interface BatchChatCompletionResponse { + id?: string; + model?: string; + choices: BatchChatCompletionChoice[]; + usage?: BatchChatCompletionUsage; +} + +export interface BatchResultEntry { + batch_request_id: string; + error_message?: string; + batch_result?: { + error?: string; + response?: { + chat_get_completion?: BatchChatCompletionResponse; + }; + }; +} + +export interface BatchResultsPage { + results: BatchResultEntry[]; + pagination_token?: string | null; +} + +export interface PollBatchRequestOptions extends BatchClientOptions { + batchId: string; + batchRequestId: string; + timeoutMs?: number; + initialPollMs?: number; + maxPollMs?: number; + pageSize?: number; +} + +export async function createBatch( + options: BatchClientOptions & { + name: string; + inputFileId?: string; + }, +): Promise { + return requestJson({ + ...options, + method: "POST", + path: "/batches", + body: { + name: options.name, + ...(options.inputFileId ? { input_file_id: options.inputFileId } : {}), + }, + retryRateLimit: true, + }); +} + +export async function addBatchRequests( + options: BatchClientOptions & { + batchId: string; + batchRequests: BatchRequestEntry[]; + }, +): Promise { + await requestJson({ + ...options, + method: "POST", + path: `/batches/${encodeURIComponent(options.batchId)}/requests`, + body: { + batch_requests: options.batchRequests, + }, + }); +} + +export async function getBatchStatus( + options: BatchClientOptions & { + batchId: string; + }, +): Promise { + return requestJson({ + ...options, + method: "GET", + path: `/batches/${encodeURIComponent(options.batchId)}`, + }); +} + +export async function listBatchResults( + options: BatchClientOptions & { + batchId: string; + pageSize?: number; + paginationToken?: string; + }, +): Promise { + const search = new URLSearchParams(); + search.set("page_size", String(options.pageSize ?? DEFAULT_BATCH_PAGE_SIZE)); + if (options.paginationToken) { + search.set("pagination_token", options.paginationToken); + } + + return requestJson({ + ...options, + method: "GET", + path: `/batches/${encodeURIComponent(options.batchId)}/results?${search.toString()}`, + }); +} + +export async function pollBatchRequestResult(options: PollBatchRequestOptions): Promise { + const timeoutMs = options.timeoutMs ?? DEFAULT_TIMEOUT_MS; + const initialPollMs = options.initialPollMs ?? DEFAULT_INITIAL_POLL_MS; + const maxPollMs = options.maxPollMs ?? DEFAULT_MAX_POLL_MS; + const pageSize = options.pageSize ?? DEFAULT_BATCH_PAGE_SIZE; + const startedAt = Date.now(); + let delayMs = initialPollMs; + + while (true) { + throwIfAborted(options.signal); + + const match = await findBatchResult({ + ...options, + pageSize, + }); + if (match && isBatchResultReady(match)) { + return match; + } + + await getBatchStatus(options); + + if (Date.now() - startedAt >= timeoutMs) { + throw new Error( + `Timed out waiting for batch request "${options.batchRequestId}" in batch "${options.batchId}" after ${Math.round(timeoutMs / 1000)}s.`, + ); + } + + await sleep(delayMs, options.signal); + delayMs = Math.min(Math.round(delayMs * 1.5), maxPollMs); + } +} + +export function getBatchChatCompletion(result: BatchResultEntry): BatchChatCompletionResponse { + if (result.error_message) { + throw new Error(`Batch request "${result.batch_request_id}" failed: ${result.error_message}`); + } + + if (result.batch_result?.error) { + throw new Error(`Batch request "${result.batch_request_id}" failed: ${result.batch_result.error}`); + } + + const response = result.batch_result?.response?.chat_get_completion; + if (!response) { + throw new Error(`Batch request "${result.batch_request_id}" did not return a chat completion result.`); + } + + return response; +} + +async function findBatchResult( + options: BatchClientOptions & { + batchId: string; + batchRequestId: string; + pageSize: number; + }, +): Promise { + let paginationToken: string | undefined; + + while (true) { + const page = await listBatchResults({ + ...options, + pageSize: options.pageSize, + paginationToken, + }); + const match = page.results.find((entry) => entry.batch_request_id === options.batchRequestId); + if (match) { + return match; + } + + paginationToken = page.pagination_token ?? undefined; + if (!paginationToken) { + return null; + } + } +} + +function isBatchResultReady(result: BatchResultEntry): boolean { + if (result.error_message) { + return true; + } + if (result.batch_result?.response?.chat_get_completion) { + return true; + } + if (result.batch_result?.error && !isTransientBatchResultError(result.batch_result.error)) { + return true; + } + return false; +} + +function isTransientBatchResultError(message: string): boolean { + return message.includes("not yet visible in ClickHouse"); +} + +async function requestJson( + options: BatchClientOptions & { + method: "GET" | "POST"; + path: string; + body?: unknown; + retryRateLimit?: boolean; + }, +): Promise { + const url = `${normalizeBaseUrl(options.baseURL)}${options.path.startsWith("/") ? options.path : `/${options.path}`}`; + let attempt = 0; + let delayMs = 500; + + while (true) { + throwIfAborted(options.signal); + + const response = await fetch(url, { + method: options.method, + headers: { + Authorization: `Bearer ${options.apiKey}`, + ...(options.body ? { "Content-Type": "application/json" } : {}), + }, + body: options.body ? JSON.stringify(options.body) : undefined, + signal: options.signal, + }); + + if (response.ok) { + return readJson(response); + } + + if (options.retryRateLimit && response.status === 429 && attempt < DEFAULT_RATE_LIMIT_RETRIES) { + const retryAfterMs = getRetryAfterMs(response) ?? delayMs; + await sleep(retryAfterMs, options.signal); + attempt += 1; + delayMs = Math.min(delayMs * 2, 5_000); + continue; + } + + const errorText = await response.text(); + const suffix = errorText.trim() ? `: ${errorText.trim()}` : ""; + throw new Error(`Batch API ${options.method} ${options.path} failed (${response.status})${suffix}`); + } +} + +async function readJson(response: Response): Promise { + const text = await response.text(); + if (!text.trim()) { + return {} as T; + } + + try { + return JSON.parse(text) as T; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + throw new Error(`Batch API returned invalid JSON: ${message}`); + } +} + +function normalizeBaseUrl(baseURL?: string): string { + return (baseURL || process.env.GROK_BASE_URL || DEFAULT_BASE_URL).replace(/\/+$/, ""); +} + +function getRetryAfterMs(response: Response): number | undefined { + const retryAfter = response.headers.get("retry-after"); + if (!retryAfter) return undefined; + + const seconds = Number(retryAfter); + if (Number.isFinite(seconds) && seconds >= 0) { + return seconds * 1000; + } + + const dateMs = Date.parse(retryAfter); + if (!Number.isNaN(dateMs)) { + return Math.max(0, dateMs - Date.now()); + } + + return undefined; +} + +function throwIfAborted(signal?: AbortSignal): void { + if (signal?.aborted) { + throw new Error("Request aborted"); + } +} + +function sleep(ms: number, signal?: AbortSignal): Promise { + if (ms <= 0) { + return Promise.resolve(); + } + + return new Promise((resolve, reject) => { + const timeout = setTimeout(() => { + cleanup(); + resolve(); + }, ms); + + const onAbort = () => { + clearTimeout(timeout); + cleanup(); + reject(new Error("Request aborted")); + }; + + const cleanup = () => { + signal?.removeEventListener("abort", onAbort); + }; + + signal?.addEventListener("abort", onAbort, { once: true }); + }); +} diff --git a/src/grok/tool-schemas.test.ts b/src/grok/tool-schemas.test.ts new file mode 100644 index 0000000..383adf3 --- /dev/null +++ b/src/grok/tool-schemas.test.ts @@ -0,0 +1,38 @@ +import { tool } from "ai"; +import { describe, expect, it } from "vitest"; +import { z } from "zod"; +import { toolSetToBatchTools } from "./tool-schemas"; + +describe("toolSetToBatchTools", () => { + it("converts AI SDK tools into batch function tools", async () => { + const tools = { + greet: tool({ + description: "Say hello", + inputSchema: z.object({ + name: z.string().describe("Who to greet"), + }), + execute: async () => ({ success: true }), + }), + }; + + const batchTools = await toolSetToBatchTools(tools); + expect(batchTools).toHaveLength(1); + expect(batchTools[0]).toMatchObject({ + type: "function", + function: { + name: "greet", + description: "Say hello", + }, + }); + expect(batchTools[0]?.function.parameters).toMatchObject({ + type: "object", + properties: { + name: { + type: "string", + description: "Who to greet", + }, + }, + required: ["name"], + }); + }); +}); diff --git a/src/grok/tool-schemas.ts b/src/grok/tool-schemas.ts new file mode 100644 index 0000000..2337115 --- /dev/null +++ b/src/grok/tool-schemas.ts @@ -0,0 +1,29 @@ +import { asSchema } from "@ai-sdk/provider-utils"; +import type { ToolSet } from "ai"; +import type { BatchFunctionTool } from "./batch"; + +export async function toolSetToBatchTools(tools: ToolSet): Promise { + const entries = Object.entries(tools); + if (entries.length === 0) { + return []; + } + + const batchTools: BatchFunctionTool[] = []; + for (const [name, tool] of entries) { + if (tool.type === "provider") { + throw new Error(`Batch mode does not support provider-defined tool "${name}".`); + } + + batchTools.push({ + type: "function", + function: { + name, + description: tool.description, + parameters: await asSchema(tool.inputSchema as never).jsonSchema, + ...(tool.strict != null ? { strict: tool.strict } : {}), + }, + }); + } + + return batchTools; +} diff --git a/src/headless/output.ts b/src/headless/output.ts index 3ec56e3..9bb5938 100644 --- a/src/headless/output.ts +++ b/src/headless/output.ts @@ -47,6 +47,7 @@ export type HeadlessJsonEvent = inputTokens?: number; outputTokens?: number; totalTokens?: number; + costUsdTicks?: number; }; } | { diff --git a/src/index.ts b/src/index.ts index 9a93257..991f267 100755 --- a/src/index.ts +++ b/src/index.ts @@ -51,12 +51,13 @@ async function startInteractive( baseURL: string, model: string, maxToolRounds: number, + batchApi: boolean, sandboxMode: SandboxMode, sandboxSettings: SandboxSettings, session?: string, initialMessage?: string, ) { - const agent = new Agent(apiKey, baseURL, model, maxToolRounds, { session, sandboxMode, sandboxSettings }); + const agent = new Agent(apiKey, baseURL, model, maxToolRounds, { session, sandboxMode, sandboxSettings, batchApi }); const { createCliRenderer } = await import("@opentui/core"); const { createRoot } = await import("@opentui/react"); const { createElement } = await import("react"); @@ -100,12 +101,18 @@ async function runHeadless( baseURL: string, model: string, maxToolRounds: number, + batchApi: boolean, sandboxMode: SandboxMode, sandboxSettings: SandboxSettings, format: HeadlessOutputFormat, session?: string, ) { - const agent = new Agent(apiKey, baseURL, model, maxToolRounds, { session, sandboxMode, sandboxSettings }); + const agent = new Agent(apiKey, baseURL, model, maxToolRounds, { + session, + sandboxMode, + sandboxSettings, + batchApi, + }); const prelude = renderHeadlessPrelude(format, agent.getSessionId() || undefined); if (prelude.stdout) process.stdout.write(prelude.stdout); if (prelude.stderr) process.stderr.write(prelude.stderr); @@ -180,6 +187,7 @@ async function runBackgroundDelegation(jobPath: string, options: CliOptions) { persistSession: false, sandboxMode, sandboxSettings, + batchApi: Boolean(delegation.batchApi ?? options.batchApi === true), }); const result = await agent.runTaskRequest({ agent: delegation.agent, @@ -271,6 +279,7 @@ program .option("-s, --session ", "Continue a saved session by id, or use 'latest'") .option("--background-task-file ", "Run a persisted background delegation") .option("--max-tool-rounds ", "Max tool execution rounds", "400") + .option("--batch-api", "Use xAI Batch API for model calls (async, lower cost)") .option("--update", "Update grok to the latest version and exit") .action(async (message: string[], options) => { if (options.update) { @@ -302,6 +311,7 @@ program config.baseURL, config.model, config.maxToolRounds, + options.batchApi === true, config.sandboxMode, config.sandboxSettings, options.format, @@ -317,6 +327,7 @@ program config.baseURL, config.model, config.maxToolRounds, + options.batchApi === true, config.sandboxMode, config.sandboxSettings, options.format, @@ -331,6 +342,7 @@ program config.baseURL, config.model, config.maxToolRounds, + options.batchApi === true, config.sandboxMode, config.sandboxSettings, options.session,