feat: add batch mode for headless Grok CLI runs (#231)
* Add batch API support for unattended headless runs. * Handle empty child batch responses as errors. * remove redundant markdown, bump package version --------- Co-authored-by: Ismail Pelaseyed <homanp@gmail.com>
This commit is contained in:
Alan ZabihiandIsmail Pelaseyed authored and GitHub committed
1 parent
19154d22c6
commit
e2ca4e3621
14 files changed
+1615
-31
No files matched your search
@@ -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.
|
||||
@@ -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
|
||||
|
||||
@@ -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=="],
|
||||
|
||||
+2
-1
@@ -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",
|
||||
|
||||
@@ -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<ToolResult>) | 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<typeof resolveModelRuntime>;
|
||||
childTools: ToolSet;
|
||||
maxSteps: number;
|
||||
initialDetail: string;
|
||||
onActivity?: (detail: string) => void;
|
||||
signal?: AbortSignal;
|
||||
}): Promise<ToolResult> {
|
||||
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<typeof resolveModelRuntime>;
|
||||
modelInfo: ReturnType<typeof getModelInfo>;
|
||||
signal: AbortSignal;
|
||||
}): AsyncGenerator<StreamChunk, void, unknown> {
|
||||
const { userModelMessage, observer, provider, subagents, system, runtime, modelInfo, signal } = args;
|
||||
let attemptedOverflowRecovery = false;
|
||||
|
||||
while (true) {
|
||||
let closeMcp: (() => Promise<void>) | 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,
|
||||
|
||||
@@ -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<typeof import("../grok/batch")>("../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<string, never>;
|
||||
maxSteps: number;
|
||||
initialDetail: string;
|
||||
}) => Promise<unknown>;
|
||||
}
|
||||
).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();
|
||||
});
|
||||
});
|
||||
@@ -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<string, unknown> {
|
||||
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<string, unknown>;
|
||||
}
|
||||
|
||||
async function importDelegationsModule(options: { home?: string; spawnMock?: ReturnType<typeof vi.fn> } = {}) {
|
||||
vi.resetModules();
|
||||
vi.doUnmock("os");
|
||||
vi.doUnmock("child_process");
|
||||
|
||||
if (options.home) {
|
||||
vi.doMock("os", async () => {
|
||||
const actual = await vi.importActual<typeof import("os")>("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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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<typeof fetch>()
|
||||
.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<typeof fetch>()
|
||||
.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);
|
||||
});
|
||||
});
|
||||
@@ -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<BatchInfo> {
|
||||
return requestJson<BatchInfo>({
|
||||
...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<void> {
|
||||
await requestJson<unknown>({
|
||||
...options,
|
||||
method: "POST",
|
||||
path: `/batches/${encodeURIComponent(options.batchId)}/requests`,
|
||||
body: {
|
||||
batch_requests: options.batchRequests,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
export async function getBatchStatus(
|
||||
options: BatchClientOptions & {
|
||||
batchId: string;
|
||||
},
|
||||
): Promise<BatchInfo> {
|
||||
return requestJson<BatchInfo>({
|
||||
...options,
|
||||
method: "GET",
|
||||
path: `/batches/${encodeURIComponent(options.batchId)}`,
|
||||
});
|
||||
}
|
||||
|
||||
export async function listBatchResults(
|
||||
options: BatchClientOptions & {
|
||||
batchId: string;
|
||||
pageSize?: number;
|
||||
paginationToken?: string;
|
||||
},
|
||||
): Promise<BatchResultsPage> {
|
||||
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<BatchResultsPage>({
|
||||
...options,
|
||||
method: "GET",
|
||||
path: `/batches/${encodeURIComponent(options.batchId)}/results?${search.toString()}`,
|
||||
});
|
||||
}
|
||||
|
||||
export async function pollBatchRequestResult(options: PollBatchRequestOptions): Promise<BatchResultEntry> {
|
||||
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<BatchResultEntry | null> {
|
||||
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<T>(
|
||||
options: BatchClientOptions & {
|
||||
method: "GET" | "POST";
|
||||
path: string;
|
||||
body?: unknown;
|
||||
retryRateLimit?: boolean;
|
||||
},
|
||||
): Promise<T> {
|
||||
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<T>(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<T>(response: Response): Promise<T> {
|
||||
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<void> {
|
||||
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 });
|
||||
});
|
||||
}
|
||||
@@ -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"],
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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<BatchFunctionTool[]> {
|
||||
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;
|
||||
}
|
||||
@@ -47,6 +47,7 @@ export type HeadlessJsonEvent =
|
||||
inputTokens?: number;
|
||||
outputTokens?: number;
|
||||
totalTokens?: number;
|
||||
costUsdTicks?: number;
|
||||
};
|
||||
}
|
||||
| {
|
||||
|
||||
+14
-2
@@ -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 <id>", "Continue a saved session by id, or use 'latest'")
|
||||
.option("--background-task-file <path>", "Run a persisted background delegation")
|
||||
.option("--max-tool-rounds <n>", "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,
|
||||
|
||||
Reference in new issue
Block a user