chore(reference): remove pi mono snapshot
This commit is contained in:
ZacharyZhang-NY committed
1 parent
e3eb45b08d
commit
d971ca7929
684 files changed
+7
-192581
No files matched your search
@@ -189,7 +189,8 @@ src/
|
||||
| `scripts/` | Build, release, and packaging Node scripts. |
|
||||
| `landing/` | Marketing site (Next.js). Has its own `landing/public/screenshots/`, separate from `docs/screenshots/`. |
|
||||
| `AgentCompany-Rust/` | Exploratory Rust port; see section below. |
|
||||
| `pi-mono-reference/` | Local reference snapshot only. Shipping runtime source of truth is the pinned `@mariozechner/*` npm dependencies in root `package.json`. |
|
||||
|
||||
Agent runtime source of truth is the pinned `@mariozechner/*` npm dependency set in root `package.json` and `pnpm-lock.yaml`. Third-party reference snapshots live outside this repository.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
> Owner: TBD（建议项目 maintainer）
|
||||
> 关联 commit：`35c1229 Pin pi-agent-core/pi-ai/pi-coding-agent to exact 0.67.68`
|
||||
> Blocking dependencies for full close-out：B5 LLM eval harness
|
||||
> Superseded note（2026-04-22）：in-tree `pi-mono-reference/` snapshot has been removed. Runtime truth now lives in root `package.json` and `pnpm-lock.yaml`.
|
||||
|
||||
---
|
||||
|
||||
@@ -36,7 +37,7 @@ caret range `^0.67.6` 在 pnpm semver 里允许 resolve 到 `>=0.67.6 <0.68.0`
|
||||
|
||||
## 实际执行的决策
|
||||
|
||||
**采纳 Linus / Jobs 立场**：pin 到精确 `0.67.68`（升级 62 patch），同时保持 in-tree reference 不变。
|
||||
**采纳 Linus / Jobs 立场**：pin 到精确 `0.67.68`（升级 62 patch），并记录 runtime package set as the source of truth.
|
||||
|
||||
执行后的事实：
|
||||
- `pnpm install --no-frozen-lockfile` 后 `pnpm typecheck` 全绿
|
||||
@@ -73,7 +74,7 @@ git add package.json pnpm-lock.yaml
|
||||
git commit -m "Revert pi-agent-core pin to 0.67.6 due to behavior regression"
|
||||
```
|
||||
|
||||
如选择回滚，**必须同步**把 `pi-mono-reference/packages/*/package.json` 也降到 `0.67.6`，避免 reference 与运行时再次分叉。
|
||||
如选择回滚，**必须同步**更新本 proposal 的目标版本记录、`package.json`、`pnpm-lock.yaml` 和行为 diff 报告。
|
||||
|
||||
## 不在范围
|
||||
|
||||
|
||||
@@ -59,7 +59,6 @@ AgentCompany/
|
||||
│ ├── package.json
|
||||
│ ├── app/
|
||||
│ └── ...
|
||||
├── pi-mono-reference/ # 保留在 root（现状）
|
||||
└── docs/
|
||||
```
|
||||
|
||||
@@ -155,7 +154,7 @@ packages:
|
||||
- **不引入 turborepo / nx**：pnpm workspace 已足够，加 task runner 是过度工程
|
||||
- **不改代码模块边界**：desktop 内部代码结构本 proposal 不动
|
||||
- **不做跨 package 共享代码**：landing 目前不需要 desktop 的类型；如果未来需要再加 `packages/shared/`
|
||||
- **不动 pi-mono-reference/**：那是第三方 reference 副本，保留在 root
|
||||
- **第三方 runtime reference snapshot 边界**：reference snapshot 已迁出本仓库，monorepo split 保持这个边界
|
||||
- **不改 git branching strategy**：单仓单分支 vs 多分支不是本 proposal 问题
|
||||
|
||||
## 参考文件清单
|
||||
@@ -198,7 +197,6 @@ AgentCompany/
|
||||
├── .gitignore
|
||||
├── CLAUDE.md / AGENTS.md
|
||||
├── docs/
|
||||
├── pi-mono-reference/ # 保留根（第三方 reference）
|
||||
├── scripts/
|
||||
│ ├── desktop/ # electron-builder / signing / release
|
||||
│ └── landing/ # cloudflare deploy / release-links
|
||||
|
||||
@@ -1,21 +0,0 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025 Mario Zechner
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -1,41 +0,0 @@
|
||||
{
|
||||
"$schema": "https://biomejs.dev/schemas/2.3.5/schema.json",
|
||||
"linter": {
|
||||
"enabled": true,
|
||||
"rules": {
|
||||
"recommended": true,
|
||||
"style": {
|
||||
"noNonNullAssertion": "off",
|
||||
"useConst": "error",
|
||||
"useNodejsImportProtocol": "off"
|
||||
},
|
||||
"suspicious": {
|
||||
"noExplicitAny": "off",
|
||||
"noControlCharactersInRegex": "off",
|
||||
"noEmptyInterface": "off"
|
||||
}
|
||||
}
|
||||
},
|
||||
"formatter": {
|
||||
"enabled": true,
|
||||
"formatWithErrors": false,
|
||||
"indentStyle": "tab",
|
||||
"indentWidth": 3,
|
||||
"lineWidth": 120
|
||||
},
|
||||
"files": {
|
||||
"includes": [
|
||||
"packages/*/src/**/*.ts",
|
||||
"packages/*/test/**/*.ts",
|
||||
"packages/coding-agent/examples/**/*.ts",
|
||||
"packages/web-ui/src/**/*.ts",
|
||||
"packages/web-ui/example/**/*.ts",
|
||||
"!**/node_modules/**/*",
|
||||
"!**/test-sessions.ts",
|
||||
"!**/models.generated.ts",
|
||||
"!packages/web-ui/src/app.css",
|
||||
"!packages/mom/data/**/*",
|
||||
"!!**/node_modules"
|
||||
]
|
||||
}
|
||||
}
|
||||
Generated
-8869
File diff suppressed because it is too large.
Load diff
@@ -1,60 +0,0 @@
|
||||
{
|
||||
"name": "pi-monorepo",
|
||||
"private": true,
|
||||
"type": "module",
|
||||
"workspaces": [
|
||||
"packages/*",
|
||||
"packages/web-ui/example",
|
||||
"packages/coding-agent/examples/extensions/with-deps",
|
||||
"packages/coding-agent/examples/extensions/custom-provider-anthropic",
|
||||
"packages/coding-agent/examples/extensions/custom-provider-gitlab-duo",
|
||||
"packages/coding-agent/examples/extensions/custom-provider-qwen-cli"
|
||||
],
|
||||
"scripts": {
|
||||
"clean": "npm run clean --workspaces",
|
||||
"build": "cd packages/tui && npm run build && cd ../ai && npm run build && cd ../agent && npm run build && cd ../coding-agent && npm run build && cd ../mom && npm run build && cd ../web-ui && npm run build && cd ../pods && npm run build",
|
||||
"dev": "concurrently --names \"ai,agent,coding-agent,mom,web-ui,tui\" --prefix-colors \"cyan,yellow,red,white,green,magenta\" \"cd packages/ai && npm run dev\" \"cd packages/agent && npm run dev\" \"cd packages/coding-agent && npm run dev\" \"cd packages/mom && npm run dev\" \"cd packages/web-ui && npm run dev\" \"cd packages/tui && npm run dev\"",
|
||||
"dev:tsc": "concurrently --names \"ai,web-ui\" --prefix-colors \"cyan,green\" \"cd packages/ai && npm run dev:tsc\" \"cd packages/web-ui && npm run dev:tsc\"",
|
||||
"check": "biome check --write --error-on-warnings . && tsgo --noEmit && npm run check:browser-smoke && cd packages/web-ui && npm run check",
|
||||
"check:browser-smoke": "node scripts/check-browser-smoke.mjs",
|
||||
"profile:tui": "node scripts/profile-coding-agent-node.mjs --mode tui",
|
||||
"profile:rpc": "node scripts/profile-coding-agent-node.mjs --mode rpc",
|
||||
"test": "npm run test --workspaces --if-present",
|
||||
"version:patch": "npm version patch -ws --no-git-tag-version && node scripts/sync-versions.js && shx rm -rf node_modules packages/*/node_modules package-lock.json && npm install",
|
||||
"version:minor": "npm version minor -ws --no-git-tag-version && node scripts/sync-versions.js && shx rm -rf node_modules packages/*/node_modules package-lock.json && npm install",
|
||||
"version:major": "npm version major -ws --no-git-tag-version && node scripts/sync-versions.js && shx rm -rf node_modules packages/*/node_modules package-lock.json && npm install",
|
||||
"version:set": "npm version -ws",
|
||||
"prepublishOnly": "npm run clean && npm run build && npm run check",
|
||||
"publish": "npm run prepublishOnly && npm publish -ws --access public",
|
||||
"publish:dry": "npm run prepublishOnly && npm publish -ws --access public --dry-run",
|
||||
"release:patch": "node scripts/release.mjs patch",
|
||||
"release:minor": "node scripts/release.mjs minor",
|
||||
"release:major": "node scripts/release.mjs major",
|
||||
"prepare": "husky"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@biomejs/biome": "2.3.5",
|
||||
"@types/node": "^22.10.5",
|
||||
"@typescript/native-preview": "7.0.0-dev.20260120.1",
|
||||
"concurrently": "^9.2.1",
|
||||
"husky": "^9.1.7",
|
||||
"tsx": "^4.20.3",
|
||||
"typescript": "^5.9.2",
|
||||
"shx": "^0.4.0"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=20.0.0"
|
||||
},
|
||||
"version": "0.0.3",
|
||||
"dependencies": {
|
||||
"@mariozechner/jiti": "^2.6.5",
|
||||
"@mariozechner/pi-coding-agent": "^0.30.2",
|
||||
"get-east-asian-width": "^1.4.0"
|
||||
},
|
||||
"overrides": {
|
||||
"rimraf": "6.1.2",
|
||||
"gaxios": {
|
||||
"rimraf": "6.1.2"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
{
|
||||
"name": "@mariozechner/pi-agent-core",
|
||||
"version": "0.67.68",
|
||||
"description": "General-purpose agent with transport abstraction, state management, and attachment support",
|
||||
"type": "module",
|
||||
"main": "./dist/index.js",
|
||||
"types": "./dist/index.d.ts",
|
||||
"files": [
|
||||
"dist",
|
||||
"README.md"
|
||||
],
|
||||
"scripts": {
|
||||
"clean": "shx rm -rf dist",
|
||||
"build": "tsgo -p tsconfig.build.json",
|
||||
"dev": "tsgo -p tsconfig.build.json --watch --preserveWatchOutput",
|
||||
"test": "vitest --run",
|
||||
"prepublishOnly": "npm run clean && npm run build"
|
||||
},
|
||||
"dependencies": {
|
||||
"@mariozechner/pi-ai": "^0.67.68"
|
||||
},
|
||||
"keywords": [
|
||||
"ai",
|
||||
"agent",
|
||||
"llm",
|
||||
"transport",
|
||||
"state-management"
|
||||
],
|
||||
"author": "Mario Zechner",
|
||||
"license": "MIT",
|
||||
"repository": {
|
||||
"type": "git",
|
||||
"url": "git+https://github.com/badlogic/pi-mono.git",
|
||||
"directory": "packages/agent"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=20.0.0"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^24.3.0",
|
||||
"typescript": "^5.7.3",
|
||||
"vitest": "^3.2.4"
|
||||
}
|
||||
}
|
||||
@@ -1,639 +0,0 @@
|
||||
/**
|
||||
* Agent loop that works with AgentMessage throughout.
|
||||
* Transforms to Message[] only at the LLM call boundary.
|
||||
*/
|
||||
|
||||
import {
|
||||
type AssistantMessage,
|
||||
type Context,
|
||||
EventStream,
|
||||
streamSimple,
|
||||
type ToolResultMessage,
|
||||
validateToolArguments,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import type {
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentLoopConfig,
|
||||
AgentMessage,
|
||||
AgentTool,
|
||||
AgentToolCall,
|
||||
AgentToolResult,
|
||||
StreamFn,
|
||||
} from "./types.js";
|
||||
|
||||
export type AgentEventSink = (event: AgentEvent) => Promise<void> | void;
|
||||
|
||||
/**
|
||||
* Start an agent loop with a new prompt message.
|
||||
* The prompt is added to the context and events are emitted for it.
|
||||
*/
|
||||
export function agentLoop(
|
||||
prompts: AgentMessage[],
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
signal?: AbortSignal,
|
||||
streamFn?: StreamFn,
|
||||
): EventStream<AgentEvent, AgentMessage[]> {
|
||||
const stream = createAgentStream();
|
||||
|
||||
void runAgentLoop(
|
||||
prompts,
|
||||
context,
|
||||
config,
|
||||
async (event) => {
|
||||
stream.push(event);
|
||||
},
|
||||
signal,
|
||||
streamFn,
|
||||
).then((messages) => {
|
||||
stream.end(messages);
|
||||
});
|
||||
|
||||
return stream;
|
||||
}
|
||||
|
||||
/**
|
||||
* Continue an agent loop from the current context without adding a new message.
|
||||
* Used for retries - context already has user message or tool results.
|
||||
*
|
||||
* **Important:** The last message in context must convert to a `user` or `toolResult` message
|
||||
* via `convertToLlm`. If it doesn't, the LLM provider will reject the request.
|
||||
* This cannot be validated here since `convertToLlm` is only called once per turn.
|
||||
*/
|
||||
export function agentLoopContinue(
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
signal?: AbortSignal,
|
||||
streamFn?: StreamFn,
|
||||
): EventStream<AgentEvent, AgentMessage[]> {
|
||||
if (context.messages.length === 0) {
|
||||
throw new Error("Cannot continue: no messages in context");
|
||||
}
|
||||
|
||||
if (context.messages[context.messages.length - 1].role === "assistant") {
|
||||
throw new Error("Cannot continue from message role: assistant");
|
||||
}
|
||||
|
||||
const stream = createAgentStream();
|
||||
|
||||
void runAgentLoopContinue(
|
||||
context,
|
||||
config,
|
||||
async (event) => {
|
||||
stream.push(event);
|
||||
},
|
||||
signal,
|
||||
streamFn,
|
||||
).then((messages) => {
|
||||
stream.end(messages);
|
||||
});
|
||||
|
||||
return stream;
|
||||
}
|
||||
|
||||
export async function runAgentLoop(
|
||||
prompts: AgentMessage[],
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
emit: AgentEventSink,
|
||||
signal?: AbortSignal,
|
||||
streamFn?: StreamFn,
|
||||
): Promise<AgentMessage[]> {
|
||||
const newMessages: AgentMessage[] = [...prompts];
|
||||
const currentContext: AgentContext = {
|
||||
...context,
|
||||
messages: [...context.messages, ...prompts],
|
||||
};
|
||||
|
||||
await emit({ type: "agent_start" });
|
||||
await emit({ type: "turn_start" });
|
||||
for (const prompt of prompts) {
|
||||
await emit({ type: "message_start", message: prompt });
|
||||
await emit({ type: "message_end", message: prompt });
|
||||
}
|
||||
|
||||
await runLoop(currentContext, newMessages, config, signal, emit, streamFn);
|
||||
return newMessages;
|
||||
}
|
||||
|
||||
export async function runAgentLoopContinue(
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
emit: AgentEventSink,
|
||||
signal?: AbortSignal,
|
||||
streamFn?: StreamFn,
|
||||
): Promise<AgentMessage[]> {
|
||||
if (context.messages.length === 0) {
|
||||
throw new Error("Cannot continue: no messages in context");
|
||||
}
|
||||
|
||||
if (context.messages[context.messages.length - 1].role === "assistant") {
|
||||
throw new Error("Cannot continue from message role: assistant");
|
||||
}
|
||||
|
||||
const newMessages: AgentMessage[] = [];
|
||||
const currentContext: AgentContext = { ...context };
|
||||
|
||||
await emit({ type: "agent_start" });
|
||||
await emit({ type: "turn_start" });
|
||||
|
||||
await runLoop(currentContext, newMessages, config, signal, emit, streamFn);
|
||||
return newMessages;
|
||||
}
|
||||
|
||||
function createAgentStream(): EventStream<AgentEvent, AgentMessage[]> {
|
||||
return new EventStream<AgentEvent, AgentMessage[]>(
|
||||
(event: AgentEvent) => event.type === "agent_end",
|
||||
(event: AgentEvent) => (event.type === "agent_end" ? event.messages : []),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Main loop logic shared by agentLoop and agentLoopContinue.
|
||||
*/
|
||||
async function runLoop(
|
||||
currentContext: AgentContext,
|
||||
newMessages: AgentMessage[],
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
streamFn?: StreamFn,
|
||||
): Promise<void> {
|
||||
let firstTurn = true;
|
||||
// Check for steering messages at start (user may have typed while waiting)
|
||||
let pendingMessages: AgentMessage[] = (await config.getSteeringMessages?.()) || [];
|
||||
|
||||
// Outer loop: continues when queued follow-up messages arrive after agent would stop
|
||||
while (true) {
|
||||
let hasMoreToolCalls = true;
|
||||
|
||||
// Inner loop: process tool calls and steering messages
|
||||
while (hasMoreToolCalls || pendingMessages.length > 0) {
|
||||
if (!firstTurn) {
|
||||
await emit({ type: "turn_start" });
|
||||
} else {
|
||||
firstTurn = false;
|
||||
}
|
||||
|
||||
// Process pending messages (inject before next assistant response)
|
||||
if (pendingMessages.length > 0) {
|
||||
for (const message of pendingMessages) {
|
||||
await emit({ type: "message_start", message });
|
||||
await emit({ type: "message_end", message });
|
||||
currentContext.messages.push(message);
|
||||
newMessages.push(message);
|
||||
}
|
||||
pendingMessages = [];
|
||||
}
|
||||
|
||||
// Stream assistant response
|
||||
const message = await streamAssistantResponse(currentContext, config, signal, emit, streamFn);
|
||||
newMessages.push(message);
|
||||
|
||||
if (message.stopReason === "error" || message.stopReason === "aborted") {
|
||||
await emit({ type: "turn_end", message, toolResults: [] });
|
||||
await emit({ type: "agent_end", messages: newMessages });
|
||||
return;
|
||||
}
|
||||
|
||||
// Check for tool calls
|
||||
const toolCalls = message.content.filter((c) => c.type === "toolCall");
|
||||
hasMoreToolCalls = toolCalls.length > 0;
|
||||
|
||||
const toolResults: ToolResultMessage[] = [];
|
||||
if (hasMoreToolCalls) {
|
||||
toolResults.push(...(await executeToolCalls(currentContext, message, config, signal, emit)));
|
||||
|
||||
for (const result of toolResults) {
|
||||
currentContext.messages.push(result);
|
||||
newMessages.push(result);
|
||||
}
|
||||
}
|
||||
|
||||
await emit({ type: "turn_end", message, toolResults });
|
||||
|
||||
pendingMessages = (await config.getSteeringMessages?.()) || [];
|
||||
}
|
||||
|
||||
// Agent would stop here. Check for follow-up messages.
|
||||
const followUpMessages = (await config.getFollowUpMessages?.()) || [];
|
||||
if (followUpMessages.length > 0) {
|
||||
// Set as pending so inner loop processes them
|
||||
pendingMessages = followUpMessages;
|
||||
continue;
|
||||
}
|
||||
|
||||
// No more messages, exit
|
||||
break;
|
||||
}
|
||||
|
||||
await emit({ type: "agent_end", messages: newMessages });
|
||||
}
|
||||
|
||||
/**
|
||||
* Stream an assistant response from the LLM.
|
||||
* This is where AgentMessage[] gets transformed to Message[] for the LLM.
|
||||
*/
|
||||
async function streamAssistantResponse(
|
||||
context: AgentContext,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
streamFn?: StreamFn,
|
||||
): Promise<AssistantMessage> {
|
||||
// Apply context transform if configured (AgentMessage[] → AgentMessage[])
|
||||
let messages = context.messages;
|
||||
if (config.transformContext) {
|
||||
messages = await config.transformContext(messages, signal);
|
||||
}
|
||||
|
||||
// Convert to LLM-compatible messages (AgentMessage[] → Message[])
|
||||
const llmMessages = await config.convertToLlm(messages);
|
||||
|
||||
// Build LLM context
|
||||
const llmContext: Context = {
|
||||
systemPrompt: context.systemPrompt,
|
||||
messages: llmMessages,
|
||||
tools: context.tools,
|
||||
};
|
||||
|
||||
const streamFunction = streamFn || streamSimple;
|
||||
|
||||
// Resolve API key (important for expiring tokens)
|
||||
const resolvedApiKey =
|
||||
(config.getApiKey ? await config.getApiKey(config.model.provider) : undefined) || config.apiKey;
|
||||
|
||||
const response = await streamFunction(config.model, llmContext, {
|
||||
...config,
|
||||
apiKey: resolvedApiKey,
|
||||
signal,
|
||||
});
|
||||
|
||||
let partialMessage: AssistantMessage | null = null;
|
||||
let addedPartial = false;
|
||||
|
||||
for await (const event of response) {
|
||||
switch (event.type) {
|
||||
case "start":
|
||||
partialMessage = event.partial;
|
||||
context.messages.push(partialMessage);
|
||||
addedPartial = true;
|
||||
await emit({ type: "message_start", message: { ...partialMessage } });
|
||||
break;
|
||||
|
||||
case "text_start":
|
||||
case "text_delta":
|
||||
case "text_end":
|
||||
case "thinking_start":
|
||||
case "thinking_delta":
|
||||
case "thinking_end":
|
||||
case "toolcall_start":
|
||||
case "toolcall_delta":
|
||||
case "toolcall_end":
|
||||
if (partialMessage) {
|
||||
partialMessage = event.partial;
|
||||
context.messages[context.messages.length - 1] = partialMessage;
|
||||
await emit({
|
||||
type: "message_update",
|
||||
assistantMessageEvent: event,
|
||||
message: { ...partialMessage },
|
||||
});
|
||||
}
|
||||
break;
|
||||
|
||||
case "done":
|
||||
case "error": {
|
||||
const finalMessage = await response.result();
|
||||
if (addedPartial) {
|
||||
context.messages[context.messages.length - 1] = finalMessage;
|
||||
} else {
|
||||
context.messages.push(finalMessage);
|
||||
}
|
||||
if (!addedPartial) {
|
||||
await emit({ type: "message_start", message: { ...finalMessage } });
|
||||
}
|
||||
await emit({ type: "message_end", message: finalMessage });
|
||||
return finalMessage;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const finalMessage = await response.result();
|
||||
if (addedPartial) {
|
||||
context.messages[context.messages.length - 1] = finalMessage;
|
||||
} else {
|
||||
context.messages.push(finalMessage);
|
||||
await emit({ type: "message_start", message: { ...finalMessage } });
|
||||
}
|
||||
await emit({ type: "message_end", message: finalMessage });
|
||||
return finalMessage;
|
||||
}
|
||||
|
||||
/**
|
||||
* Execute tool calls from an assistant message.
|
||||
*/
|
||||
async function executeToolCalls(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ToolResultMessage[]> {
|
||||
const toolCalls = assistantMessage.content.filter((c) => c.type === "toolCall");
|
||||
const hasSequentialToolCall = toolCalls.some(
|
||||
(tc) => currentContext.tools?.find((t) => t.name === tc.name)?.executionMode === "sequential",
|
||||
);
|
||||
if (config.toolExecution === "sequential" || hasSequentialToolCall) {
|
||||
return executeToolCallsSequential(currentContext, assistantMessage, toolCalls, config, signal, emit);
|
||||
}
|
||||
return executeToolCallsParallel(currentContext, assistantMessage, toolCalls, config, signal, emit);
|
||||
}
|
||||
|
||||
async function executeToolCallsSequential(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
toolCalls: AgentToolCall[],
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ToolResultMessage[]> {
|
||||
const results: ToolResultMessage[] = [];
|
||||
|
||||
for (const toolCall of toolCalls) {
|
||||
await emit({
|
||||
type: "tool_execution_start",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args: toolCall.arguments,
|
||||
});
|
||||
|
||||
const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal);
|
||||
if (preparation.kind === "immediate") {
|
||||
results.push(await emitToolCallOutcome(toolCall, preparation.result, preparation.isError, emit));
|
||||
} else {
|
||||
const executed = await executePreparedToolCall(preparation, signal, emit);
|
||||
results.push(
|
||||
await finalizeExecutedToolCall(
|
||||
currentContext,
|
||||
assistantMessage,
|
||||
preparation,
|
||||
executed,
|
||||
config,
|
||||
signal,
|
||||
emit,
|
||||
),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
async function executeToolCallsParallel(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
toolCalls: AgentToolCall[],
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ToolResultMessage[]> {
|
||||
const results: ToolResultMessage[] = [];
|
||||
const runnableCalls: PreparedToolCall[] = [];
|
||||
|
||||
for (const toolCall of toolCalls) {
|
||||
await emit({
|
||||
type: "tool_execution_start",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
args: toolCall.arguments,
|
||||
});
|
||||
|
||||
const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal);
|
||||
if (preparation.kind === "immediate") {
|
||||
results.push(await emitToolCallOutcome(toolCall, preparation.result, preparation.isError, emit));
|
||||
} else {
|
||||
runnableCalls.push(preparation);
|
||||
}
|
||||
}
|
||||
|
||||
const runningCalls = runnableCalls.map((prepared) => ({
|
||||
prepared,
|
||||
execution: executePreparedToolCall(prepared, signal, emit),
|
||||
}));
|
||||
|
||||
for (const running of runningCalls) {
|
||||
const executed = await running.execution;
|
||||
results.push(
|
||||
await finalizeExecutedToolCall(
|
||||
currentContext,
|
||||
assistantMessage,
|
||||
running.prepared,
|
||||
executed,
|
||||
config,
|
||||
signal,
|
||||
emit,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
type PreparedToolCall = {
|
||||
kind: "prepared";
|
||||
toolCall: AgentToolCall;
|
||||
tool: AgentTool<any>;
|
||||
args: unknown;
|
||||
};
|
||||
|
||||
type ImmediateToolCallOutcome = {
|
||||
kind: "immediate";
|
||||
result: AgentToolResult<any>;
|
||||
isError: boolean;
|
||||
};
|
||||
|
||||
type ExecutedToolCallOutcome = {
|
||||
result: AgentToolResult<any>;
|
||||
isError: boolean;
|
||||
};
|
||||
|
||||
function prepareToolCallArguments(tool: AgentTool<any>, toolCall: AgentToolCall): AgentToolCall {
|
||||
if (!tool.prepareArguments) {
|
||||
return toolCall;
|
||||
}
|
||||
const preparedArguments = tool.prepareArguments(toolCall.arguments);
|
||||
if (preparedArguments === toolCall.arguments) {
|
||||
return toolCall;
|
||||
}
|
||||
return {
|
||||
...toolCall,
|
||||
arguments: preparedArguments as Record<string, any>,
|
||||
};
|
||||
}
|
||||
|
||||
async function prepareToolCall(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
toolCall: AgentToolCall,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
): Promise<PreparedToolCall | ImmediateToolCallOutcome> {
|
||||
const tool = currentContext.tools?.find((t) => t.name === toolCall.name);
|
||||
if (!tool) {
|
||||
return {
|
||||
kind: "immediate",
|
||||
result: createErrorToolResult(`Tool ${toolCall.name} not found`),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
|
||||
try {
|
||||
const preparedToolCall = prepareToolCallArguments(tool, toolCall);
|
||||
const validatedArgs = validateToolArguments(tool, preparedToolCall);
|
||||
if (config.beforeToolCall) {
|
||||
const beforeResult = await config.beforeToolCall(
|
||||
{
|
||||
assistantMessage,
|
||||
toolCall,
|
||||
args: validatedArgs,
|
||||
context: currentContext,
|
||||
},
|
||||
signal,
|
||||
);
|
||||
if (beforeResult?.block) {
|
||||
return {
|
||||
kind: "immediate",
|
||||
result: createErrorToolResult(beforeResult.reason || "Tool execution was blocked"),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
return {
|
||||
kind: "prepared",
|
||||
toolCall,
|
||||
tool,
|
||||
args: validatedArgs,
|
||||
};
|
||||
} catch (error) {
|
||||
return {
|
||||
kind: "immediate",
|
||||
result: createErrorToolResult(error instanceof Error ? error.message : String(error)),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
async function executePreparedToolCall(
|
||||
prepared: PreparedToolCall,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ExecutedToolCallOutcome> {
|
||||
const updateEvents: Promise<void>[] = [];
|
||||
|
||||
try {
|
||||
const result = await prepared.tool.execute(
|
||||
prepared.toolCall.id,
|
||||
prepared.args as never,
|
||||
signal,
|
||||
(partialResult) => {
|
||||
updateEvents.push(
|
||||
Promise.resolve(
|
||||
emit({
|
||||
type: "tool_execution_update",
|
||||
toolCallId: prepared.toolCall.id,
|
||||
toolName: prepared.toolCall.name,
|
||||
args: prepared.toolCall.arguments,
|
||||
partialResult,
|
||||
}),
|
||||
),
|
||||
);
|
||||
},
|
||||
);
|
||||
await Promise.all(updateEvents);
|
||||
return { result, isError: false };
|
||||
} catch (error) {
|
||||
await Promise.all(updateEvents);
|
||||
return {
|
||||
result: createErrorToolResult(error instanceof Error ? error.message : String(error)),
|
||||
isError: true,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
async function finalizeExecutedToolCall(
|
||||
currentContext: AgentContext,
|
||||
assistantMessage: AssistantMessage,
|
||||
prepared: PreparedToolCall,
|
||||
executed: ExecutedToolCallOutcome,
|
||||
config: AgentLoopConfig,
|
||||
signal: AbortSignal | undefined,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ToolResultMessage> {
|
||||
let result = executed.result;
|
||||
let isError = executed.isError;
|
||||
|
||||
if (config.afterToolCall) {
|
||||
try {
|
||||
const afterResult = await config.afterToolCall(
|
||||
{
|
||||
assistantMessage,
|
||||
toolCall: prepared.toolCall,
|
||||
args: prepared.args,
|
||||
result,
|
||||
isError,
|
||||
context: currentContext,
|
||||
},
|
||||
signal,
|
||||
);
|
||||
if (afterResult) {
|
||||
result = {
|
||||
content: afterResult.content ?? result.content,
|
||||
details: afterResult.details ?? result.details,
|
||||
};
|
||||
isError = afterResult.isError ?? isError;
|
||||
}
|
||||
} catch (error) {
|
||||
result = createErrorToolResult(error instanceof Error ? error.message : String(error));
|
||||
isError = true;
|
||||
}
|
||||
}
|
||||
|
||||
return await emitToolCallOutcome(prepared.toolCall, result, isError, emit);
|
||||
}
|
||||
|
||||
function createErrorToolResult(message: string): AgentToolResult<any> {
|
||||
return {
|
||||
content: [{ type: "text", text: message }],
|
||||
details: {},
|
||||
};
|
||||
}
|
||||
|
||||
async function emitToolCallOutcome(
|
||||
toolCall: AgentToolCall,
|
||||
result: AgentToolResult<any>,
|
||||
isError: boolean,
|
||||
emit: AgentEventSink,
|
||||
): Promise<ToolResultMessage> {
|
||||
await emit({
|
||||
type: "tool_execution_end",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
result,
|
||||
isError,
|
||||
});
|
||||
|
||||
const toolResultMessage: ToolResultMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
content: result.content,
|
||||
details: result.details,
|
||||
isError,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
await emit({ type: "message_start", message: toolResultMessage });
|
||||
await emit({ type: "message_end", message: toolResultMessage });
|
||||
return toolResultMessage;
|
||||
}
|
||||
@@ -1,543 +0,0 @@
|
||||
import {
|
||||
type ImageContent,
|
||||
type Message,
|
||||
type Model,
|
||||
type SimpleStreamOptions,
|
||||
streamSimple,
|
||||
type TextContent,
|
||||
type ThinkingBudgets,
|
||||
type Transport,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.js";
|
||||
import type {
|
||||
AfterToolCallContext,
|
||||
AfterToolCallResult,
|
||||
AgentContext,
|
||||
AgentEvent,
|
||||
AgentLoopConfig,
|
||||
AgentMessage,
|
||||
AgentState,
|
||||
AgentTool,
|
||||
BeforeToolCallContext,
|
||||
BeforeToolCallResult,
|
||||
StreamFn,
|
||||
ToolExecutionMode,
|
||||
} from "./types.js";
|
||||
|
||||
function defaultConvertToLlm(messages: AgentMessage[]): Message[] {
|
||||
return messages.filter(
|
||||
(message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult",
|
||||
);
|
||||
}
|
||||
|
||||
const EMPTY_USAGE = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
|
||||
const DEFAULT_MODEL = {
|
||||
id: "unknown",
|
||||
name: "unknown",
|
||||
api: "unknown",
|
||||
provider: "unknown",
|
||||
baseUrl: "",
|
||||
reasoning: false,
|
||||
input: [],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 0,
|
||||
maxTokens: 0,
|
||||
} satisfies Model<any>;
|
||||
|
||||
type QueueMode = "all" | "one-at-a-time";
|
||||
|
||||
type MutableAgentState = Omit<AgentState, "isStreaming" | "streamingMessage" | "pendingToolCalls" | "errorMessage"> & {
|
||||
isStreaming: boolean;
|
||||
streamingMessage?: AgentMessage;
|
||||
pendingToolCalls: Set<string>;
|
||||
errorMessage?: string;
|
||||
};
|
||||
|
||||
function createMutableAgentState(
|
||||
initialState?: Partial<Omit<AgentState, "pendingToolCalls" | "isStreaming" | "streamingMessage" | "errorMessage">>,
|
||||
): MutableAgentState {
|
||||
let tools = initialState?.tools?.slice() ?? [];
|
||||
let messages = initialState?.messages?.slice() ?? [];
|
||||
|
||||
return {
|
||||
systemPrompt: initialState?.systemPrompt ?? "",
|
||||
model: initialState?.model ?? DEFAULT_MODEL,
|
||||
thinkingLevel: initialState?.thinkingLevel ?? "off",
|
||||
get tools() {
|
||||
return tools;
|
||||
},
|
||||
set tools(nextTools: AgentTool<any>[]) {
|
||||
tools = nextTools.slice();
|
||||
},
|
||||
get messages() {
|
||||
return messages;
|
||||
},
|
||||
set messages(nextMessages: AgentMessage[]) {
|
||||
messages = nextMessages.slice();
|
||||
},
|
||||
isStreaming: false,
|
||||
streamingMessage: undefined,
|
||||
pendingToolCalls: new Set<string>(),
|
||||
errorMessage: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
/** Options for constructing an {@link Agent}. */
|
||||
export interface AgentOptions {
|
||||
initialState?: Partial<Omit<AgentState, "pendingToolCalls" | "isStreaming" | "streamingMessage" | "errorMessage">>;
|
||||
convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
|
||||
transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
|
||||
streamFn?: StreamFn;
|
||||
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
||||
onPayload?: SimpleStreamOptions["onPayload"];
|
||||
onResponse?: SimpleStreamOptions["onResponse"];
|
||||
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
|
||||
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
|
||||
steeringMode?: QueueMode;
|
||||
followUpMode?: QueueMode;
|
||||
sessionId?: string;
|
||||
thinkingBudgets?: ThinkingBudgets;
|
||||
transport?: Transport;
|
||||
maxRetryDelayMs?: number;
|
||||
toolExecution?: ToolExecutionMode;
|
||||
}
|
||||
|
||||
class PendingMessageQueue {
|
||||
private messages: AgentMessage[] = [];
|
||||
|
||||
constructor(public mode: QueueMode) {}
|
||||
|
||||
enqueue(message: AgentMessage): void {
|
||||
this.messages.push(message);
|
||||
}
|
||||
|
||||
hasItems(): boolean {
|
||||
return this.messages.length > 0;
|
||||
}
|
||||
|
||||
drain(): AgentMessage[] {
|
||||
if (this.mode === "all") {
|
||||
const drained = this.messages.slice();
|
||||
this.messages = [];
|
||||
return drained;
|
||||
}
|
||||
|
||||
const first = this.messages[0];
|
||||
if (!first) {
|
||||
return [];
|
||||
}
|
||||
this.messages = this.messages.slice(1);
|
||||
return [first];
|
||||
}
|
||||
|
||||
clear(): void {
|
||||
this.messages = [];
|
||||
}
|
||||
}
|
||||
|
||||
type ActiveRun = {
|
||||
promise: Promise<void>;
|
||||
resolve: () => void;
|
||||
abortController: AbortController;
|
||||
};
|
||||
|
||||
/**
|
||||
* Stateful wrapper around the low-level agent loop.
|
||||
*
|
||||
* `Agent` owns the current transcript, emits lifecycle events, executes tools,
|
||||
* and exposes queueing APIs for steering and follow-up messages.
|
||||
*/
|
||||
export class Agent {
|
||||
private _state: MutableAgentState;
|
||||
private readonly listeners = new Set<(event: AgentEvent, signal: AbortSignal) => Promise<void> | void>();
|
||||
private readonly steeringQueue: PendingMessageQueue;
|
||||
private readonly followUpQueue: PendingMessageQueue;
|
||||
|
||||
public convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
|
||||
public transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
|
||||
public streamFn: StreamFn;
|
||||
public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
||||
public onPayload?: SimpleStreamOptions["onPayload"];
|
||||
public onResponse?: SimpleStreamOptions["onResponse"];
|
||||
public beforeToolCall?: (
|
||||
context: BeforeToolCallContext,
|
||||
signal?: AbortSignal,
|
||||
) => Promise<BeforeToolCallResult | undefined>;
|
||||
public afterToolCall?: (
|
||||
context: AfterToolCallContext,
|
||||
signal?: AbortSignal,
|
||||
) => Promise<AfterToolCallResult | undefined>;
|
||||
private activeRun?: ActiveRun;
|
||||
/** Session identifier forwarded to providers for cache-aware backends. */
|
||||
public sessionId?: string;
|
||||
/** Optional per-level thinking token budgets forwarded to the stream function. */
|
||||
public thinkingBudgets?: ThinkingBudgets;
|
||||
/** Preferred transport forwarded to the stream function. */
|
||||
public transport: Transport;
|
||||
/** Optional cap for provider-requested retry delays. */
|
||||
public maxRetryDelayMs?: number;
|
||||
/** Tool execution strategy for assistant messages that contain multiple tool calls. */
|
||||
public toolExecution: ToolExecutionMode;
|
||||
|
||||
constructor(options: AgentOptions = {}) {
|
||||
this._state = createMutableAgentState(options.initialState);
|
||||
this.convertToLlm = options.convertToLlm ?? defaultConvertToLlm;
|
||||
this.transformContext = options.transformContext;
|
||||
this.streamFn = options.streamFn ?? streamSimple;
|
||||
this.getApiKey = options.getApiKey;
|
||||
this.onPayload = options.onPayload;
|
||||
this.onResponse = options.onResponse;
|
||||
this.beforeToolCall = options.beforeToolCall;
|
||||
this.afterToolCall = options.afterToolCall;
|
||||
this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time");
|
||||
this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time");
|
||||
this.sessionId = options.sessionId;
|
||||
this.thinkingBudgets = options.thinkingBudgets;
|
||||
this.transport = options.transport ?? "sse";
|
||||
this.maxRetryDelayMs = options.maxRetryDelayMs;
|
||||
this.toolExecution = options.toolExecution ?? "parallel";
|
||||
}
|
||||
|
||||
/**
|
||||
* Subscribe to agent lifecycle events.
|
||||
*
|
||||
* Listener promises are awaited in subscription order and are included in
|
||||
* the current run's settlement. Listeners also receive the active abort
|
||||
* signal for the current run.
|
||||
*
|
||||
* `agent_end` is the final emitted event for a run, but the agent does not
|
||||
* become idle until all awaited listeners for that event have settled.
|
||||
*/
|
||||
subscribe(listener: (event: AgentEvent, signal: AbortSignal) => Promise<void> | void): () => void {
|
||||
this.listeners.add(listener);
|
||||
return () => this.listeners.delete(listener);
|
||||
}
|
||||
|
||||
/**
|
||||
* Current agent state.
|
||||
*
|
||||
* Assigning `state.tools` or `state.messages` copies the provided top-level array.
|
||||
*/
|
||||
get state(): AgentState {
|
||||
return this._state;
|
||||
}
|
||||
|
||||
/** Controls how queued steering messages are drained. */
|
||||
set steeringMode(mode: QueueMode) {
|
||||
this.steeringQueue.mode = mode;
|
||||
}
|
||||
|
||||
get steeringMode(): QueueMode {
|
||||
return this.steeringQueue.mode;
|
||||
}
|
||||
|
||||
/** Controls how queued follow-up messages are drained. */
|
||||
set followUpMode(mode: QueueMode) {
|
||||
this.followUpQueue.mode = mode;
|
||||
}
|
||||
|
||||
get followUpMode(): QueueMode {
|
||||
return this.followUpQueue.mode;
|
||||
}
|
||||
|
||||
/** Queue a message to be injected after the current assistant turn finishes. */
|
||||
steer(message: AgentMessage): void {
|
||||
this.steeringQueue.enqueue(message);
|
||||
}
|
||||
|
||||
/** Queue a message to run only after the agent would otherwise stop. */
|
||||
followUp(message: AgentMessage): void {
|
||||
this.followUpQueue.enqueue(message);
|
||||
}
|
||||
|
||||
/** Remove all queued steering messages. */
|
||||
clearSteeringQueue(): void {
|
||||
this.steeringQueue.clear();
|
||||
}
|
||||
|
||||
/** Remove all queued follow-up messages. */
|
||||
clearFollowUpQueue(): void {
|
||||
this.followUpQueue.clear();
|
||||
}
|
||||
|
||||
/** Remove all queued steering and follow-up messages. */
|
||||
clearAllQueues(): void {
|
||||
this.clearSteeringQueue();
|
||||
this.clearFollowUpQueue();
|
||||
}
|
||||
|
||||
/** Returns true when either queue still contains pending messages. */
|
||||
hasQueuedMessages(): boolean {
|
||||
return this.steeringQueue.hasItems() || this.followUpQueue.hasItems();
|
||||
}
|
||||
|
||||
/** Active abort signal for the current run, if any. */
|
||||
get signal(): AbortSignal | undefined {
|
||||
return this.activeRun?.abortController.signal;
|
||||
}
|
||||
|
||||
/** Abort the current run, if one is active. */
|
||||
abort(): void {
|
||||
this.activeRun?.abortController.abort();
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve when the current run and all awaited event listeners have finished.
|
||||
*
|
||||
* This resolves after `agent_end` listeners settle.
|
||||
*/
|
||||
waitForIdle(): Promise<void> {
|
||||
return this.activeRun?.promise ?? Promise.resolve();
|
||||
}
|
||||
|
||||
/** Clear transcript state, runtime state, and queued messages. */
|
||||
reset(): void {
|
||||
this._state.messages = [];
|
||||
this._state.isStreaming = false;
|
||||
this._state.streamingMessage = undefined;
|
||||
this._state.pendingToolCalls = new Set<string>();
|
||||
this._state.errorMessage = undefined;
|
||||
this.clearFollowUpQueue();
|
||||
this.clearSteeringQueue();
|
||||
}
|
||||
|
||||
/** Start a new prompt from text, a single message, or a batch of messages. */
|
||||
async prompt(message: AgentMessage | AgentMessage[]): Promise<void>;
|
||||
async prompt(input: string, images?: ImageContent[]): Promise<void>;
|
||||
async prompt(input: string | AgentMessage | AgentMessage[], images?: ImageContent[]): Promise<void> {
|
||||
if (this.activeRun) {
|
||||
throw new Error(
|
||||
"Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.",
|
||||
);
|
||||
}
|
||||
const messages = this.normalizePromptInput(input, images);
|
||||
await this.runPromptMessages(messages);
|
||||
}
|
||||
|
||||
/** Continue from the current transcript. The last message must be a user or tool-result message. */
|
||||
async continue(): Promise<void> {
|
||||
if (this.activeRun) {
|
||||
throw new Error("Agent is already processing. Wait for completion before continuing.");
|
||||
}
|
||||
|
||||
const lastMessage = this._state.messages[this._state.messages.length - 1];
|
||||
if (!lastMessage) {
|
||||
throw new Error("No messages to continue from");
|
||||
}
|
||||
|
||||
if (lastMessage.role === "assistant") {
|
||||
const queuedSteering = this.steeringQueue.drain();
|
||||
if (queuedSteering.length > 0) {
|
||||
await this.runPromptMessages(queuedSteering, { skipInitialSteeringPoll: true });
|
||||
return;
|
||||
}
|
||||
|
||||
const queuedFollowUps = this.followUpQueue.drain();
|
||||
if (queuedFollowUps.length > 0) {
|
||||
await this.runPromptMessages(queuedFollowUps);
|
||||
return;
|
||||
}
|
||||
|
||||
throw new Error("Cannot continue from message role: assistant");
|
||||
}
|
||||
|
||||
await this.runContinuation();
|
||||
}
|
||||
|
||||
private normalizePromptInput(
|
||||
input: string | AgentMessage | AgentMessage[],
|
||||
images?: ImageContent[],
|
||||
): AgentMessage[] {
|
||||
if (Array.isArray(input)) {
|
||||
return input;
|
||||
}
|
||||
|
||||
if (typeof input !== "string") {
|
||||
return [input];
|
||||
}
|
||||
|
||||
const content: Array<TextContent | ImageContent> = [{ type: "text", text: input }];
|
||||
if (images && images.length > 0) {
|
||||
content.push(...images);
|
||||
}
|
||||
return [{ role: "user", content, timestamp: Date.now() }];
|
||||
}
|
||||
|
||||
private async runPromptMessages(
|
||||
messages: AgentMessage[],
|
||||
options: { skipInitialSteeringPoll?: boolean } = {},
|
||||
): Promise<void> {
|
||||
await this.runWithLifecycle(async (signal) => {
|
||||
await runAgentLoop(
|
||||
messages,
|
||||
this.createContextSnapshot(),
|
||||
this.createLoopConfig(options),
|
||||
(event) => this.processEvents(event),
|
||||
signal,
|
||||
this.streamFn,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
private async runContinuation(): Promise<void> {
|
||||
await this.runWithLifecycle(async (signal) => {
|
||||
await runAgentLoopContinue(
|
||||
this.createContextSnapshot(),
|
||||
this.createLoopConfig(),
|
||||
(event) => this.processEvents(event),
|
||||
signal,
|
||||
this.streamFn,
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
private createContextSnapshot(): AgentContext {
|
||||
return {
|
||||
systemPrompt: this._state.systemPrompt,
|
||||
messages: this._state.messages.slice(),
|
||||
tools: this._state.tools.slice(),
|
||||
};
|
||||
}
|
||||
|
||||
private createLoopConfig(options: { skipInitialSteeringPoll?: boolean } = {}): AgentLoopConfig {
|
||||
let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true;
|
||||
return {
|
||||
model: this._state.model,
|
||||
reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel,
|
||||
sessionId: this.sessionId,
|
||||
onPayload: this.onPayload,
|
||||
onResponse: this.onResponse,
|
||||
transport: this.transport,
|
||||
thinkingBudgets: this.thinkingBudgets,
|
||||
maxRetryDelayMs: this.maxRetryDelayMs,
|
||||
toolExecution: this.toolExecution,
|
||||
beforeToolCall: this.beforeToolCall,
|
||||
afterToolCall: this.afterToolCall,
|
||||
convertToLlm: this.convertToLlm,
|
||||
transformContext: this.transformContext,
|
||||
getApiKey: this.getApiKey,
|
||||
getSteeringMessages: async () => {
|
||||
if (skipInitialSteeringPoll) {
|
||||
skipInitialSteeringPoll = false;
|
||||
return [];
|
||||
}
|
||||
return this.steeringQueue.drain();
|
||||
},
|
||||
getFollowUpMessages: async () => this.followUpQueue.drain(),
|
||||
};
|
||||
}
|
||||
|
||||
private async runWithLifecycle(executor: (signal: AbortSignal) => Promise<void>): Promise<void> {
|
||||
if (this.activeRun) {
|
||||
throw new Error("Agent is already processing.");
|
||||
}
|
||||
|
||||
const abortController = new AbortController();
|
||||
let resolvePromise = () => {};
|
||||
const promise = new Promise<void>((resolve) => {
|
||||
resolvePromise = resolve;
|
||||
});
|
||||
this.activeRun = { promise, resolve: resolvePromise, abortController };
|
||||
|
||||
this._state.isStreaming = true;
|
||||
this._state.streamingMessage = undefined;
|
||||
this._state.errorMessage = undefined;
|
||||
|
||||
try {
|
||||
await executor(abortController.signal);
|
||||
} catch (error) {
|
||||
await this.handleRunFailure(error, abortController.signal.aborted);
|
||||
} finally {
|
||||
this.finishRun();
|
||||
}
|
||||
}
|
||||
|
||||
private async handleRunFailure(error: unknown, aborted: boolean): Promise<void> {
|
||||
const failureMessage = {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "" }],
|
||||
api: this._state.model.api,
|
||||
provider: this._state.model.provider,
|
||||
model: this._state.model.id,
|
||||
usage: EMPTY_USAGE,
|
||||
stopReason: aborted ? "aborted" : "error",
|
||||
errorMessage: error instanceof Error ? error.message : String(error),
|
||||
timestamp: Date.now(),
|
||||
} satisfies AgentMessage;
|
||||
this._state.messages.push(failureMessage);
|
||||
this._state.errorMessage = failureMessage.errorMessage;
|
||||
await this.processEvents({ type: "agent_end", messages: [failureMessage] });
|
||||
}
|
||||
|
||||
private finishRun(): void {
|
||||
this._state.isStreaming = false;
|
||||
this._state.streamingMessage = undefined;
|
||||
this._state.pendingToolCalls = new Set<string>();
|
||||
this.activeRun?.resolve();
|
||||
this.activeRun = undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Reduce internal state for a loop event, then await listeners.
|
||||
*
|
||||
* `agent_end` only means no further loop events will be emitted. The run is
|
||||
* considered idle later, after all awaited listeners for `agent_end` finish
|
||||
* and `finishRun()` clears runtime-owned state.
|
||||
*/
|
||||
private async processEvents(event: AgentEvent): Promise<void> {
|
||||
switch (event.type) {
|
||||
case "message_start":
|
||||
this._state.streamingMessage = event.message;
|
||||
break;
|
||||
|
||||
case "message_update":
|
||||
this._state.streamingMessage = event.message;
|
||||
break;
|
||||
|
||||
case "message_end":
|
||||
this._state.streamingMessage = undefined;
|
||||
this._state.messages.push(event.message);
|
||||
break;
|
||||
|
||||
case "tool_execution_start": {
|
||||
const pendingToolCalls = new Set(this._state.pendingToolCalls);
|
||||
pendingToolCalls.add(event.toolCallId);
|
||||
this._state.pendingToolCalls = pendingToolCalls;
|
||||
break;
|
||||
}
|
||||
|
||||
case "tool_execution_end": {
|
||||
const pendingToolCalls = new Set(this._state.pendingToolCalls);
|
||||
pendingToolCalls.delete(event.toolCallId);
|
||||
this._state.pendingToolCalls = pendingToolCalls;
|
||||
break;
|
||||
}
|
||||
|
||||
case "turn_end":
|
||||
if (event.message.role === "assistant" && event.message.errorMessage) {
|
||||
this._state.errorMessage = event.message.errorMessage;
|
||||
}
|
||||
break;
|
||||
|
||||
case "agent_end":
|
||||
this._state.streamingMessage = undefined;
|
||||
break;
|
||||
}
|
||||
|
||||
const signal = this.activeRun?.abortController.signal;
|
||||
if (!signal) {
|
||||
throw new Error("Agent listener invoked outside active run");
|
||||
}
|
||||
for (const listener of this.listeners) {
|
||||
await listener(event, signal);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,8 +0,0 @@
|
||||
// Core Agent
|
||||
export * from "./agent.js";
|
||||
// Loop functions
|
||||
export * from "./agent-loop.js";
|
||||
// Proxy utilities
|
||||
export * from "./proxy.js";
|
||||
// Types
|
||||
export * from "./types.js";
|
||||
@@ -1,340 +0,0 @@
|
||||
/**
|
||||
* Proxy stream function for apps that route LLM calls through a server.
|
||||
* The server manages auth and proxies requests to LLM providers.
|
||||
*/
|
||||
|
||||
// Internal import for JSON parsing utility
|
||||
import {
|
||||
type AssistantMessage,
|
||||
type AssistantMessageEvent,
|
||||
type Context,
|
||||
EventStream,
|
||||
type Model,
|
||||
parseStreamingJson,
|
||||
type SimpleStreamOptions,
|
||||
type StopReason,
|
||||
type ToolCall,
|
||||
} from "@mariozechner/pi-ai";
|
||||
|
||||
// Create stream class matching ProxyMessageEventStream
|
||||
class ProxyMessageEventStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
|
||||
constructor() {
|
||||
super(
|
||||
(event) => event.type === "done" || event.type === "error",
|
||||
(event) => {
|
||||
if (event.type === "done") return event.message;
|
||||
if (event.type === "error") return event.error;
|
||||
throw new Error("Unexpected event type");
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Proxy event types - server sends these with partial field stripped to reduce bandwidth.
|
||||
*/
|
||||
export type ProxyAssistantMessageEvent =
|
||||
| { type: "start" }
|
||||
| { type: "text_start"; contentIndex: number }
|
||||
| { type: "text_delta"; contentIndex: number; delta: string }
|
||||
| { type: "text_end"; contentIndex: number; contentSignature?: string }
|
||||
| { type: "thinking_start"; contentIndex: number }
|
||||
| { type: "thinking_delta"; contentIndex: number; delta: string }
|
||||
| { type: "thinking_end"; contentIndex: number; contentSignature?: string }
|
||||
| { type: "toolcall_start"; contentIndex: number; id: string; toolName: string }
|
||||
| { type: "toolcall_delta"; contentIndex: number; delta: string }
|
||||
| { type: "toolcall_end"; contentIndex: number }
|
||||
| {
|
||||
type: "done";
|
||||
reason: Extract<StopReason, "stop" | "length" | "toolUse">;
|
||||
usage: AssistantMessage["usage"];
|
||||
}
|
||||
| {
|
||||
type: "error";
|
||||
reason: Extract<StopReason, "aborted" | "error">;
|
||||
errorMessage?: string;
|
||||
usage: AssistantMessage["usage"];
|
||||
};
|
||||
|
||||
export interface ProxyStreamOptions extends SimpleStreamOptions {
|
||||
/** Auth token for the proxy server */
|
||||
authToken: string;
|
||||
/** Proxy server URL (e.g., "https://genai.example.com") */
|
||||
proxyUrl: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Stream function that proxies through a server instead of calling LLM providers directly.
|
||||
* The server strips the partial field from delta events to reduce bandwidth.
|
||||
* We reconstruct the partial message client-side.
|
||||
*
|
||||
* Use this as the `streamFn` option when creating an Agent that needs to go through a proxy.
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* const agent = new Agent({
|
||||
* streamFn: (model, context, options) =>
|
||||
* streamProxy(model, context, {
|
||||
* ...options,
|
||||
* authToken: await getAuthToken(),
|
||||
* proxyUrl: "https://genai.example.com",
|
||||
* }),
|
||||
* });
|
||||
* ```
|
||||
*/
|
||||
export function streamProxy(model: Model<any>, context: Context, options: ProxyStreamOptions): ProxyMessageEventStream {
|
||||
const stream = new ProxyMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
// Initialize the partial message that we'll build up from events
|
||||
const partial: AssistantMessage = {
|
||||
role: "assistant",
|
||||
stopReason: "stop",
|
||||
content: [],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
let reader: ReadableStreamDefaultReader<Uint8Array> | undefined;
|
||||
|
||||
const abortHandler = () => {
|
||||
if (reader) {
|
||||
reader.cancel("Request aborted by user").catch(() => {});
|
||||
}
|
||||
};
|
||||
|
||||
if (options.signal) {
|
||||
options.signal.addEventListener("abort", abortHandler);
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await fetch(`${options.proxyUrl}/api/stream`, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
Authorization: `Bearer ${options.authToken}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model,
|
||||
context,
|
||||
options: {
|
||||
temperature: options.temperature,
|
||||
maxTokens: options.maxTokens,
|
||||
reasoning: options.reasoning,
|
||||
},
|
||||
}),
|
||||
signal: options.signal,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
let errorMessage = `Proxy error: ${response.status} ${response.statusText}`;
|
||||
try {
|
||||
const errorData = (await response.json()) as { error?: string };
|
||||
if (errorData.error) {
|
||||
errorMessage = `Proxy error: ${errorData.error}`;
|
||||
}
|
||||
} catch {
|
||||
// Couldn't parse error response
|
||||
}
|
||||
throw new Error(errorMessage);
|
||||
}
|
||||
|
||||
reader = response.body!.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
if (options.signal?.aborted) {
|
||||
throw new Error("Request aborted by user");
|
||||
}
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split("\n");
|
||||
buffer = lines.pop() || "";
|
||||
|
||||
for (const line of lines) {
|
||||
if (line.startsWith("data: ")) {
|
||||
const data = line.slice(6).trim();
|
||||
if (data) {
|
||||
const proxyEvent = JSON.parse(data) as ProxyAssistantMessageEvent;
|
||||
const event = processProxyEvent(proxyEvent, partial);
|
||||
if (event) {
|
||||
stream.push(event);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (options.signal?.aborted) {
|
||||
throw new Error("Request aborted by user");
|
||||
}
|
||||
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : String(error);
|
||||
const reason = options.signal?.aborted ? "aborted" : "error";
|
||||
partial.stopReason = reason;
|
||||
partial.errorMessage = errorMessage;
|
||||
stream.push({
|
||||
type: "error",
|
||||
reason,
|
||||
error: partial,
|
||||
});
|
||||
stream.end();
|
||||
} finally {
|
||||
if (options.signal) {
|
||||
options.signal.removeEventListener("abort", abortHandler);
|
||||
}
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
}
|
||||
|
||||
/**
|
||||
* Process a proxy event and update the partial message.
|
||||
*/
|
||||
function processProxyEvent(
|
||||
proxyEvent: ProxyAssistantMessageEvent,
|
||||
partial: AssistantMessage,
|
||||
): AssistantMessageEvent | undefined {
|
||||
switch (proxyEvent.type) {
|
||||
case "start":
|
||||
return { type: "start", partial };
|
||||
|
||||
case "text_start":
|
||||
partial.content[proxyEvent.contentIndex] = { type: "text", text: "" };
|
||||
return { type: "text_start", contentIndex: proxyEvent.contentIndex, partial };
|
||||
|
||||
case "text_delta": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "text") {
|
||||
content.text += proxyEvent.delta;
|
||||
return {
|
||||
type: "text_delta",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
delta: proxyEvent.delta,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received text_delta for non-text content");
|
||||
}
|
||||
|
||||
case "text_end": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "text") {
|
||||
content.textSignature = proxyEvent.contentSignature;
|
||||
return {
|
||||
type: "text_end",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
content: content.text,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received text_end for non-text content");
|
||||
}
|
||||
|
||||
case "thinking_start":
|
||||
partial.content[proxyEvent.contentIndex] = { type: "thinking", thinking: "" };
|
||||
return { type: "thinking_start", contentIndex: proxyEvent.contentIndex, partial };
|
||||
|
||||
case "thinking_delta": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "thinking") {
|
||||
content.thinking += proxyEvent.delta;
|
||||
return {
|
||||
type: "thinking_delta",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
delta: proxyEvent.delta,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received thinking_delta for non-thinking content");
|
||||
}
|
||||
|
||||
case "thinking_end": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "thinking") {
|
||||
content.thinkingSignature = proxyEvent.contentSignature;
|
||||
return {
|
||||
type: "thinking_end",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
content: content.thinking,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received thinking_end for non-thinking content");
|
||||
}
|
||||
|
||||
case "toolcall_start":
|
||||
partial.content[proxyEvent.contentIndex] = {
|
||||
type: "toolCall",
|
||||
id: proxyEvent.id,
|
||||
name: proxyEvent.toolName,
|
||||
arguments: {},
|
||||
partialJson: "",
|
||||
} satisfies ToolCall & { partialJson: string } as ToolCall;
|
||||
return { type: "toolcall_start", contentIndex: proxyEvent.contentIndex, partial };
|
||||
|
||||
case "toolcall_delta": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "toolCall") {
|
||||
(content as any).partialJson += proxyEvent.delta;
|
||||
content.arguments = parseStreamingJson((content as any).partialJson) || {};
|
||||
partial.content[proxyEvent.contentIndex] = { ...content }; // Trigger reactivity
|
||||
return {
|
||||
type: "toolcall_delta",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
delta: proxyEvent.delta,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
throw new Error("Received toolcall_delta for non-toolCall content");
|
||||
}
|
||||
|
||||
case "toolcall_end": {
|
||||
const content = partial.content[proxyEvent.contentIndex];
|
||||
if (content?.type === "toolCall") {
|
||||
delete (content as any).partialJson;
|
||||
return {
|
||||
type: "toolcall_end",
|
||||
contentIndex: proxyEvent.contentIndex,
|
||||
toolCall: content,
|
||||
partial,
|
||||
};
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
case "done":
|
||||
partial.stopReason = proxyEvent.reason;
|
||||
partial.usage = proxyEvent.usage;
|
||||
return { type: "done", reason: proxyEvent.reason, message: partial };
|
||||
|
||||
case "error":
|
||||
partial.stopReason = proxyEvent.reason;
|
||||
partial.errorMessage = proxyEvent.errorMessage;
|
||||
partial.usage = proxyEvent.usage;
|
||||
return { type: "error", reason: proxyEvent.reason, error: partial };
|
||||
|
||||
default: {
|
||||
const _exhaustiveCheck: never = proxyEvent;
|
||||
console.warn(`Unhandled proxy event type: ${(proxyEvent as any).type}`);
|
||||
return undefined;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,349 +0,0 @@
|
||||
import type {
|
||||
AssistantMessage,
|
||||
AssistantMessageEvent,
|
||||
ImageContent,
|
||||
Message,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
streamSimple,
|
||||
TextContent,
|
||||
Tool,
|
||||
ToolResultMessage,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import type { Static, TSchema } from "@sinclair/typebox";
|
||||
|
||||
/**
|
||||
* Stream function used by the agent loop.
|
||||
*
|
||||
* Contract:
|
||||
* - Must not throw or return a rejected promise for request/model/runtime failures.
|
||||
* - Must return an AssistantMessageEventStream.
|
||||
* - Failures must be encoded in the returned stream via protocol events and a
|
||||
* final AssistantMessage with stopReason "error" or "aborted" and errorMessage.
|
||||
*/
|
||||
export type StreamFn = (
|
||||
...args: Parameters<typeof streamSimple>
|
||||
) => ReturnType<typeof streamSimple> | Promise<ReturnType<typeof streamSimple>>;
|
||||
|
||||
/**
|
||||
* Configuration for how tool calls from a single assistant message are executed.
|
||||
*
|
||||
* - "sequential": each tool call is prepared, executed, and finalized before the next one starts.
|
||||
* - "parallel": tool calls are prepared sequentially, then allowed tools execute concurrently.
|
||||
* Final tool results are still emitted in assistant source order.
|
||||
*/
|
||||
export type ToolExecutionMode = "sequential" | "parallel";
|
||||
|
||||
/** A single tool call content block emitted by an assistant message. */
|
||||
export type AgentToolCall = Extract<AssistantMessage["content"][number], { type: "toolCall" }>;
|
||||
|
||||
/**
|
||||
* Result returned from `beforeToolCall`.
|
||||
*
|
||||
* Returning `{ block: true }` prevents the tool from executing. The loop emits an error tool result instead.
|
||||
* `reason` becomes the text shown in that error result. If omitted, a default blocked message is used.
|
||||
*/
|
||||
export interface BeforeToolCallResult {
|
||||
block?: boolean;
|
||||
reason?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Partial override returned from `afterToolCall`.
|
||||
*
|
||||
* Merge semantics are field-by-field:
|
||||
* - `content`: if provided, replaces the tool result content array in full
|
||||
* - `details`: if provided, replaces the tool result details value in full
|
||||
* - `isError`: if provided, replaces the tool result error flag
|
||||
*
|
||||
* Omitted fields keep the original executed tool result values.
|
||||
* There is no deep merge for `content` or `details`.
|
||||
*/
|
||||
export interface AfterToolCallResult {
|
||||
content?: (TextContent | ImageContent)[];
|
||||
details?: unknown;
|
||||
isError?: boolean;
|
||||
}
|
||||
|
||||
/** Context passed to `beforeToolCall`. */
|
||||
export interface BeforeToolCallContext {
|
||||
/** The assistant message that requested the tool call. */
|
||||
assistantMessage: AssistantMessage;
|
||||
/** The raw tool call block from `assistantMessage.content`. */
|
||||
toolCall: AgentToolCall;
|
||||
/** Validated tool arguments for the target tool schema. */
|
||||
args: unknown;
|
||||
/** Current agent context at the time the tool call is prepared. */
|
||||
context: AgentContext;
|
||||
}
|
||||
|
||||
/** Context passed to `afterToolCall`. */
|
||||
export interface AfterToolCallContext {
|
||||
/** The assistant message that requested the tool call. */
|
||||
assistantMessage: AssistantMessage;
|
||||
/** The raw tool call block from `assistantMessage.content`. */
|
||||
toolCall: AgentToolCall;
|
||||
/** Validated tool arguments for the target tool schema. */
|
||||
args: unknown;
|
||||
/** The executed tool result before any `afterToolCall` overrides are applied. */
|
||||
result: AgentToolResult<any>;
|
||||
/** Whether the executed tool result is currently treated as an error. */
|
||||
isError: boolean;
|
||||
/** Current agent context at the time the tool call is finalized. */
|
||||
context: AgentContext;
|
||||
}
|
||||
|
||||
export interface AgentLoopConfig extends SimpleStreamOptions {
|
||||
model: Model<any>;
|
||||
|
||||
/**
|
||||
* Converts AgentMessage[] to LLM-compatible Message[] before each LLM call.
|
||||
*
|
||||
* Each AgentMessage must be converted to a UserMessage, AssistantMessage, or ToolResultMessage
|
||||
* that the LLM can understand. AgentMessages that cannot be converted (e.g., UI-only notifications,
|
||||
* status messages) should be filtered out.
|
||||
*
|
||||
* Contract: must not throw or reject. Return a safe fallback value instead.
|
||||
* Throwing interrupts the low-level agent loop without producing a normal event sequence.
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* convertToLlm: (messages) => messages.flatMap(m => {
|
||||
* if (m.role === "custom") {
|
||||
* // Convert custom message to user message
|
||||
* return [{ role: "user", content: m.content, timestamp: m.timestamp }];
|
||||
* }
|
||||
* if (m.role === "notification") {
|
||||
* // Filter out UI-only messages
|
||||
* return [];
|
||||
* }
|
||||
* // Pass through standard LLM messages
|
||||
* return [m];
|
||||
* })
|
||||
* ```
|
||||
*/
|
||||
convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
|
||||
|
||||
/**
|
||||
* Optional transform applied to the context before `convertToLlm`.
|
||||
*
|
||||
* Use this for operations that work at the AgentMessage level:
|
||||
* - Context window management (pruning old messages)
|
||||
* - Injecting context from external sources
|
||||
*
|
||||
* Contract: must not throw or reject. Return the original messages or another
|
||||
* safe fallback value instead.
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* transformContext: async (messages) => {
|
||||
* if (estimateTokens(messages) > MAX_TOKENS) {
|
||||
* return pruneOldMessages(messages);
|
||||
* }
|
||||
* return messages;
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
|
||||
|
||||
/**
|
||||
* Resolves an API key dynamically for each LLM call.
|
||||
*
|
||||
* Useful for short-lived OAuth tokens (e.g., GitHub Copilot) that may expire
|
||||
* during long-running tool execution phases.
|
||||
*
|
||||
* Contract: must not throw or reject. Return undefined when no key is available.
|
||||
*/
|
||||
getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
|
||||
|
||||
/**
|
||||
* Returns steering messages to inject into the conversation mid-run.
|
||||
*
|
||||
* Called after the current assistant turn finishes executing its tool calls.
|
||||
* If messages are returned, they are added to the context before the next LLM call.
|
||||
* Tool calls from the current assistant message are not skipped.
|
||||
*
|
||||
* Use this for "steering" the agent while it's working.
|
||||
*
|
||||
* Contract: must not throw or reject. Return [] when no steering messages are available.
|
||||
*/
|
||||
getSteeringMessages?: () => Promise<AgentMessage[]>;
|
||||
|
||||
/**
|
||||
* Returns follow-up messages to process after the agent would otherwise stop.
|
||||
*
|
||||
* Called when the agent has no more tool calls and no steering messages.
|
||||
* If messages are returned, they're added to the context and the agent
|
||||
* continues with another turn.
|
||||
*
|
||||
* Use this for follow-up messages that should wait until the agent finishes.
|
||||
*
|
||||
* Contract: must not throw or reject. Return [] when no follow-up messages are available.
|
||||
*/
|
||||
getFollowUpMessages?: () => Promise<AgentMessage[]>;
|
||||
|
||||
/**
|
||||
* Tool execution mode.
|
||||
* - "sequential": execute tool calls one by one
|
||||
* - "parallel": preflight tool calls sequentially, then execute allowed tools concurrently
|
||||
*
|
||||
* Default: "parallel"
|
||||
*/
|
||||
toolExecution?: ToolExecutionMode;
|
||||
|
||||
/**
|
||||
* Called before a tool is executed, after arguments have been validated.
|
||||
*
|
||||
* Return `{ block: true }` to prevent execution. The loop emits an error tool result instead.
|
||||
* The hook receives the agent abort signal and is responsible for honoring it.
|
||||
*/
|
||||
beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
|
||||
|
||||
/**
|
||||
* Called after a tool finishes executing, before final tool events are emitted.
|
||||
*
|
||||
* Return an `AfterToolCallResult` to override parts of the executed tool result:
|
||||
* - `content` replaces the full content array
|
||||
* - `details` replaces the full details payload
|
||||
* - `isError` replaces the error flag
|
||||
*
|
||||
* Any omitted fields keep their original values. No deep merge is performed.
|
||||
* The hook receives the agent abort signal and is responsible for honoring it.
|
||||
*/
|
||||
afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Thinking/reasoning level for models that support it.
|
||||
* Note: "xhigh" is only supported by OpenAI gpt-5.1-codex-max, gpt-5.2, gpt-5.2-codex, gpt-5.3, and gpt-5.3-codex models.
|
||||
*/
|
||||
export type ThinkingLevel = "off" | "minimal" | "low" | "medium" | "high" | "xhigh";
|
||||
|
||||
/**
|
||||
* Extensible interface for custom app messages.
|
||||
* Apps can extend via declaration merging:
|
||||
*
|
||||
* @example
|
||||
* ```typescript
|
||||
* declare module "@mariozechner/agent" {
|
||||
* interface CustomAgentMessages {
|
||||
* artifact: ArtifactMessage;
|
||||
* notification: NotificationMessage;
|
||||
* }
|
||||
* }
|
||||
* ```
|
||||
*/
|
||||
export interface CustomAgentMessages {
|
||||
// Empty by default - apps extend via declaration merging
|
||||
}
|
||||
|
||||
/**
|
||||
* AgentMessage: Union of LLM messages + custom messages.
|
||||
* This abstraction allows apps to add custom message types while maintaining
|
||||
* type safety and compatibility with the base LLM messages.
|
||||
*/
|
||||
export type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages];
|
||||
|
||||
/**
|
||||
* Public agent state.
|
||||
*
|
||||
* `tools` and `messages` use accessor properties so implementations can copy
|
||||
* assigned arrays before storing them.
|
||||
*/
|
||||
export interface AgentState {
|
||||
/** System prompt sent with each model request. */
|
||||
systemPrompt: string;
|
||||
/** Active model used for future turns. */
|
||||
model: Model<any>;
|
||||
/** Requested reasoning level for future turns. */
|
||||
thinkingLevel: ThinkingLevel;
|
||||
/** Available tools. Assigning a new array copies the top-level array. */
|
||||
set tools(tools: AgentTool<any>[]);
|
||||
get tools(): AgentTool<any>[];
|
||||
/** Conversation transcript. Assigning a new array copies the top-level array. */
|
||||
set messages(messages: AgentMessage[]);
|
||||
get messages(): AgentMessage[];
|
||||
/**
|
||||
* True while the agent is processing a prompt or continuation.
|
||||
*
|
||||
* This remains true until awaited `agent_end` listeners settle.
|
||||
*/
|
||||
readonly isStreaming: boolean;
|
||||
/** Partial assistant message for the current streamed response, if any. */
|
||||
readonly streamingMessage?: AgentMessage;
|
||||
/** Tool call ids currently executing. */
|
||||
readonly pendingToolCalls: ReadonlySet<string>;
|
||||
/** Error message from the most recent failed or aborted assistant turn, if any. */
|
||||
readonly errorMessage?: string;
|
||||
}
|
||||
|
||||
/** Final or partial result produced by a tool. */
|
||||
export interface AgentToolResult<T> {
|
||||
/** Text or image content returned to the model. */
|
||||
content: (TextContent | ImageContent)[];
|
||||
/** Arbitrary structured details for logs or UI rendering. */
|
||||
details: T;
|
||||
}
|
||||
|
||||
/** Callback used by tools to stream partial execution updates. */
|
||||
export type AgentToolUpdateCallback<T = any> = (partialResult: AgentToolResult<T>) => void;
|
||||
|
||||
/** Tool definition used by the agent runtime. */
|
||||
export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any> extends Tool<TParameters> {
|
||||
/** Human-readable label for UI display. */
|
||||
label: string;
|
||||
/**
|
||||
* Optional compatibility shim for raw tool-call arguments before schema validation.
|
||||
* Must return an object that matches `TParameters`.
|
||||
*/
|
||||
prepareArguments?: (args: unknown) => Static<TParameters>;
|
||||
/** Execute the tool call. Throw on failure instead of encoding errors in `content`. */
|
||||
execute: (
|
||||
toolCallId: string,
|
||||
params: Static<TParameters>,
|
||||
signal?: AbortSignal,
|
||||
onUpdate?: AgentToolUpdateCallback<TDetails>,
|
||||
) => Promise<AgentToolResult<TDetails>>;
|
||||
/**
|
||||
* Per-tool execution mode override.
|
||||
* - "sequential": this tool must execute one at a time with other tool calls.
|
||||
* - "parallel": this tool can execute concurrently with other tool calls.
|
||||
*
|
||||
* If omitted, the default execution mode applies.
|
||||
*/
|
||||
executionMode?: ToolExecutionMode;
|
||||
}
|
||||
|
||||
/** Context snapshot passed into the low-level agent loop. */
|
||||
export interface AgentContext {
|
||||
/** System prompt included with the request. */
|
||||
systemPrompt: string;
|
||||
/** Transcript visible to the model. */
|
||||
messages: AgentMessage[];
|
||||
/** Tools available for this run. */
|
||||
tools?: AgentTool<any>[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Events emitted by the Agent for UI updates.
|
||||
*
|
||||
* `agent_end` is the last event emitted for a run, but awaited `Agent.subscribe()`
|
||||
* listeners for that event are still part of run settlement. The agent becomes
|
||||
* idle only after those listeners finish.
|
||||
*/
|
||||
export type AgentEvent =
|
||||
// Agent lifecycle
|
||||
| { type: "agent_start" }
|
||||
| { type: "agent_end"; messages: AgentMessage[] }
|
||||
// Turn lifecycle - a turn is one assistant response + any tool calls/results
|
||||
| { type: "turn_start" }
|
||||
| { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[] }
|
||||
// Message lifecycle - emitted for user, assistant, and toolResult messages
|
||||
| { type: "message_start"; message: AgentMessage }
|
||||
// Only emitted for assistant messages during streaming
|
||||
| { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent }
|
||||
| { type: "message_end"; message: AgentMessage }
|
||||
// Tool execution lifecycle
|
||||
| { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any }
|
||||
| { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any }
|
||||
| { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean };
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1,463 +0,0 @@
|
||||
import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@mariozechner/pi-ai";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { Agent } from "../src/index.js";
|
||||
|
||||
// Mock stream that mimics AssistantMessageEventStream
|
||||
class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
|
||||
constructor() {
|
||||
super(
|
||||
(event) => event.type === "done" || event.type === "error",
|
||||
(event) => {
|
||||
if (event.type === "done") return event.message;
|
||||
if (event.type === "error") return event.error;
|
||||
throw new Error("Unexpected event type");
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
function createAssistantMessage(text: string): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text }],
|
||||
api: "openai-responses",
|
||||
provider: "openai",
|
||||
model: "mock",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createDeferred(): {
|
||||
promise: Promise<void>;
|
||||
resolve: () => void;
|
||||
} {
|
||||
let resolve = () => {};
|
||||
const promise = new Promise<void>((resolvePromise) => {
|
||||
resolve = resolvePromise;
|
||||
});
|
||||
return { promise, resolve };
|
||||
}
|
||||
|
||||
describe("Agent", () => {
|
||||
it("should create an agent instance with default state", () => {
|
||||
const agent = new Agent();
|
||||
|
||||
expect(agent.state).toBeDefined();
|
||||
expect(agent.state.systemPrompt).toBe("");
|
||||
expect(agent.state.model).toBeDefined();
|
||||
expect(agent.state.thinkingLevel).toBe("off");
|
||||
expect(agent.state.tools).toEqual([]);
|
||||
expect(agent.state.messages).toEqual([]);
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.streamingMessage).toBe(undefined);
|
||||
expect(agent.state.pendingToolCalls).toEqual(new Set());
|
||||
expect(agent.state.errorMessage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("should create an agent instance with custom initial state", () => {
|
||||
const customModel = getModel("openai", "gpt-4o-mini");
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model: customModel,
|
||||
thinkingLevel: "low",
|
||||
},
|
||||
});
|
||||
|
||||
expect(agent.state.systemPrompt).toBe("You are a helpful assistant.");
|
||||
expect(agent.state.model).toBe(customModel);
|
||||
expect(agent.state.thinkingLevel).toBe("low");
|
||||
});
|
||||
|
||||
it("should subscribe to events", () => {
|
||||
const agent = new Agent();
|
||||
|
||||
let eventCount = 0;
|
||||
const unsubscribe = agent.subscribe((_event) => {
|
||||
eventCount++;
|
||||
});
|
||||
|
||||
// No initial event on subscribe
|
||||
expect(eventCount).toBe(0);
|
||||
|
||||
// State mutators don't emit events
|
||||
agent.state.systemPrompt = "Test prompt";
|
||||
expect(eventCount).toBe(0);
|
||||
expect(agent.state.systemPrompt).toBe("Test prompt");
|
||||
|
||||
// Unsubscribe should work
|
||||
unsubscribe();
|
||||
agent.state.systemPrompt = "Another prompt";
|
||||
expect(eventCount).toBe(0); // Should not increase
|
||||
});
|
||||
|
||||
it("should await async subscribers before prompt resolves", async () => {
|
||||
const barrier = createDeferred();
|
||||
const agent = new Agent({
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
let listenerFinished = false;
|
||||
agent.subscribe(async (event) => {
|
||||
if (event.type === "agent_end") {
|
||||
await barrier.promise;
|
||||
listenerFinished = true;
|
||||
}
|
||||
});
|
||||
|
||||
let promptResolved = false;
|
||||
const promptPromise = agent.prompt("hello").then(() => {
|
||||
promptResolved = true;
|
||||
});
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
expect(promptResolved).toBe(false);
|
||||
expect(listenerFinished).toBe(false);
|
||||
expect(agent.state.isStreaming).toBe(true);
|
||||
|
||||
barrier.resolve();
|
||||
await promptPromise;
|
||||
|
||||
expect(listenerFinished).toBe(true);
|
||||
expect(promptResolved).toBe(true);
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
});
|
||||
|
||||
it("waitForIdle should wait for async subscribers", async () => {
|
||||
const barrier = createDeferred();
|
||||
const agent = new Agent({
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
agent.subscribe(async (event) => {
|
||||
if (event.type === "message_end" && event.message.role === "assistant") {
|
||||
await barrier.promise;
|
||||
}
|
||||
});
|
||||
|
||||
const promptPromise = agent.prompt("hello");
|
||||
let idleResolved = false;
|
||||
const idlePromise = agent.waitForIdle().then(() => {
|
||||
idleResolved = true;
|
||||
});
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
expect(idleResolved).toBe(false);
|
||||
expect(agent.state.isStreaming).toBe(true);
|
||||
|
||||
barrier.resolve();
|
||||
await Promise.all([promptPromise, idlePromise]);
|
||||
|
||||
expect(idleResolved).toBe(true);
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
});
|
||||
|
||||
it("should pass the active abort signal to subscribers", async () => {
|
||||
let receivedSignal: AbortSignal | undefined;
|
||||
const agent = new Agent({
|
||||
streamFn: (_model, _context, options) => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
const checkAbort = () => {
|
||||
if (options?.signal?.aborted) {
|
||||
stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
|
||||
} else {
|
||||
setTimeout(checkAbort, 5);
|
||||
}
|
||||
};
|
||||
checkAbort();
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
agent.subscribe((event, signal) => {
|
||||
if (event.type === "agent_start") {
|
||||
receivedSignal = signal;
|
||||
}
|
||||
});
|
||||
|
||||
const promptPromise = agent.prompt("hello");
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
|
||||
expect(receivedSignal).toBeDefined();
|
||||
expect(receivedSignal?.aborted).toBe(false);
|
||||
|
||||
agent.abort();
|
||||
await promptPromise;
|
||||
|
||||
expect(receivedSignal?.aborted).toBe(true);
|
||||
});
|
||||
|
||||
it("should update state with mutators", () => {
|
||||
const agent = new Agent();
|
||||
|
||||
// Test setSystemPrompt
|
||||
agent.state.systemPrompt = "Custom prompt";
|
||||
expect(agent.state.systemPrompt).toBe("Custom prompt");
|
||||
|
||||
// Test setModel
|
||||
const newModel = getModel("google", "gemini-2.5-flash");
|
||||
agent.state.model = newModel;
|
||||
expect(agent.state.model).toBe(newModel);
|
||||
|
||||
// Test setThinkingLevel
|
||||
agent.state.thinkingLevel = "high";
|
||||
expect(agent.state.thinkingLevel).toBe("high");
|
||||
|
||||
// Test setTools
|
||||
const tools = [{ name: "test", description: "test tool" } as any];
|
||||
agent.state.tools = tools;
|
||||
expect(agent.state.tools).toEqual(tools);
|
||||
expect(agent.state.tools).not.toBe(tools); // Should be a copy
|
||||
|
||||
// Test replaceMessages
|
||||
const messages = [{ role: "user" as const, content: "Hello", timestamp: Date.now() }];
|
||||
agent.state.messages = messages;
|
||||
expect(agent.state.messages).toEqual(messages);
|
||||
expect(agent.state.messages).not.toBe(messages); // Should be a copy
|
||||
|
||||
// Test appendMessage
|
||||
const newMessage = { role: "assistant" as const, content: [{ type: "text" as const, text: "Hi" }] };
|
||||
agent.state.messages.push(newMessage as any);
|
||||
expect(agent.state.messages).toHaveLength(2);
|
||||
expect(agent.state.messages[1]).toBe(newMessage);
|
||||
|
||||
// Test clearMessages
|
||||
agent.state.messages = [];
|
||||
expect(agent.state.messages).toEqual([]);
|
||||
});
|
||||
|
||||
it("should support steering message queue", async () => {
|
||||
const agent = new Agent();
|
||||
|
||||
const message = { role: "user" as const, content: "Steering message", timestamp: Date.now() };
|
||||
agent.steer(message);
|
||||
|
||||
// The message is queued but not yet in state.messages
|
||||
expect(agent.state.messages).not.toContainEqual(message);
|
||||
});
|
||||
|
||||
it("should support follow-up message queue", async () => {
|
||||
const agent = new Agent();
|
||||
|
||||
const message = { role: "user" as const, content: "Follow-up message", timestamp: Date.now() };
|
||||
agent.followUp(message);
|
||||
|
||||
// The message is queued but not yet in state.messages
|
||||
expect(agent.state.messages).not.toContainEqual(message);
|
||||
});
|
||||
|
||||
it("should handle abort controller", () => {
|
||||
const agent = new Agent();
|
||||
|
||||
// Should not throw even if nothing is running
|
||||
expect(() => agent.abort()).not.toThrow();
|
||||
});
|
||||
|
||||
it("should throw when prompt() called while streaming", async () => {
|
||||
let abortSignal: AbortSignal | undefined;
|
||||
const agent = new Agent({
|
||||
// Use a stream function that responds to abort
|
||||
streamFn: (_model, _context, options) => {
|
||||
abortSignal = options?.signal;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
// Check abort signal periodically
|
||||
const checkAbort = () => {
|
||||
if (abortSignal?.aborted) {
|
||||
stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
|
||||
} else {
|
||||
setTimeout(checkAbort, 5);
|
||||
}
|
||||
};
|
||||
checkAbort();
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
// Start first prompt (don't await, it will block until abort)
|
||||
const firstPrompt = agent.prompt("First message");
|
||||
|
||||
// Wait a tick for isStreaming to be set
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
expect(agent.state.isStreaming).toBe(true);
|
||||
|
||||
// Second prompt should reject
|
||||
await expect(agent.prompt("Second message")).rejects.toThrow(
|
||||
"Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.",
|
||||
);
|
||||
|
||||
// Cleanup - abort to stop the stream
|
||||
agent.abort();
|
||||
await firstPrompt.catch(() => {}); // Ignore abort error
|
||||
});
|
||||
|
||||
it("should throw when continue() called while streaming", async () => {
|
||||
let abortSignal: AbortSignal | undefined;
|
||||
const agent = new Agent({
|
||||
streamFn: (_model, _context, options) => {
|
||||
abortSignal = options?.signal;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
const checkAbort = () => {
|
||||
if (abortSignal?.aborted) {
|
||||
stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
|
||||
} else {
|
||||
setTimeout(checkAbort, 5);
|
||||
}
|
||||
};
|
||||
checkAbort();
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
// Start first prompt
|
||||
const firstPrompt = agent.prompt("First message");
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
expect(agent.state.isStreaming).toBe(true);
|
||||
|
||||
// continue() should reject
|
||||
await expect(agent.continue()).rejects.toThrow(
|
||||
"Agent is already processing. Wait for completion before continuing.",
|
||||
);
|
||||
|
||||
// Cleanup
|
||||
agent.abort();
|
||||
await firstPrompt.catch(() => {});
|
||||
});
|
||||
|
||||
it("continue() should process queued follow-up messages after an assistant turn", async () => {
|
||||
const agent = new Agent({
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Processed") });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
agent.state.messages = [
|
||||
{
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Initial" }],
|
||||
timestamp: Date.now() - 10,
|
||||
},
|
||||
createAssistantMessage("Initial response"),
|
||||
];
|
||||
|
||||
agent.followUp({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Queued follow-up" }],
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
|
||||
await expect(agent.continue()).resolves.toBeUndefined();
|
||||
|
||||
const hasQueuedFollowUp = agent.state.messages.some((message) => {
|
||||
if (message.role !== "user") return false;
|
||||
if (typeof message.content === "string") return message.content === "Queued follow-up";
|
||||
return message.content.some((part) => part.type === "text" && part.text === "Queued follow-up");
|
||||
});
|
||||
|
||||
expect(hasQueuedFollowUp).toBe(true);
|
||||
expect(agent.state.messages[agent.state.messages.length - 1].role).toBe("assistant");
|
||||
});
|
||||
|
||||
it("continue() should keep one-at-a-time steering semantics from assistant tail", async () => {
|
||||
let responseCount = 0;
|
||||
const agent = new Agent({
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
responseCount++;
|
||||
queueMicrotask(() => {
|
||||
stream.push({
|
||||
type: "done",
|
||||
reason: "stop",
|
||||
message: createAssistantMessage(`Processed ${responseCount}`),
|
||||
});
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
agent.state.messages = [
|
||||
{
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Initial" }],
|
||||
timestamp: Date.now() - 10,
|
||||
},
|
||||
createAssistantMessage("Initial response"),
|
||||
];
|
||||
|
||||
agent.steer({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Steering 1" }],
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
agent.steer({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Steering 2" }],
|
||||
timestamp: Date.now() + 1,
|
||||
});
|
||||
|
||||
await expect(agent.continue()).resolves.toBeUndefined();
|
||||
|
||||
const recentMessages = agent.state.messages.slice(-4);
|
||||
expect(recentMessages.map((m) => m.role)).toEqual(["user", "assistant", "user", "assistant"]);
|
||||
expect(responseCount).toBe(2);
|
||||
});
|
||||
|
||||
it("forwards sessionId to streamFn options", async () => {
|
||||
let receivedSessionId: string | undefined;
|
||||
const agent = new Agent({
|
||||
sessionId: "session-abc",
|
||||
streamFn: (_model, _context, options) => {
|
||||
receivedSessionId = options?.sessionId;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage("ok");
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
await agent.prompt("hello");
|
||||
expect(receivedSessionId).toBe("session-abc");
|
||||
|
||||
// Test setter
|
||||
agent.sessionId = "session-def";
|
||||
expect(agent.sessionId).toBe("session-def");
|
||||
|
||||
await agent.prompt("hello again");
|
||||
expect(receivedSessionId).toBe("session-def");
|
||||
});
|
||||
});
|
||||
@@ -1,404 +0,0 @@
|
||||
import {
|
||||
type AssistantMessage,
|
||||
type FauxProviderRegistration,
|
||||
fauxAssistantMessage,
|
||||
fauxText,
|
||||
fauxThinking,
|
||||
fauxToolCall,
|
||||
type Model,
|
||||
registerFauxProvider,
|
||||
type ToolResultMessage,
|
||||
type UserMessage,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { Agent, type AgentEvent } from "../src/index.js";
|
||||
import { calculateTool } from "./utils/calculate.js";
|
||||
|
||||
const registrations: FauxProviderRegistration[] = [];
|
||||
|
||||
function createFauxRegistration(options: Parameters<typeof registerFauxProvider>[0] = {}): FauxProviderRegistration {
|
||||
const registration = registerFauxProvider(options);
|
||||
registrations.push(registration);
|
||||
return registration;
|
||||
}
|
||||
|
||||
function getTextContent(message: AssistantMessage | ToolResultMessage): string {
|
||||
return message.content
|
||||
.filter((block) => block.type === "text")
|
||||
.map((block) => block.text)
|
||||
.join("\n");
|
||||
}
|
||||
|
||||
afterEach(() => {
|
||||
while (registrations.length > 0) {
|
||||
registrations.pop()?.unregister();
|
||||
}
|
||||
});
|
||||
|
||||
async function basicPrompt(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant. Keep your responses concise.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
await agent.prompt("What is 2+2? Answer with just the number.");
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBe(2);
|
||||
expect(agent.state.messages[0].role).toBe("user");
|
||||
expect(agent.state.messages[1].role).toBe("assistant");
|
||||
|
||||
const assistantMessage = agent.state.messages[1];
|
||||
if (assistantMessage.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(getTextContent(assistantMessage)).toContain("4");
|
||||
}
|
||||
|
||||
async function toolExecution(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant. Always use the calculator tool for math.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [calculateTool],
|
||||
},
|
||||
});
|
||||
|
||||
const pendingToolCallsDuringEvents: Array<{ type: AgentEvent["type"]; ids: string[] }> = [];
|
||||
agent.subscribe((event) => {
|
||||
if (event.type === "tool_execution_start" || event.type === "tool_execution_end") {
|
||||
pendingToolCallsDuringEvents.push({
|
||||
type: event.type,
|
||||
ids: [...agent.state.pendingToolCalls],
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
await agent.prompt("Calculate 123 * 456 using the calculator tool.");
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBeGreaterThanOrEqual(4);
|
||||
const toolResultMsg = agent.state.messages.find((message) => message.role === "toolResult");
|
||||
expect(toolResultMsg).toBeDefined();
|
||||
if (toolResultMsg?.role !== "toolResult") throw new Error("Expected tool result message");
|
||||
expect(getTextContent(toolResultMsg)).toContain("123 * 456 = 56088");
|
||||
|
||||
const finalMessage = agent.state.messages[agent.state.messages.length - 1];
|
||||
if (finalMessage.role !== "assistant") throw new Error("Expected final assistant message");
|
||||
expect(getTextContent(finalMessage)).toContain("56088");
|
||||
expect(agent.state.pendingToolCalls.size).toBe(0);
|
||||
expect(pendingToolCallsDuringEvents).toEqual([
|
||||
{ type: "tool_execution_start", ids: ["calc-1"] },
|
||||
{ type: "tool_execution_end", ids: [] },
|
||||
]);
|
||||
}
|
||||
|
||||
async function abortExecution(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
const promptPromise = agent.prompt("Count slowly from 1 to 20.");
|
||||
setTimeout(() => {
|
||||
agent.abort();
|
||||
}, 30);
|
||||
|
||||
await promptPromise;
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBeGreaterThanOrEqual(2);
|
||||
|
||||
const lastMessage = agent.state.messages[agent.state.messages.length - 1];
|
||||
if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(lastMessage.stopReason).toBe("aborted");
|
||||
expect(lastMessage.errorMessage).toBeDefined();
|
||||
expect(agent.state.errorMessage).toBe(lastMessage.errorMessage);
|
||||
}
|
||||
|
||||
async function stateUpdates(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
const events: AgentEvent["type"][] = [];
|
||||
agent.subscribe((event) => {
|
||||
events.push(event.type);
|
||||
});
|
||||
|
||||
await agent.prompt("Count from 1 to 5.");
|
||||
|
||||
expect(events).toContain("agent_start");
|
||||
expect(events).toContain("turn_start");
|
||||
expect(events).toContain("message_start");
|
||||
expect(events).toContain("message_update");
|
||||
expect(events).toContain("message_end");
|
||||
expect(events).toContain("turn_end");
|
||||
expect(events).toContain("agent_end");
|
||||
expect(events.indexOf("agent_start")).toBeLessThan(events.indexOf("message_start"));
|
||||
expect(events.indexOf("message_start")).toBeLessThan(events.indexOf("message_end"));
|
||||
expect(events.indexOf("message_end")).toBeLessThan(events.lastIndexOf("agent_end"));
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBe(2);
|
||||
}
|
||||
|
||||
async function multiTurnConversation(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
await agent.prompt("My name is Alice.");
|
||||
expect(agent.state.messages.length).toBe(2);
|
||||
|
||||
await agent.prompt("What is my name?");
|
||||
expect(agent.state.messages.length).toBe(4);
|
||||
|
||||
const lastMessage = agent.state.messages[3];
|
||||
if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(getTextContent(lastMessage).toLowerCase()).toContain("alice");
|
||||
}
|
||||
|
||||
describe("Agent integration with faux provider", () => {
|
||||
it("handles a basic text prompt", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
faux.setResponses([fauxAssistantMessage("4")]);
|
||||
await basicPrompt(faux.getModel());
|
||||
});
|
||||
|
||||
it("executes tools and tracks pending tool calls", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
faux.setResponses([
|
||||
fauxAssistantMessage(
|
||||
[
|
||||
fauxText("Let me calculate that."),
|
||||
fauxToolCall("calculate", { expression: "123 * 456" }, { id: "calc-1" }),
|
||||
],
|
||||
{ stopReason: "toolUse" },
|
||||
),
|
||||
fauxAssistantMessage("The result is 56088."),
|
||||
]);
|
||||
await toolExecution(faux.getModel());
|
||||
});
|
||||
|
||||
it("handles abort during streaming", async () => {
|
||||
const faux = createFauxRegistration({
|
||||
tokensPerSecond: 20,
|
||||
tokenSize: { min: 2, max: 2 },
|
||||
});
|
||||
faux.setResponses([
|
||||
fauxAssistantMessage(
|
||||
"one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen",
|
||||
),
|
||||
]);
|
||||
await abortExecution(faux.getModel());
|
||||
});
|
||||
|
||||
it("emits lifecycle updates while streaming", async () => {
|
||||
const faux = createFauxRegistration({ tokenSize: { min: 1, max: 1 } });
|
||||
faux.setResponses([fauxAssistantMessage("1 2 3 4 5")]);
|
||||
await stateUpdates(faux.getModel());
|
||||
});
|
||||
|
||||
it("maintains context across multiple turns", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
faux.setResponses([
|
||||
fauxAssistantMessage("Nice to meet you, Alice."),
|
||||
(context) => {
|
||||
const hasAlice = context.messages.some((message) => {
|
||||
if (message.role !== "user") return false;
|
||||
if (typeof message.content === "string") return message.content.includes("Alice");
|
||||
return message.content.some((block) => block.type === "text" && block.text.includes("Alice"));
|
||||
});
|
||||
return fauxAssistantMessage(hasAlice ? "Your name is Alice." : "I do not know your name.");
|
||||
},
|
||||
]);
|
||||
await multiTurnConversation(faux.getModel());
|
||||
});
|
||||
|
||||
it("preserves thinking content blocks", async () => {
|
||||
const faux = createFauxRegistration({ models: [{ id: "faux-reasoning", reasoning: true }] });
|
||||
faux.setResponses([fauxAssistantMessage([fauxThinking("step by step"), fauxText("4")])]);
|
||||
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model: faux.getModel(),
|
||||
thinkingLevel: "low",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
await agent.prompt("What is 2+2?");
|
||||
|
||||
const assistantMessage = agent.state.messages[1];
|
||||
if (assistantMessage?.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(assistantMessage.content).toEqual([
|
||||
{ type: "thinking", thinking: "step by step" },
|
||||
{ type: "text", text: "4" },
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Agent.continue() with faux provider", () => {
|
||||
describe("validation", () => {
|
||||
it("throws when no messages in context", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "Test",
|
||||
model: faux.getModel(),
|
||||
},
|
||||
});
|
||||
|
||||
await expect(agent.continue()).rejects.toThrow("No messages to continue from");
|
||||
});
|
||||
|
||||
it("throws when last message is assistant", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
const model = faux.getModel();
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "Test",
|
||||
model,
|
||||
},
|
||||
});
|
||||
|
||||
const assistantMessage: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "Hello" }],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
agent.state.messages = [assistantMessage];
|
||||
|
||||
await expect(agent.continue()).rejects.toThrow("Cannot continue from message role: assistant");
|
||||
});
|
||||
});
|
||||
|
||||
describe("continue from user message", () => {
|
||||
it("continues and gets a response when last message is user", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
faux.setResponses([fauxAssistantMessage("HELLO WORLD")]);
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant. Follow instructions exactly.",
|
||||
model: faux.getModel(),
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
const userMessage: UserMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Say exactly: HELLO WORLD" }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
agent.state.messages = [userMessage];
|
||||
|
||||
await agent.continue();
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBe(2);
|
||||
expect(agent.state.messages[0].role).toBe("user");
|
||||
expect(agent.state.messages[1].role).toBe("assistant");
|
||||
|
||||
const assistantMsg = agent.state.messages[1];
|
||||
if (assistantMsg.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(getTextContent(assistantMsg).toUpperCase()).toContain("HELLO WORLD");
|
||||
});
|
||||
});
|
||||
|
||||
describe("continue from tool result", () => {
|
||||
it("continues and processes tool results", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
const model = faux.getModel();
|
||||
faux.setResponses([fauxAssistantMessage("The answer is 8.")]);
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt:
|
||||
"You are a helpful assistant. After getting a calculation result, state the answer clearly.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [calculateTool],
|
||||
},
|
||||
});
|
||||
|
||||
const userMessage: UserMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "What is 5 + 3?" }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const assistantMessage: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [
|
||||
{ type: "text", text: "Let me calculate that." },
|
||||
{ type: "toolCall", id: "calc-1", name: "calculate", arguments: { expression: "5 + 3" } },
|
||||
],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "toolUse",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const toolResult: ToolResultMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: "calc-1",
|
||||
toolName: "calculate",
|
||||
content: [{ type: "text", text: "5 + 3 = 8" }],
|
||||
isError: false,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
agent.state.messages = [userMessage, assistantMessage, toolResult];
|
||||
|
||||
await agent.continue();
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBeGreaterThanOrEqual(4);
|
||||
|
||||
const lastMessage = agent.state.messages[agent.state.messages.length - 1];
|
||||
expect(lastMessage.role).toBe("assistant");
|
||||
if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(getTextContent(lastMessage)).toContain("8");
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,32 +0,0 @@
|
||||
import { type Static, Type } from "@sinclair/typebox";
|
||||
import type { AgentTool, AgentToolResult } from "../../src/types.js";
|
||||
|
||||
export interface CalculateResult extends AgentToolResult<undefined> {
|
||||
content: Array<{ type: "text"; text: string }>;
|
||||
details: undefined;
|
||||
}
|
||||
|
||||
export function calculate(expression: string): CalculateResult {
|
||||
try {
|
||||
const result = new Function(`return ${expression}`)();
|
||||
return { content: [{ type: "text", text: `${expression} = ${result}` }], details: undefined };
|
||||
} catch (e: any) {
|
||||
throw new Error(e.message || String(e));
|
||||
}
|
||||
}
|
||||
|
||||
const calculateSchema = Type.Object({
|
||||
expression: Type.String({ description: "The mathematical expression to evaluate" }),
|
||||
});
|
||||
|
||||
type CalculateParams = Static<typeof calculateSchema>;
|
||||
|
||||
export const calculateTool: AgentTool<typeof calculateSchema, undefined> = {
|
||||
label: "Calculator",
|
||||
name: "calculate",
|
||||
description: "Evaluate mathematical expressions",
|
||||
parameters: calculateSchema,
|
||||
execute: async (_toolCallId: string, args: CalculateParams) => {
|
||||
return calculate(args.expression);
|
||||
},
|
||||
};
|
||||
@@ -1,46 +0,0 @@
|
||||
import { type Static, Type } from "@sinclair/typebox";
|
||||
import type { AgentTool, AgentToolResult } from "../../src/types.js";
|
||||
|
||||
export interface GetCurrentTimeResult extends AgentToolResult<{ utcTimestamp: number }> {}
|
||||
|
||||
export async function getCurrentTime(timezone?: string): Promise<GetCurrentTimeResult> {
|
||||
const date = new Date();
|
||||
if (timezone) {
|
||||
try {
|
||||
const timeStr = date.toLocaleString("en-US", {
|
||||
timeZone: timezone,
|
||||
dateStyle: "full",
|
||||
timeStyle: "long",
|
||||
});
|
||||
return {
|
||||
content: [{ type: "text", text: timeStr }],
|
||||
details: { utcTimestamp: date.getTime() },
|
||||
};
|
||||
} catch (_e) {
|
||||
throw new Error(`Invalid timezone: ${timezone}. Current UTC time: ${date.toISOString()}`);
|
||||
}
|
||||
}
|
||||
const timeStr = date.toLocaleString("en-US", { dateStyle: "full", timeStyle: "long" });
|
||||
return {
|
||||
content: [{ type: "text", text: timeStr }],
|
||||
details: { utcTimestamp: date.getTime() },
|
||||
};
|
||||
}
|
||||
|
||||
const getCurrentTimeSchema = Type.Object({
|
||||
timezone: Type.Optional(
|
||||
Type.String({ description: "Optional timezone (e.g., 'America/New_York', 'Europe/London')" }),
|
||||
),
|
||||
});
|
||||
|
||||
type GetCurrentTimeParams = Static<typeof getCurrentTimeSchema>;
|
||||
|
||||
export const getCurrentTimeTool: AgentTool<typeof getCurrentTimeSchema, { utcTimestamp: number }> = {
|
||||
label: "Current Time",
|
||||
name: "get_current_time",
|
||||
description: "Get the current date and time",
|
||||
parameters: getCurrentTimeSchema,
|
||||
execute: async (_toolCallId: string, args: GetCurrentTimeParams) => {
|
||||
return getCurrentTime(args.timezone);
|
||||
},
|
||||
};
|
||||
@@ -1,9 +0,0 @@
|
||||
{
|
||||
"extends": "../../tsconfig.base.json",
|
||||
"compilerOptions": {
|
||||
"outDir": "./dist",
|
||||
"rootDir": "./src"
|
||||
},
|
||||
"include": ["src/**/*.ts"],
|
||||
"exclude": ["node_modules", "dist", "**/*.d.ts", "src/**/*.d.ts"]
|
||||
}
|
||||
@@ -1,9 +0,0 @@
|
||||
import { defineConfig } from "vitest/config";
|
||||
|
||||
export default defineConfig({
|
||||
test: {
|
||||
globals: true,
|
||||
environment: "node",
|
||||
testTimeout: 30000, // 30 seconds for API calls
|
||||
},
|
||||
});
|
||||
@@ -1 +0,0 @@
|
||||
export * from "./dist/bedrock-provider.js";
|
||||
@@ -1 +0,0 @@
|
||||
export * from "./dist/bedrock-provider.js";
|
||||
@@ -1,114 +0,0 @@
|
||||
{
|
||||
"name": "@mariozechner/pi-ai",
|
||||
"version": "0.67.68",
|
||||
"description": "Unified LLM API with automatic model discovery and provider configuration",
|
||||
"type": "module",
|
||||
"main": "./dist/index.js",
|
||||
"types": "./dist/index.d.ts",
|
||||
"exports": {
|
||||
".": {
|
||||
"types": "./dist/index.d.ts",
|
||||
"import": "./dist/index.js"
|
||||
},
|
||||
"./anthropic": {
|
||||
"types": "./dist/providers/anthropic.d.ts",
|
||||
"import": "./dist/providers/anthropic.js"
|
||||
},
|
||||
"./azure-openai-responses": {
|
||||
"types": "./dist/providers/azure-openai-responses.d.ts",
|
||||
"import": "./dist/providers/azure-openai-responses.js"
|
||||
},
|
||||
"./google": {
|
||||
"types": "./dist/providers/google.d.ts",
|
||||
"import": "./dist/providers/google.js"
|
||||
},
|
||||
"./google-gemini-cli": {
|
||||
"types": "./dist/providers/google-gemini-cli.d.ts",
|
||||
"import": "./dist/providers/google-gemini-cli.js"
|
||||
},
|
||||
"./google-vertex": {
|
||||
"types": "./dist/providers/google-vertex.d.ts",
|
||||
"import": "./dist/providers/google-vertex.js"
|
||||
},
|
||||
"./mistral": {
|
||||
"types": "./dist/providers/mistral.d.ts",
|
||||
"import": "./dist/providers/mistral.js"
|
||||
},
|
||||
"./openai-codex-responses": {
|
||||
"types": "./dist/providers/openai-codex-responses.d.ts",
|
||||
"import": "./dist/providers/openai-codex-responses.js"
|
||||
},
|
||||
"./openai-completions": {
|
||||
"types": "./dist/providers/openai-completions.d.ts",
|
||||
"import": "./dist/providers/openai-completions.js"
|
||||
},
|
||||
"./openai-responses": {
|
||||
"types": "./dist/providers/openai-responses.d.ts",
|
||||
"import": "./dist/providers/openai-responses.js"
|
||||
},
|
||||
"./oauth": {
|
||||
"types": "./dist/oauth.d.ts",
|
||||
"import": "./dist/oauth.js"
|
||||
},
|
||||
"./bedrock-provider": {
|
||||
"types": "./dist/bedrock-provider.d.ts",
|
||||
"import": "./dist/bedrock-provider.js"
|
||||
}
|
||||
},
|
||||
"bin": {
|
||||
"pi-ai": "./dist/cli.js"
|
||||
},
|
||||
"files": [
|
||||
"dist",
|
||||
"README.md"
|
||||
],
|
||||
"scripts": {
|
||||
"clean": "shx rm -rf dist",
|
||||
"generate-models": "npx tsx scripts/generate-models.ts",
|
||||
"build": "npm run generate-models && tsgo -p tsconfig.build.json",
|
||||
"dev": "tsgo -p tsconfig.build.json --watch --preserveWatchOutput",
|
||||
"dev:tsc": "tsgo -p tsconfig.build.json --watch --preserveWatchOutput",
|
||||
"test": "vitest --run",
|
||||
"prepublishOnly": "npm run clean && npm run build"
|
||||
},
|
||||
"dependencies": {
|
||||
"@anthropic-ai/sdk": "^0.90.0",
|
||||
"@aws-sdk/client-bedrock-runtime": "^3.1030.0",
|
||||
"@google/genai": "^1.40.0",
|
||||
"@mistralai/mistralai": "^2.2.0",
|
||||
"@sinclair/typebox": "^0.34.41",
|
||||
"ajv": "^8.17.1",
|
||||
"ajv-formats": "^3.0.1",
|
||||
"chalk": "^5.6.2",
|
||||
"openai": "6.26.0",
|
||||
"partial-json": "^0.1.7",
|
||||
"proxy-agent": "^6.5.0",
|
||||
"undici": "^7.19.1",
|
||||
"zod-to-json-schema": "^3.24.6"
|
||||
},
|
||||
"keywords": [
|
||||
"ai",
|
||||
"llm",
|
||||
"openai",
|
||||
"anthropic",
|
||||
"gemini",
|
||||
"bedrock",
|
||||
"unified",
|
||||
"api"
|
||||
],
|
||||
"author": "Mario Zechner",
|
||||
"license": "MIT",
|
||||
"repository": {
|
||||
"type": "git",
|
||||
"url": "git+https://github.com/badlogic/pi-mono.git",
|
||||
"directory": "packages/ai"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=20.0.0"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@types/node": "^24.3.0",
|
||||
"canvas": "^3.2.0",
|
||||
"vitest": "^3.2.4"
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1,34 +0,0 @@
|
||||
#!/usr/bin/env tsx
|
||||
|
||||
import { createCanvas } from "canvas";
|
||||
import { writeFileSync } from "fs";
|
||||
import { join, dirname } from "path";
|
||||
import { fileURLToPath } from "url";
|
||||
|
||||
const __filename = fileURLToPath(import.meta.url);
|
||||
const __dirname = dirname(__filename);
|
||||
|
||||
// Create a 200x200 canvas
|
||||
const canvas = createCanvas(200, 200);
|
||||
const ctx = canvas.getContext("2d");
|
||||
|
||||
// Fill background with white
|
||||
ctx.fillStyle = "white";
|
||||
ctx.fillRect(0, 0, 200, 200);
|
||||
|
||||
// Draw a red circle in the center
|
||||
ctx.fillStyle = "red";
|
||||
ctx.beginPath();
|
||||
ctx.arc(100, 100, 50, 0, Math.PI * 2);
|
||||
ctx.fill();
|
||||
|
||||
// Save the image
|
||||
const buffer = canvas.toBuffer("image/png");
|
||||
const outputPath = join(__dirname, "..", "test", "data", "red-circle.png");
|
||||
|
||||
// Ensure the directory exists
|
||||
import { mkdirSync } from "fs";
|
||||
mkdirSync(join(__dirname, "..", "test", "data"), { recursive: true });
|
||||
|
||||
writeFileSync(outputPath, buffer);
|
||||
console.log(`Generated test image at: ${outputPath}`);
|
||||
@@ -1,98 +0,0 @@
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessageEventStream,
|
||||
Context,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
} from "./types.js";
|
||||
|
||||
export type ApiStreamFunction = (
|
||||
model: Model<Api>,
|
||||
context: Context,
|
||||
options?: StreamOptions,
|
||||
) => AssistantMessageEventStream;
|
||||
|
||||
export type ApiStreamSimpleFunction = (
|
||||
model: Model<Api>,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
) => AssistantMessageEventStream;
|
||||
|
||||
export interface ApiProvider<TApi extends Api = Api, TOptions extends StreamOptions = StreamOptions> {
|
||||
api: TApi;
|
||||
stream: StreamFunction<TApi, TOptions>;
|
||||
streamSimple: StreamFunction<TApi, SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface ApiProviderInternal {
|
||||
api: Api;
|
||||
stream: ApiStreamFunction;
|
||||
streamSimple: ApiStreamSimpleFunction;
|
||||
}
|
||||
|
||||
type RegisteredApiProvider = {
|
||||
provider: ApiProviderInternal;
|
||||
sourceId?: string;
|
||||
};
|
||||
|
||||
const apiProviderRegistry = new Map<string, RegisteredApiProvider>();
|
||||
|
||||
function wrapStream<TApi extends Api, TOptions extends StreamOptions>(
|
||||
api: TApi,
|
||||
stream: StreamFunction<TApi, TOptions>,
|
||||
): ApiStreamFunction {
|
||||
return (model, context, options) => {
|
||||
if (model.api !== api) {
|
||||
throw new Error(`Mismatched api: ${model.api} expected ${api}`);
|
||||
}
|
||||
return stream(model as Model<TApi>, context, options as TOptions);
|
||||
};
|
||||
}
|
||||
|
||||
function wrapStreamSimple<TApi extends Api>(
|
||||
api: TApi,
|
||||
streamSimple: StreamFunction<TApi, SimpleStreamOptions>,
|
||||
): ApiStreamSimpleFunction {
|
||||
return (model, context, options) => {
|
||||
if (model.api !== api) {
|
||||
throw new Error(`Mismatched api: ${model.api} expected ${api}`);
|
||||
}
|
||||
return streamSimple(model as Model<TApi>, context, options);
|
||||
};
|
||||
}
|
||||
|
||||
export function registerApiProvider<TApi extends Api, TOptions extends StreamOptions>(
|
||||
provider: ApiProvider<TApi, TOptions>,
|
||||
sourceId?: string,
|
||||
): void {
|
||||
apiProviderRegistry.set(provider.api, {
|
||||
provider: {
|
||||
api: provider.api,
|
||||
stream: wrapStream(provider.api, provider.stream),
|
||||
streamSimple: wrapStreamSimple(provider.api, provider.streamSimple),
|
||||
},
|
||||
sourceId,
|
||||
});
|
||||
}
|
||||
|
||||
export function getApiProvider(api: Api): ApiProviderInternal | undefined {
|
||||
return apiProviderRegistry.get(api)?.provider;
|
||||
}
|
||||
|
||||
export function getApiProviders(): ApiProviderInternal[] {
|
||||
return Array.from(apiProviderRegistry.values(), (entry) => entry.provider);
|
||||
}
|
||||
|
||||
export function unregisterApiProviders(sourceId: string): void {
|
||||
for (const [api, entry] of apiProviderRegistry.entries()) {
|
||||
if (entry.sourceId === sourceId) {
|
||||
apiProviderRegistry.delete(api);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export function clearApiProviders(): void {
|
||||
apiProviderRegistry.clear();
|
||||
}
|
||||
@@ -1,6 +0,0 @@
|
||||
import { streamBedrock, streamSimpleBedrock } from "./providers/amazon-bedrock.js";
|
||||
|
||||
export const bedrockProviderModule = {
|
||||
streamBedrock,
|
||||
streamSimpleBedrock,
|
||||
};
|
||||
@@ -1,133 +0,0 @@
|
||||
#!/usr/bin/env node
|
||||
|
||||
import { createInterface } from "node:readline";
|
||||
import { existsSync, readFileSync, writeFileSync } from "fs";
|
||||
import { getOAuthProvider, getOAuthProviders } from "./utils/oauth/index.js";
|
||||
import type { OAuthCredentials, OAuthProviderId } from "./utils/oauth/types.js";
|
||||
|
||||
const AUTH_FILE = "auth.json";
|
||||
const PROVIDERS = getOAuthProviders();
|
||||
|
||||
function prompt(rl: ReturnType<typeof createInterface>, question: string): Promise<string> {
|
||||
return new Promise((resolve) => rl.question(question, resolve));
|
||||
}
|
||||
|
||||
function loadAuth(): Record<string, { type: "oauth" } & OAuthCredentials> {
|
||||
if (!existsSync(AUTH_FILE)) return {};
|
||||
try {
|
||||
return JSON.parse(readFileSync(AUTH_FILE, "utf-8"));
|
||||
} catch {
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
function saveAuth(auth: Record<string, { type: "oauth" } & OAuthCredentials>): void {
|
||||
writeFileSync(AUTH_FILE, JSON.stringify(auth, null, 2), "utf-8");
|
||||
}
|
||||
|
||||
async function login(providerId: OAuthProviderId): Promise<void> {
|
||||
const provider = getOAuthProvider(providerId);
|
||||
if (!provider) {
|
||||
console.error(`Unknown provider: ${providerId}`);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const rl = createInterface({ input: process.stdin, output: process.stdout });
|
||||
const promptFn = (msg: string) => prompt(rl, `${msg} `);
|
||||
|
||||
try {
|
||||
const credentials = await provider.login({
|
||||
onAuth: (info) => {
|
||||
console.log(`\nOpen this URL in your browser:\n${info.url}`);
|
||||
if (info.instructions) console.log(info.instructions);
|
||||
console.log();
|
||||
},
|
||||
onPrompt: async (p) => {
|
||||
return await promptFn(`${p.message}${p.placeholder ? ` (${p.placeholder})` : ""}:`);
|
||||
},
|
||||
onProgress: (msg) => console.log(msg),
|
||||
});
|
||||
|
||||
const auth = loadAuth();
|
||||
auth[providerId] = { type: "oauth", ...credentials };
|
||||
saveAuth(auth);
|
||||
|
||||
console.log(`\nCredentials saved to ${AUTH_FILE}`);
|
||||
} finally {
|
||||
rl.close();
|
||||
}
|
||||
}
|
||||
|
||||
async function main(): Promise<void> {
|
||||
const args = process.argv.slice(2);
|
||||
const command = args[0];
|
||||
|
||||
if (!command || command === "help" || command === "--help" || command === "-h") {
|
||||
const providerList = PROVIDERS.map((p) => ` ${p.id.padEnd(20)} ${p.name}`).join("\n");
|
||||
console.log(`Usage: npx @mariozechner/pi-ai <command> [provider]
|
||||
|
||||
Commands:
|
||||
login [provider] Login to an OAuth provider
|
||||
list List available providers
|
||||
|
||||
Providers:
|
||||
${providerList}
|
||||
|
||||
Examples:
|
||||
npx @mariozechner/pi-ai login # interactive provider selection
|
||||
npx @mariozechner/pi-ai login anthropic # login to specific provider
|
||||
npx @mariozechner/pi-ai list # list providers
|
||||
`);
|
||||
return;
|
||||
}
|
||||
|
||||
if (command === "list") {
|
||||
console.log("Available OAuth providers:\n");
|
||||
for (const p of PROVIDERS) {
|
||||
console.log(` ${p.id.padEnd(20)} ${p.name}`);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (command === "login") {
|
||||
let provider = args[1] as OAuthProviderId | undefined;
|
||||
|
||||
if (!provider) {
|
||||
const rl = createInterface({ input: process.stdin, output: process.stdout });
|
||||
console.log("Select a provider:\n");
|
||||
for (let i = 0; i < PROVIDERS.length; i++) {
|
||||
console.log(` ${i + 1}. ${PROVIDERS[i].name}`);
|
||||
}
|
||||
console.log();
|
||||
|
||||
const choice = await prompt(rl, `Enter number (1-${PROVIDERS.length}): `);
|
||||
rl.close();
|
||||
|
||||
const index = parseInt(choice, 10) - 1;
|
||||
if (index < 0 || index >= PROVIDERS.length) {
|
||||
console.error("Invalid selection");
|
||||
process.exit(1);
|
||||
}
|
||||
provider = PROVIDERS[index].id;
|
||||
}
|
||||
|
||||
if (!PROVIDERS.some((p) => p.id === provider)) {
|
||||
console.error(`Unknown provider: ${provider}`);
|
||||
console.error(`Use 'npx @mariozechner/pi-ai list' to see available providers`);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
console.log(`Logging in to ${provider}...`);
|
||||
await login(provider);
|
||||
return;
|
||||
}
|
||||
|
||||
console.error(`Unknown command: ${command}`);
|
||||
console.error(`Use 'npx @mariozechner/pi-ai --help' for usage`);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
main().catch((err) => {
|
||||
console.error("Error:", err.message);
|
||||
process.exit(1);
|
||||
});
|
||||
@@ -1,133 +0,0 @@
|
||||
// NEVER convert to top-level imports - breaks browser/Vite builds (web-ui)
|
||||
let _existsSync: typeof import("node:fs").existsSync | null = null;
|
||||
let _homedir: typeof import("node:os").homedir | null = null;
|
||||
let _join: typeof import("node:path").join | null = null;
|
||||
|
||||
type DynamicImport = (specifier: string) => Promise<unknown>;
|
||||
|
||||
const dynamicImport: DynamicImport = (specifier) => import(specifier);
|
||||
const NODE_FS_SPECIFIER = "node:" + "fs";
|
||||
const NODE_OS_SPECIFIER = "node:" + "os";
|
||||
const NODE_PATH_SPECIFIER = "node:" + "path";
|
||||
|
||||
// Eagerly load in Node.js/Bun environment only
|
||||
if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) {
|
||||
dynamicImport(NODE_FS_SPECIFIER).then((m) => {
|
||||
_existsSync = (m as typeof import("node:fs")).existsSync;
|
||||
});
|
||||
dynamicImport(NODE_OS_SPECIFIER).then((m) => {
|
||||
_homedir = (m as typeof import("node:os")).homedir;
|
||||
});
|
||||
dynamicImport(NODE_PATH_SPECIFIER).then((m) => {
|
||||
_join = (m as typeof import("node:path")).join;
|
||||
});
|
||||
}
|
||||
|
||||
import type { KnownProvider } from "./types.js";
|
||||
|
||||
let cachedVertexAdcCredentialsExists: boolean | null = null;
|
||||
|
||||
function hasVertexAdcCredentials(): boolean {
|
||||
if (cachedVertexAdcCredentialsExists === null) {
|
||||
// If node modules haven't loaded yet (async import race at startup),
|
||||
// return false WITHOUT caching so the next call retries once they're ready.
|
||||
// Only cache false permanently in a browser environment where fs is never available.
|
||||
if (!_existsSync || !_homedir || !_join) {
|
||||
const isNode = typeof process !== "undefined" && (process.versions?.node || process.versions?.bun);
|
||||
if (!isNode) {
|
||||
// Definitively in a browser — safe to cache false permanently
|
||||
cachedVertexAdcCredentialsExists = false;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check GOOGLE_APPLICATION_CREDENTIALS env var first (standard way)
|
||||
const gacPath = process.env.GOOGLE_APPLICATION_CREDENTIALS;
|
||||
if (gacPath) {
|
||||
cachedVertexAdcCredentialsExists = _existsSync(gacPath);
|
||||
} else {
|
||||
// Fall back to default ADC path (lazy evaluation)
|
||||
cachedVertexAdcCredentialsExists = _existsSync(
|
||||
_join(_homedir(), ".config", "gcloud", "application_default_credentials.json"),
|
||||
);
|
||||
}
|
||||
}
|
||||
return cachedVertexAdcCredentialsExists;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get API key for provider from known environment variables, e.g. OPENAI_API_KEY.
|
||||
*
|
||||
* Will not return API keys for providers that require OAuth tokens.
|
||||
*/
|
||||
export function getEnvApiKey(provider: KnownProvider): string | undefined;
|
||||
export function getEnvApiKey(provider: string): string | undefined;
|
||||
export function getEnvApiKey(provider: any): string | undefined {
|
||||
// Fall back to environment variables
|
||||
if (provider === "github-copilot") {
|
||||
return process.env.COPILOT_GITHUB_TOKEN || process.env.GH_TOKEN || process.env.GITHUB_TOKEN;
|
||||
}
|
||||
|
||||
// ANTHROPIC_OAUTH_TOKEN takes precedence over ANTHROPIC_API_KEY
|
||||
if (provider === "anthropic") {
|
||||
return process.env.ANTHROPIC_OAUTH_TOKEN || process.env.ANTHROPIC_API_KEY;
|
||||
}
|
||||
|
||||
// Vertex AI supports either an explicit API key or Application Default Credentials
|
||||
// Auth is configured via `gcloud auth application-default login`
|
||||
if (provider === "google-vertex") {
|
||||
if (process.env.GOOGLE_CLOUD_API_KEY) {
|
||||
return process.env.GOOGLE_CLOUD_API_KEY;
|
||||
}
|
||||
|
||||
const hasCredentials = hasVertexAdcCredentials();
|
||||
const hasProject = !!(process.env.GOOGLE_CLOUD_PROJECT || process.env.GCLOUD_PROJECT);
|
||||
const hasLocation = !!process.env.GOOGLE_CLOUD_LOCATION;
|
||||
|
||||
if (hasCredentials && hasProject && hasLocation) {
|
||||
return "<authenticated>";
|
||||
}
|
||||
}
|
||||
|
||||
if (provider === "amazon-bedrock") {
|
||||
// Amazon Bedrock supports multiple credential sources:
|
||||
// 1. AWS_PROFILE - named profile from ~/.aws/credentials
|
||||
// 2. AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY - standard IAM keys
|
||||
// 3. AWS_BEARER_TOKEN_BEDROCK - Bedrock API keys (bearer token)
|
||||
// 4. AWS_CONTAINER_CREDENTIALS_RELATIVE_URI - ECS task roles
|
||||
// 5. AWS_CONTAINER_CREDENTIALS_FULL_URI - ECS task roles (full URI)
|
||||
// 6. AWS_WEB_IDENTITY_TOKEN_FILE - IRSA (IAM Roles for Service Accounts)
|
||||
if (
|
||||
process.env.AWS_PROFILE ||
|
||||
(process.env.AWS_ACCESS_KEY_ID && process.env.AWS_SECRET_ACCESS_KEY) ||
|
||||
process.env.AWS_BEARER_TOKEN_BEDROCK ||
|
||||
process.env.AWS_CONTAINER_CREDENTIALS_RELATIVE_URI ||
|
||||
process.env.AWS_CONTAINER_CREDENTIALS_FULL_URI ||
|
||||
process.env.AWS_WEB_IDENTITY_TOKEN_FILE
|
||||
) {
|
||||
return "<authenticated>";
|
||||
}
|
||||
}
|
||||
|
||||
const envMap: Record<string, string> = {
|
||||
openai: "OPENAI_API_KEY",
|
||||
"azure-openai-responses": "AZURE_OPENAI_API_KEY",
|
||||
google: "GEMINI_API_KEY",
|
||||
groq: "GROQ_API_KEY",
|
||||
cerebras: "CEREBRAS_API_KEY",
|
||||
xai: "XAI_API_KEY",
|
||||
openrouter: "OPENROUTER_API_KEY",
|
||||
"vercel-ai-gateway": "AI_GATEWAY_API_KEY",
|
||||
zai: "ZAI_API_KEY",
|
||||
mistral: "MISTRAL_API_KEY",
|
||||
minimax: "MINIMAX_API_KEY",
|
||||
"minimax-cn": "MINIMAX_CN_API_KEY",
|
||||
huggingface: "HF_TOKEN",
|
||||
opencode: "OPENCODE_API_KEY",
|
||||
"opencode-go": "OPENCODE_API_KEY",
|
||||
"kimi-coding": "KIMI_API_KEY",
|
||||
};
|
||||
|
||||
const envVar = envMap[provider];
|
||||
return envVar ? process.env[envVar] : undefined;
|
||||
}
|
||||
@@ -1,35 +0,0 @@
|
||||
export type { Static, TSchema } from "@sinclair/typebox";
|
||||
export { Type } from "@sinclair/typebox";
|
||||
|
||||
export * from "./api-registry.js";
|
||||
export * from "./env-api-keys.js";
|
||||
export * from "./models.js";
|
||||
export type { BedrockOptions, BedrockThinkingDisplay } from "./providers/amazon-bedrock.js";
|
||||
export type { AnthropicEffort, AnthropicOptions, AnthropicThinkingDisplay } from "./providers/anthropic.js";
|
||||
export type { AzureOpenAIResponsesOptions } from "./providers/azure-openai-responses.js";
|
||||
export * from "./providers/faux.js";
|
||||
export type { GoogleOptions } from "./providers/google.js";
|
||||
export type { GoogleGeminiCliOptions, GoogleThinkingLevel } from "./providers/google-gemini-cli.js";
|
||||
export type { GoogleVertexOptions } from "./providers/google-vertex.js";
|
||||
export type { MistralOptions } from "./providers/mistral.js";
|
||||
export type { OpenAICodexResponsesOptions } from "./providers/openai-codex-responses.js";
|
||||
export type { OpenAICompletionsOptions } from "./providers/openai-completions.js";
|
||||
export type { OpenAIResponsesOptions } from "./providers/openai-responses.js";
|
||||
export * from "./providers/register-builtins.js";
|
||||
export * from "./stream.js";
|
||||
export * from "./types.js";
|
||||
export * from "./utils/event-stream.js";
|
||||
export * from "./utils/json-parse.js";
|
||||
export type {
|
||||
OAuthAuthInfo,
|
||||
OAuthCredentials,
|
||||
OAuthLoginCallbacks,
|
||||
OAuthPrompt,
|
||||
OAuthProvider,
|
||||
OAuthProviderId,
|
||||
OAuthProviderInfo,
|
||||
OAuthProviderInterface,
|
||||
} from "./utils/oauth/types.js";
|
||||
export * from "./utils/overflow.js";
|
||||
export * from "./utils/typebox-helpers.js";
|
||||
export * from "./utils/validation.js";
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1,82 +0,0 @@
|
||||
import { MODELS } from "./models.generated.js";
|
||||
import type { Api, KnownProvider, Model, Usage } from "./types.js";
|
||||
|
||||
const modelRegistry: Map<string, Map<string, Model<Api>>> = new Map();
|
||||
|
||||
// Initialize registry from MODELS on module load
|
||||
for (const [provider, models] of Object.entries(MODELS)) {
|
||||
const providerModels = new Map<string, Model<Api>>();
|
||||
for (const [id, model] of Object.entries(models)) {
|
||||
providerModels.set(id, model as Model<Api>);
|
||||
}
|
||||
modelRegistry.set(provider, providerModels);
|
||||
}
|
||||
|
||||
type ModelApi<
|
||||
TProvider extends KnownProvider,
|
||||
TModelId extends keyof (typeof MODELS)[TProvider],
|
||||
> = (typeof MODELS)[TProvider][TModelId] extends { api: infer TApi } ? (TApi extends Api ? TApi : never) : never;
|
||||
|
||||
export function getModel<TProvider extends KnownProvider, TModelId extends keyof (typeof MODELS)[TProvider]>(
|
||||
provider: TProvider,
|
||||
modelId: TModelId,
|
||||
): Model<ModelApi<TProvider, TModelId>> {
|
||||
const providerModels = modelRegistry.get(provider);
|
||||
return providerModels?.get(modelId as string) as Model<ModelApi<TProvider, TModelId>>;
|
||||
}
|
||||
|
||||
export function getProviders(): KnownProvider[] {
|
||||
return Array.from(modelRegistry.keys()) as KnownProvider[];
|
||||
}
|
||||
|
||||
export function getModels<TProvider extends KnownProvider>(
|
||||
provider: TProvider,
|
||||
): Model<ModelApi<TProvider, keyof (typeof MODELS)[TProvider]>>[] {
|
||||
const models = modelRegistry.get(provider);
|
||||
return models ? (Array.from(models.values()) as Model<ModelApi<TProvider, keyof (typeof MODELS)[TProvider]>>[]) : [];
|
||||
}
|
||||
|
||||
export function calculateCost<TApi extends Api>(model: Model<TApi>, usage: Usage): Usage["cost"] {
|
||||
usage.cost.input = (model.cost.input / 1000000) * usage.input;
|
||||
usage.cost.output = (model.cost.output / 1000000) * usage.output;
|
||||
usage.cost.cacheRead = (model.cost.cacheRead / 1000000) * usage.cacheRead;
|
||||
usage.cost.cacheWrite = (model.cost.cacheWrite / 1000000) * usage.cacheWrite;
|
||||
usage.cost.total = usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite;
|
||||
return usage.cost;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a model supports xhigh thinking level.
|
||||
*
|
||||
* Supported today:
|
||||
* - GPT-5.2 / GPT-5.3 / GPT-5.4 model families
|
||||
* - Opus 4.6+ models (xhigh maps to adaptive effort "max" on Anthropic-compatible providers)
|
||||
*/
|
||||
export function supportsXhigh<TApi extends Api>(model: Model<TApi>): boolean {
|
||||
if (model.id.includes("gpt-5.2") || model.id.includes("gpt-5.3") || model.id.includes("gpt-5.4")) {
|
||||
return true;
|
||||
}
|
||||
|
||||
if (
|
||||
model.id.includes("opus-4-6") ||
|
||||
model.id.includes("opus-4.6") ||
|
||||
model.id.includes("opus-4-7") ||
|
||||
model.id.includes("opus-4.7")
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if two models are equal by comparing both their id and provider.
|
||||
* Returns false if either model is null or undefined.
|
||||
*/
|
||||
export function modelsAreEqual<TApi extends Api>(
|
||||
a: Model<TApi> | null | undefined,
|
||||
b: Model<TApi> | null | undefined,
|
||||
): boolean {
|
||||
if (!a || !b) return false;
|
||||
return a.id === b.id && a.provider === b.provider;
|
||||
}
|
||||
@@ -1 +0,0 @@
|
||||
export * from "./utils/oauth/index.js";
|
||||
@@ -1,865 +0,0 @@
|
||||
import {
|
||||
BedrockRuntimeClient,
|
||||
type BedrockRuntimeClientConfig,
|
||||
BedrockRuntimeServiceException,
|
||||
StopReason as BedrockStopReason,
|
||||
type Tool as BedrockTool,
|
||||
CachePointType,
|
||||
CacheTTL,
|
||||
type ContentBlock,
|
||||
type ContentBlockDeltaEvent,
|
||||
type ContentBlockStartEvent,
|
||||
type ContentBlockStopEvent,
|
||||
ConversationRole,
|
||||
ConverseStreamCommand,
|
||||
type ConverseStreamMetadataEvent,
|
||||
ImageFormat,
|
||||
type Message,
|
||||
type SystemContentBlock,
|
||||
type ToolChoice,
|
||||
type ToolConfiguration,
|
||||
ToolResultStatus,
|
||||
} from "@aws-sdk/client-bedrock-runtime";
|
||||
import { calculateCost } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
CacheRetention,
|
||||
Context,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StopReason,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
TextContent,
|
||||
ThinkingBudgets,
|
||||
ThinkingContent,
|
||||
ThinkingLevel,
|
||||
Tool,
|
||||
ToolCall,
|
||||
ToolResultMessage,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { parseStreamingJson } from "../utils/json-parse.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import { adjustMaxTokensForThinking, buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
|
||||
export type BedrockThinkingDisplay = "summarized" | "omitted";
|
||||
|
||||
export interface BedrockOptions extends StreamOptions {
|
||||
region?: string;
|
||||
profile?: string;
|
||||
toolChoice?: "auto" | "any" | "none" | { type: "tool"; name: string };
|
||||
/* See https://docs.aws.amazon.com/bedrock/latest/userguide/inference-reasoning.html for supported models. */
|
||||
reasoning?: ThinkingLevel;
|
||||
/* Custom token budgets per thinking level. Overrides default budgets. */
|
||||
thinkingBudgets?: ThinkingBudgets;
|
||||
/* Only supported by Claude 4.x models, see https://docs.aws.amazon.com/bedrock/latest/userguide/claude-messages-extended-thinking.html#claude-messages-extended-thinking-tool-use-interleaved */
|
||||
interleavedThinking?: boolean;
|
||||
/**
|
||||
* Controls how Claude's thinking content is returned in responses.
|
||||
* - "summarized": Thinking blocks contain summarized thinking text (default here).
|
||||
* - "omitted": Thinking content is redacted but the signature still travels back
|
||||
* for multi-turn continuity, reducing time-to-first-text-token.
|
||||
*
|
||||
* Note: Anthropic's API default for Claude Opus 4.7 and Mythos Preview is
|
||||
* "omitted". We default to "summarized" here to keep behavior consistent with
|
||||
* older Claude 4 models. Only applies to Claude models on Bedrock.
|
||||
*/
|
||||
thinkingDisplay?: BedrockThinkingDisplay;
|
||||
/** Key-value pairs attached to the inference request for cost allocation tagging.
|
||||
* Keys: max 64 chars, no `aws:` prefix. Values: max 256 chars. Max 50 pairs.
|
||||
* Tags appear in AWS Cost Explorer split cost allocation data.
|
||||
* @see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseStream.html */
|
||||
requestMetadata?: Record<string, string>;
|
||||
/** Bearer token for Bedrock API key authentication.
|
||||
* When set, bypasses SigV4 signing and sends Authorization: Bearer <token> instead.
|
||||
* Requires `bedrock:CallWithBearerToken` IAM permission on the token's identity.
|
||||
* Set via AWS_BEARER_TOKEN_BEDROCK env var or pass directly.
|
||||
* @see https://docs.aws.amazon.com/service-authorization/latest/reference/list_amazonbedrock.html */
|
||||
bearerToken?: string;
|
||||
}
|
||||
|
||||
type Block = (TextContent | ThinkingContent | ToolCall) & { index?: number; partialJson?: string };
|
||||
|
||||
export const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOptions> = (
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
context: Context,
|
||||
options: BedrockOptions = {},
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: "bedrock-converse-stream" as Api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const blocks = output.content as Block[];
|
||||
|
||||
const config: BedrockRuntimeClientConfig = {
|
||||
profile: options.profile,
|
||||
};
|
||||
|
||||
// Resolve bearer token for Bedrock API key auth.
|
||||
const bearerToken = options.bearerToken || process.env.AWS_BEARER_TOKEN_BEDROCK || undefined;
|
||||
const useBearerToken = bearerToken !== undefined && process.env.AWS_BEDROCK_SKIP_AUTH !== "1";
|
||||
|
||||
// in Node.js/Bun environment only
|
||||
if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) {
|
||||
// Region resolution: explicit option > env vars > SDK default chain.
|
||||
// When AWS_PROFILE is set, we leave region undefined so the SDK can
|
||||
// resovle it from aws profile configs. Otherwise fall back to us-east-1.
|
||||
const explicitRegion = options.region || process.env.AWS_REGION || process.env.AWS_DEFAULT_REGION;
|
||||
if (explicitRegion) {
|
||||
config.region = explicitRegion;
|
||||
} else if (!process.env.AWS_PROFILE) {
|
||||
config.region = "us-east-1";
|
||||
}
|
||||
|
||||
// Support proxies that don't need authentication
|
||||
if (process.env.AWS_BEDROCK_SKIP_AUTH === "1") {
|
||||
config.credentials = {
|
||||
accessKeyId: "dummy-access-key",
|
||||
secretAccessKey: "dummy-secret-key",
|
||||
};
|
||||
}
|
||||
|
||||
if (
|
||||
process.env.HTTP_PROXY ||
|
||||
process.env.HTTPS_PROXY ||
|
||||
process.env.NO_PROXY ||
|
||||
process.env.http_proxy ||
|
||||
process.env.https_proxy ||
|
||||
process.env.no_proxy
|
||||
) {
|
||||
const nodeHttpHandler = await import("@smithy/node-http-handler");
|
||||
const proxyAgent = await import("proxy-agent");
|
||||
|
||||
const agent = new proxyAgent.ProxyAgent();
|
||||
|
||||
// Bedrock runtime uses NodeHttp2Handler by default since v3.798.0, which is based
|
||||
// on `http2` module and has no support for http agent.
|
||||
// Use NodeHttpHandler to support http agent.
|
||||
config.requestHandler = new nodeHttpHandler.NodeHttpHandler({
|
||||
httpAgent: agent,
|
||||
httpsAgent: agent,
|
||||
});
|
||||
} else if (process.env.AWS_BEDROCK_FORCE_HTTP1 === "1") {
|
||||
// Some custom endpoints require HTTP/1.1 instead of HTTP/2
|
||||
const nodeHttpHandler = await import("@smithy/node-http-handler");
|
||||
config.requestHandler = new nodeHttpHandler.NodeHttpHandler();
|
||||
}
|
||||
} else {
|
||||
// Non-Node environment (browser): fall back to us-east-1 since
|
||||
// there's no config file resolution available.
|
||||
config.region = options.region || "us-east-1";
|
||||
}
|
||||
|
||||
if (useBearerToken) {
|
||||
config.token = { token: bearerToken };
|
||||
config.authSchemePreference = ["httpBearerAuth"];
|
||||
}
|
||||
|
||||
try {
|
||||
const client = new BedrockRuntimeClient(config);
|
||||
const cacheRetention = resolveCacheRetention(options.cacheRetention);
|
||||
let commandInput = {
|
||||
modelId: model.id,
|
||||
messages: convertMessages(context, model, cacheRetention),
|
||||
system: buildSystemPrompt(context.systemPrompt, model, cacheRetention),
|
||||
inferenceConfig: { maxTokens: options.maxTokens, temperature: options.temperature },
|
||||
toolConfig: convertToolConfig(context.tools, options.toolChoice),
|
||||
additionalModelRequestFields: buildAdditionalModelRequestFields(model, options),
|
||||
...(options.requestMetadata !== undefined && { requestMetadata: options.requestMetadata }),
|
||||
};
|
||||
const nextCommandInput = await options?.onPayload?.(commandInput, model);
|
||||
if (nextCommandInput !== undefined) {
|
||||
commandInput = nextCommandInput as typeof commandInput;
|
||||
}
|
||||
const command = new ConverseStreamCommand(commandInput);
|
||||
|
||||
const response = await client.send(command, { abortSignal: options.signal });
|
||||
if (response.$metadata.httpStatusCode !== undefined) {
|
||||
const responseHeaders: Record<string, string> = {};
|
||||
if (response.$metadata.requestId) {
|
||||
responseHeaders["x-amzn-requestid"] = response.$metadata.requestId;
|
||||
}
|
||||
await options?.onResponse?.({ status: response.$metadata.httpStatusCode, headers: responseHeaders }, model);
|
||||
}
|
||||
|
||||
for await (const item of response.stream!) {
|
||||
if (item.messageStart) {
|
||||
if (item.messageStart.role !== ConversationRole.ASSISTANT) {
|
||||
throw new Error("Unexpected assistant message start but got user message start instead");
|
||||
}
|
||||
stream.push({ type: "start", partial: output });
|
||||
} else if (item.contentBlockStart) {
|
||||
handleContentBlockStart(item.contentBlockStart, blocks, output, stream);
|
||||
} else if (item.contentBlockDelta) {
|
||||
handleContentBlockDelta(item.contentBlockDelta, blocks, output, stream);
|
||||
} else if (item.contentBlockStop) {
|
||||
handleContentBlockStop(item.contentBlockStop, blocks, output, stream);
|
||||
} else if (item.messageStop) {
|
||||
output.stopReason = mapStopReason(item.messageStop.stopReason);
|
||||
} else if (item.metadata) {
|
||||
handleMetadata(item.metadata, model, output);
|
||||
} else if (item.internalServerException) {
|
||||
throw item.internalServerException;
|
||||
} else if (item.modelStreamErrorException) {
|
||||
throw item.modelStreamErrorException;
|
||||
} else if (item.validationException) {
|
||||
throw item.validationException;
|
||||
} else if (item.throttlingException) {
|
||||
throw item.throttlingException;
|
||||
} else if (item.serviceUnavailableException) {
|
||||
throw item.serviceUnavailableException;
|
||||
}
|
||||
}
|
||||
|
||||
if (options.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "error" || output.stopReason === "aborted") {
|
||||
throw new Error("An unknown error occurred");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
for (const block of output.content) {
|
||||
delete (block as Block).index;
|
||||
// partialJson is only a streaming scratch buffer; never persist it.
|
||||
delete (block as Block).partialJson;
|
||||
}
|
||||
output.stopReason = options.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = formatBedrockError(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
/**
|
||||
* Human-readable prefixes for Bedrock SDK exception names.
|
||||
* The downstream retry logic in agent-session matches patterns like
|
||||
* `server.?error` and `service.?unavailable`, so we preserve the legacy
|
||||
* prefix format rather than using the raw SDK exception name.
|
||||
*/
|
||||
const BEDROCK_ERROR_PREFIXES: Record<string, string> = {
|
||||
InternalServerException: "Internal server error",
|
||||
ModelStreamErrorException: "Model stream error",
|
||||
ValidationException: "Validation error",
|
||||
ThrottlingException: "Throttling error",
|
||||
ServiceUnavailableException: "Service unavailable",
|
||||
};
|
||||
|
||||
/**
|
||||
* Format a Bedrock error with a human-readable prefix.
|
||||
* AWS SDK exceptions (both from `client.send()` and from stream event items)
|
||||
* extend BedrockRuntimeServiceException. We map the `.name` to a stable
|
||||
* human-readable prefix so downstream consumers (retry logic, context-overflow
|
||||
* detection) can distinguish error categories via simple string matching.
|
||||
*/
|
||||
function formatBedrockError(error: unknown): string {
|
||||
const message = error instanceof Error ? error.message : JSON.stringify(error);
|
||||
if (error instanceof BedrockRuntimeServiceException) {
|
||||
const prefix = BEDROCK_ERROR_PREFIXES[error.name] ?? error.name;
|
||||
return `${prefix}: ${message}`;
|
||||
}
|
||||
return message;
|
||||
}
|
||||
|
||||
export const streamSimpleBedrock: StreamFunction<"bedrock-converse-stream", SimpleStreamOptions> = (
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const base = buildBaseOptions(model, options, undefined);
|
||||
if (!options?.reasoning) {
|
||||
return streamBedrock(model, context, { ...base, reasoning: undefined } satisfies BedrockOptions);
|
||||
}
|
||||
|
||||
if (model.id.includes("anthropic.claude") || model.id.includes("anthropic/claude")) {
|
||||
if (supportsAdaptiveThinking(model.id)) {
|
||||
return streamBedrock(model, context, {
|
||||
...base,
|
||||
reasoning: options.reasoning,
|
||||
thinkingBudgets: options.thinkingBudgets,
|
||||
} satisfies BedrockOptions);
|
||||
}
|
||||
|
||||
const adjusted = adjustMaxTokensForThinking(
|
||||
base.maxTokens || 0,
|
||||
model.maxTokens,
|
||||
options.reasoning,
|
||||
options.thinkingBudgets,
|
||||
);
|
||||
|
||||
return streamBedrock(model, context, {
|
||||
...base,
|
||||
maxTokens: adjusted.maxTokens,
|
||||
reasoning: options.reasoning,
|
||||
thinkingBudgets: {
|
||||
...(options.thinkingBudgets || {}),
|
||||
[clampReasoning(options.reasoning)!]: adjusted.thinkingBudget,
|
||||
},
|
||||
} satisfies BedrockOptions);
|
||||
}
|
||||
|
||||
return streamBedrock(model, context, {
|
||||
...base,
|
||||
reasoning: options.reasoning,
|
||||
thinkingBudgets: options.thinkingBudgets,
|
||||
} satisfies BedrockOptions);
|
||||
};
|
||||
|
||||
function handleContentBlockStart(
|
||||
event: ContentBlockStartEvent,
|
||||
blocks: Block[],
|
||||
output: AssistantMessage,
|
||||
stream: AssistantMessageEventStream,
|
||||
): void {
|
||||
const index = event.contentBlockIndex!;
|
||||
const start = event.start;
|
||||
|
||||
if (start?.toolUse) {
|
||||
const block: Block = {
|
||||
type: "toolCall",
|
||||
id: start.toolUse.toolUseId || "",
|
||||
name: start.toolUse.name || "",
|
||||
arguments: {},
|
||||
partialJson: "",
|
||||
index,
|
||||
};
|
||||
output.content.push(block);
|
||||
stream.push({ type: "toolcall_start", contentIndex: blocks.length - 1, partial: output });
|
||||
}
|
||||
}
|
||||
|
||||
function handleContentBlockDelta(
|
||||
event: ContentBlockDeltaEvent,
|
||||
blocks: Block[],
|
||||
output: AssistantMessage,
|
||||
stream: AssistantMessageEventStream,
|
||||
): void {
|
||||
const contentBlockIndex = event.contentBlockIndex!;
|
||||
const delta = event.delta;
|
||||
let index = blocks.findIndex((b) => b.index === contentBlockIndex);
|
||||
let block = blocks[index];
|
||||
|
||||
if (delta?.text !== undefined) {
|
||||
// If no text block exists yet, create one, as `handleContentBlockStart` is not sent for text blocks
|
||||
if (!block) {
|
||||
const newBlock: Block = { type: "text", text: "", index: contentBlockIndex };
|
||||
output.content.push(newBlock);
|
||||
index = blocks.length - 1;
|
||||
block = blocks[index];
|
||||
stream.push({ type: "text_start", contentIndex: index, partial: output });
|
||||
}
|
||||
if (block.type === "text") {
|
||||
block.text += delta.text;
|
||||
stream.push({ type: "text_delta", contentIndex: index, delta: delta.text, partial: output });
|
||||
}
|
||||
} else if (delta?.toolUse && block?.type === "toolCall") {
|
||||
block.partialJson = (block.partialJson || "") + (delta.toolUse.input || "");
|
||||
block.arguments = parseStreamingJson(block.partialJson);
|
||||
stream.push({ type: "toolcall_delta", contentIndex: index, delta: delta.toolUse.input || "", partial: output });
|
||||
} else if (delta?.reasoningContent) {
|
||||
let thinkingBlock = block;
|
||||
let thinkingIndex = index;
|
||||
|
||||
if (!thinkingBlock) {
|
||||
const newBlock: Block = { type: "thinking", thinking: "", thinkingSignature: "", index: contentBlockIndex };
|
||||
output.content.push(newBlock);
|
||||
thinkingIndex = blocks.length - 1;
|
||||
thinkingBlock = blocks[thinkingIndex];
|
||||
stream.push({ type: "thinking_start", contentIndex: thinkingIndex, partial: output });
|
||||
}
|
||||
|
||||
if (thinkingBlock?.type === "thinking") {
|
||||
if (delta.reasoningContent.text) {
|
||||
thinkingBlock.thinking += delta.reasoningContent.text;
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: thinkingIndex,
|
||||
delta: delta.reasoningContent.text,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
if (delta.reasoningContent.signature) {
|
||||
thinkingBlock.thinkingSignature =
|
||||
(thinkingBlock.thinkingSignature || "") + delta.reasoningContent.signature;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function handleMetadata(
|
||||
event: ConverseStreamMetadataEvent,
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
output: AssistantMessage,
|
||||
): void {
|
||||
if (event.usage) {
|
||||
output.usage.input = event.usage.inputTokens || 0;
|
||||
output.usage.output = event.usage.outputTokens || 0;
|
||||
output.usage.cacheRead = event.usage.cacheReadInputTokens || 0;
|
||||
output.usage.cacheWrite = event.usage.cacheWriteInputTokens || 0;
|
||||
output.usage.totalTokens = event.usage.totalTokens || output.usage.input + output.usage.output;
|
||||
calculateCost(model, output.usage);
|
||||
}
|
||||
}
|
||||
|
||||
function handleContentBlockStop(
|
||||
event: ContentBlockStopEvent,
|
||||
blocks: Block[],
|
||||
output: AssistantMessage,
|
||||
stream: AssistantMessageEventStream,
|
||||
): void {
|
||||
const index = blocks.findIndex((b) => b.index === event.contentBlockIndex);
|
||||
const block = blocks[index];
|
||||
if (!block) return;
|
||||
delete (block as Block).index;
|
||||
|
||||
switch (block.type) {
|
||||
case "text":
|
||||
stream.push({ type: "text_end", contentIndex: index, content: block.text, partial: output });
|
||||
break;
|
||||
case "thinking":
|
||||
stream.push({ type: "thinking_end", contentIndex: index, content: block.thinking, partial: output });
|
||||
break;
|
||||
case "toolCall":
|
||||
block.arguments = parseStreamingJson(block.partialJson);
|
||||
// Finalize in-place and strip the scratch buffer so replay only
|
||||
// carries parsed arguments.
|
||||
delete (block as Block).partialJson;
|
||||
stream.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: output });
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the model supports adaptive thinking (Opus 4.6+, Sonnet 4.6).
|
||||
*/
|
||||
function supportsAdaptiveThinking(modelId: string): boolean {
|
||||
return (
|
||||
modelId.includes("opus-4-6") ||
|
||||
modelId.includes("opus-4.6") ||
|
||||
modelId.includes("opus-4-7") ||
|
||||
modelId.includes("opus-4.7") ||
|
||||
modelId.includes("sonnet-4-6") ||
|
||||
modelId.includes("sonnet-4.6")
|
||||
);
|
||||
}
|
||||
|
||||
function mapThinkingLevelToEffort(
|
||||
level: SimpleStreamOptions["reasoning"],
|
||||
modelId: string,
|
||||
): "low" | "medium" | "high" | "xhigh" | "max" {
|
||||
switch (level) {
|
||||
case "minimal":
|
||||
case "low":
|
||||
return "low";
|
||||
case "medium":
|
||||
return "medium";
|
||||
case "high":
|
||||
return "high";
|
||||
case "xhigh":
|
||||
if (modelId.includes("opus-4-6") || modelId.includes("opus-4.6")) {
|
||||
return "max";
|
||||
}
|
||||
if (modelId.includes("opus-4-7") || modelId.includes("opus-4.7")) {
|
||||
return "xhigh";
|
||||
}
|
||||
return "high";
|
||||
default:
|
||||
return "high";
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolve cache retention preference.
|
||||
* Defaults to "short" and uses PI_CACHE_RETENTION for backward compatibility.
|
||||
*/
|
||||
function resolveCacheRetention(cacheRetention?: CacheRetention): CacheRetention {
|
||||
if (cacheRetention) {
|
||||
return cacheRetention;
|
||||
}
|
||||
if (typeof process !== "undefined" && process.env.PI_CACHE_RETENTION === "long") {
|
||||
return "long";
|
||||
}
|
||||
return "short";
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the model supports prompt caching.
|
||||
* Supported: Claude 3.5 Haiku, Claude 3.7 Sonnet, Claude 4.x models
|
||||
*
|
||||
* For base models and system-defined inference profiles the model ID / ARN
|
||||
* contains the model name, so we can decide locally.
|
||||
*
|
||||
* For application inference profiles (whose ARNs don't contain the model name),
|
||||
* set AWS_BEDROCK_FORCE_CACHE=1 to enable cache points. Amazon Nova models
|
||||
* have automatic caching and don't need explicit cache points.
|
||||
*/
|
||||
function supportsPromptCaching(model: Model<"bedrock-converse-stream">): boolean {
|
||||
const id = model.id.toLowerCase();
|
||||
if (!id.includes("claude")) {
|
||||
// Application inference profiles don't contain the model name in the ARN.
|
||||
// Allow users to force cache points via environment variable.
|
||||
if (typeof process !== "undefined" && process.env.AWS_BEDROCK_FORCE_CACHE === "1") return true;
|
||||
return false;
|
||||
}
|
||||
// Claude 4.x models (opus-4, sonnet-4, haiku-4)
|
||||
if (id.includes("-4-") || id.includes("-4.")) return true;
|
||||
// Claude 3.7 Sonnet
|
||||
if (id.includes("claude-3-7-sonnet")) return true;
|
||||
// Claude 3.5 Haiku
|
||||
if (id.includes("claude-3-5-haiku")) return true;
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if the model supports thinking signatures in reasoningContent.
|
||||
* Only Anthropic Claude models support the signature field.
|
||||
* Other models (OpenAI, Qwen, Minimax, Moonshot, etc.) reject it with:
|
||||
* "This model doesn't support the reasoningContent.reasoningText.signature field"
|
||||
*/
|
||||
function supportsThinkingSignature(model: Model<"bedrock-converse-stream">): boolean {
|
||||
const id = model.id.toLowerCase();
|
||||
return id.includes("anthropic.claude") || id.includes("anthropic/claude");
|
||||
}
|
||||
|
||||
function buildSystemPrompt(
|
||||
systemPrompt: string | undefined,
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
cacheRetention: CacheRetention,
|
||||
): SystemContentBlock[] | undefined {
|
||||
if (!systemPrompt) return undefined;
|
||||
|
||||
const blocks: SystemContentBlock[] = [{ text: sanitizeSurrogates(systemPrompt) }];
|
||||
|
||||
// Add cache point for supported Claude models when caching is enabled
|
||||
if (cacheRetention !== "none" && supportsPromptCaching(model)) {
|
||||
blocks.push({
|
||||
cachePoint: { type: CachePointType.DEFAULT, ...(cacheRetention === "long" ? { ttl: CacheTTL.ONE_HOUR } : {}) },
|
||||
});
|
||||
}
|
||||
|
||||
return blocks;
|
||||
}
|
||||
|
||||
function normalizeToolCallId(id: string): string {
|
||||
const sanitized = id.replace(/[^a-zA-Z0-9_-]/g, "_");
|
||||
return sanitized.length > 64 ? sanitized.slice(0, 64) : sanitized;
|
||||
}
|
||||
|
||||
function convertMessages(
|
||||
context: Context,
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
cacheRetention: CacheRetention,
|
||||
): Message[] {
|
||||
const result: Message[] = [];
|
||||
const transformedMessages = transformMessages(context.messages, model, normalizeToolCallId);
|
||||
|
||||
for (let i = 0; i < transformedMessages.length; i++) {
|
||||
const m = transformedMessages[i];
|
||||
|
||||
switch (m.role) {
|
||||
case "user":
|
||||
result.push({
|
||||
role: ConversationRole.USER,
|
||||
content:
|
||||
typeof m.content === "string"
|
||||
? [{ text: sanitizeSurrogates(m.content) }]
|
||||
: m.content.map((c) => {
|
||||
switch (c.type) {
|
||||
case "text":
|
||||
return { text: sanitizeSurrogates(c.text) };
|
||||
case "image":
|
||||
return { image: createImageBlock(c.mimeType, c.data) };
|
||||
default:
|
||||
throw new Error("Unknown user content type");
|
||||
}
|
||||
}),
|
||||
});
|
||||
break;
|
||||
case "assistant": {
|
||||
// Skip assistant messages with empty content (e.g., from aborted requests)
|
||||
// Bedrock rejects messages with empty content arrays
|
||||
if (m.content.length === 0) {
|
||||
continue;
|
||||
}
|
||||
const contentBlocks: ContentBlock[] = [];
|
||||
for (const c of m.content) {
|
||||
switch (c.type) {
|
||||
case "text":
|
||||
// Skip empty text blocks
|
||||
if (c.text.trim().length === 0) continue;
|
||||
contentBlocks.push({ text: sanitizeSurrogates(c.text) });
|
||||
break;
|
||||
case "toolCall":
|
||||
contentBlocks.push({
|
||||
toolUse: { toolUseId: c.id, name: c.name, input: c.arguments },
|
||||
});
|
||||
break;
|
||||
case "thinking":
|
||||
// Skip empty thinking blocks
|
||||
if (c.thinking.trim().length === 0) continue;
|
||||
// Only Anthropic models support the signature field in reasoningText.
|
||||
// For other models, we omit the signature to avoid errors like:
|
||||
// "This model doesn't support the reasoningContent.reasoningText.signature field"
|
||||
if (supportsThinkingSignature(model)) {
|
||||
// Signatures arrive after thinking deltas. If a partial or externally
|
||||
// persisted message lacks a signature, Bedrock rejects the replayed
|
||||
// reasoning block. Fall back to plain text, matching Anthropic.
|
||||
if (!c.thinkingSignature || c.thinkingSignature.trim().length === 0) {
|
||||
contentBlocks.push({ text: sanitizeSurrogates(c.thinking) });
|
||||
} else {
|
||||
contentBlocks.push({
|
||||
reasoningContent: {
|
||||
reasoningText: {
|
||||
text: sanitizeSurrogates(c.thinking),
|
||||
signature: c.thinkingSignature,
|
||||
},
|
||||
},
|
||||
});
|
||||
}
|
||||
} else {
|
||||
contentBlocks.push({
|
||||
reasoningContent: {
|
||||
reasoningText: { text: sanitizeSurrogates(c.thinking) },
|
||||
},
|
||||
});
|
||||
}
|
||||
break;
|
||||
default:
|
||||
throw new Error("Unknown assistant content type");
|
||||
}
|
||||
}
|
||||
// Skip if all content blocks were filtered out
|
||||
if (contentBlocks.length === 0) {
|
||||
continue;
|
||||
}
|
||||
result.push({
|
||||
role: ConversationRole.ASSISTANT,
|
||||
content: contentBlocks,
|
||||
});
|
||||
break;
|
||||
}
|
||||
case "toolResult": {
|
||||
// Collect all consecutive toolResult messages into a single user message
|
||||
// Bedrock requires all tool results to be in one message
|
||||
const toolResults: ContentBlock.ToolResultMember[] = [];
|
||||
|
||||
// Add current tool result with all content blocks combined
|
||||
toolResults.push({
|
||||
toolResult: {
|
||||
toolUseId: m.toolCallId,
|
||||
content: m.content.map((c) =>
|
||||
c.type === "image"
|
||||
? { image: createImageBlock(c.mimeType, c.data) }
|
||||
: { text: sanitizeSurrogates(c.text) },
|
||||
),
|
||||
status: m.isError ? ToolResultStatus.ERROR : ToolResultStatus.SUCCESS,
|
||||
},
|
||||
});
|
||||
|
||||
// Look ahead for consecutive toolResult messages
|
||||
let j = i + 1;
|
||||
while (j < transformedMessages.length && transformedMessages[j].role === "toolResult") {
|
||||
const nextMsg = transformedMessages[j] as ToolResultMessage;
|
||||
toolResults.push({
|
||||
toolResult: {
|
||||
toolUseId: nextMsg.toolCallId,
|
||||
content: nextMsg.content.map((c) =>
|
||||
c.type === "image"
|
||||
? { image: createImageBlock(c.mimeType, c.data) }
|
||||
: { text: sanitizeSurrogates(c.text) },
|
||||
),
|
||||
status: nextMsg.isError ? ToolResultStatus.ERROR : ToolResultStatus.SUCCESS,
|
||||
},
|
||||
});
|
||||
j++;
|
||||
}
|
||||
|
||||
// Skip the messages we've already processed
|
||||
i = j - 1;
|
||||
|
||||
result.push({
|
||||
role: ConversationRole.USER,
|
||||
content: toolResults,
|
||||
});
|
||||
break;
|
||||
}
|
||||
default:
|
||||
throw new Error("Unknown message role");
|
||||
}
|
||||
}
|
||||
|
||||
// Add cache point to the last user message for supported Claude models when caching is enabled
|
||||
if (cacheRetention !== "none" && supportsPromptCaching(model) && result.length > 0) {
|
||||
const lastMessage = result[result.length - 1];
|
||||
if (lastMessage.role === ConversationRole.USER && lastMessage.content) {
|
||||
(lastMessage.content as ContentBlock[]).push({
|
||||
cachePoint: {
|
||||
type: CachePointType.DEFAULT,
|
||||
...(cacheRetention === "long" ? { ttl: CacheTTL.ONE_HOUR } : {}),
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
function convertToolConfig(
|
||||
tools: Tool[] | undefined,
|
||||
toolChoice: BedrockOptions["toolChoice"],
|
||||
): ToolConfiguration | undefined {
|
||||
if (!tools?.length || toolChoice === "none") return undefined;
|
||||
|
||||
const bedrockTools: BedrockTool[] = tools.map((tool) => ({
|
||||
toolSpec: {
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
inputSchema: { json: tool.parameters },
|
||||
},
|
||||
}));
|
||||
|
||||
let bedrockToolChoice: ToolChoice | undefined;
|
||||
switch (toolChoice) {
|
||||
case "auto":
|
||||
bedrockToolChoice = { auto: {} };
|
||||
break;
|
||||
case "any":
|
||||
bedrockToolChoice = { any: {} };
|
||||
break;
|
||||
default:
|
||||
if (toolChoice?.type === "tool") {
|
||||
bedrockToolChoice = { tool: { name: toolChoice.name } };
|
||||
}
|
||||
}
|
||||
|
||||
return { tools: bedrockTools, toolChoice: bedrockToolChoice };
|
||||
}
|
||||
|
||||
function mapStopReason(reason: string | undefined): StopReason {
|
||||
switch (reason) {
|
||||
case BedrockStopReason.END_TURN:
|
||||
case BedrockStopReason.STOP_SEQUENCE:
|
||||
return "stop";
|
||||
case BedrockStopReason.MAX_TOKENS:
|
||||
case BedrockStopReason.MODEL_CONTEXT_WINDOW_EXCEEDED:
|
||||
return "length";
|
||||
case BedrockStopReason.TOOL_USE:
|
||||
return "toolUse";
|
||||
default:
|
||||
return "error";
|
||||
}
|
||||
}
|
||||
|
||||
function isGovCloudBedrockTarget(model: Model<"bedrock-converse-stream">, options: BedrockOptions): boolean {
|
||||
const region = options.region || process.env.AWS_REGION || process.env.AWS_DEFAULT_REGION;
|
||||
if (region?.toLowerCase().startsWith("us-gov-")) {
|
||||
return true;
|
||||
}
|
||||
|
||||
const modelId = model.id.toLowerCase();
|
||||
return modelId.startsWith("us-gov.") || modelId.startsWith("arn:aws-us-gov:");
|
||||
}
|
||||
|
||||
function buildAdditionalModelRequestFields(
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
options: BedrockOptions,
|
||||
): Record<string, any> | undefined {
|
||||
if (!options.reasoning || !model.reasoning) {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
if (model.id.includes("anthropic.claude") || model.id.includes("anthropic/claude")) {
|
||||
// GovCloud Bedrock currently rejects the Claude thinking.display field.
|
||||
// Omit it there until the GovCloud Converse schema catches up.
|
||||
const display = isGovCloudBedrockTarget(model, options) ? undefined : (options.thinkingDisplay ?? "summarized");
|
||||
const result: Record<string, any> = supportsAdaptiveThinking(model.id)
|
||||
? {
|
||||
thinking: { type: "adaptive", ...(display !== undefined ? { display } : {}) },
|
||||
output_config: { effort: mapThinkingLevelToEffort(options.reasoning, model.id) },
|
||||
}
|
||||
: (() => {
|
||||
const defaultBudgets: Record<ThinkingLevel, number> = {
|
||||
minimal: 1024,
|
||||
low: 2048,
|
||||
medium: 8192,
|
||||
high: 16384,
|
||||
xhigh: 16384, // Claude doesn't support xhigh, clamp to high
|
||||
};
|
||||
|
||||
// Custom budgets override defaults (xhigh not in ThinkingBudgets, use high)
|
||||
const level = options.reasoning === "xhigh" ? "high" : options.reasoning;
|
||||
const budget = options.thinkingBudgets?.[level] ?? defaultBudgets[options.reasoning];
|
||||
|
||||
return {
|
||||
thinking: {
|
||||
type: "enabled",
|
||||
budget_tokens: budget,
|
||||
...(display !== undefined ? { display } : {}),
|
||||
},
|
||||
};
|
||||
})();
|
||||
|
||||
if (!supportsAdaptiveThinking(model.id) && (options.interleavedThinking ?? true)) {
|
||||
result.anthropic_beta = ["interleaved-thinking-2025-05-14"];
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function createImageBlock(mimeType: string, data: string) {
|
||||
let format: ImageFormat;
|
||||
switch (mimeType) {
|
||||
case "image/jpeg":
|
||||
case "image/jpg":
|
||||
format = ImageFormat.JPEG;
|
||||
break;
|
||||
case "image/png":
|
||||
format = ImageFormat.PNG;
|
||||
break;
|
||||
case "image/gif":
|
||||
format = ImageFormat.GIF;
|
||||
break;
|
||||
case "image/webp":
|
||||
format = ImageFormat.WEBP;
|
||||
break;
|
||||
default:
|
||||
throw new Error(`Unknown image type: ${mimeType}`);
|
||||
}
|
||||
|
||||
const binaryString = atob(data);
|
||||
const bytes = new Uint8Array(binaryString.length);
|
||||
for (let i = 0; i < binaryString.length; i++) {
|
||||
bytes[i] = binaryString.charCodeAt(i);
|
||||
}
|
||||
|
||||
return { source: { bytes }, format };
|
||||
}
|
||||
@@ -1,954 +0,0 @@
|
||||
import Anthropic from "@anthropic-ai/sdk";
|
||||
import type {
|
||||
CacheControlEphemeral,
|
||||
ContentBlockParam,
|
||||
MessageCreateParamsStreaming,
|
||||
MessageParam,
|
||||
} from "@anthropic-ai/sdk/resources/messages.js";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { calculateCost } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
CacheRetention,
|
||||
Context,
|
||||
ImageContent,
|
||||
Message,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StopReason,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
TextContent,
|
||||
ThinkingContent,
|
||||
Tool,
|
||||
ToolCall,
|
||||
ToolResultMessage,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { parseStreamingJson } from "../utils/json-parse.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
|
||||
import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js";
|
||||
import { adjustMaxTokensForThinking, buildBaseOptions } from "./simple-options.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
|
||||
/**
|
||||
* Resolve cache retention preference.
|
||||
* Defaults to "short" and uses PI_CACHE_RETENTION for backward compatibility.
|
||||
*/
|
||||
function resolveCacheRetention(cacheRetention?: CacheRetention): CacheRetention {
|
||||
if (cacheRetention) {
|
||||
return cacheRetention;
|
||||
}
|
||||
if (typeof process !== "undefined" && process.env.PI_CACHE_RETENTION === "long") {
|
||||
return "long";
|
||||
}
|
||||
return "short";
|
||||
}
|
||||
|
||||
function getCacheControl(
|
||||
baseUrl: string,
|
||||
cacheRetention?: CacheRetention,
|
||||
): { retention: CacheRetention; cacheControl?: CacheControlEphemeral } {
|
||||
const retention = resolveCacheRetention(cacheRetention);
|
||||
if (retention === "none") {
|
||||
return { retention };
|
||||
}
|
||||
const ttl = retention === "long" && baseUrl.includes("api.anthropic.com") ? "1h" : undefined;
|
||||
return {
|
||||
retention,
|
||||
cacheControl: { type: "ephemeral", ...(ttl && { ttl }) },
|
||||
};
|
||||
}
|
||||
|
||||
// Stealth mode: Mimic Claude Code's tool naming exactly
|
||||
const claudeCodeVersion = "2.1.75";
|
||||
|
||||
// Claude Code 2.x tool names (canonical casing)
|
||||
// Source: https://cchistory.mariozechner.at/data/prompts-2.1.11.md
|
||||
// To update: https://github.com/badlogic/cchistory
|
||||
const claudeCodeTools = [
|
||||
"Read",
|
||||
"Write",
|
||||
"Edit",
|
||||
"Bash",
|
||||
"Grep",
|
||||
"Glob",
|
||||
"AskUserQuestion",
|
||||
"EnterPlanMode",
|
||||
"ExitPlanMode",
|
||||
"KillShell",
|
||||
"NotebookEdit",
|
||||
"Skill",
|
||||
"Task",
|
||||
"TaskOutput",
|
||||
"TodoWrite",
|
||||
"WebFetch",
|
||||
"WebSearch",
|
||||
];
|
||||
|
||||
const ccToolLookup = new Map(claudeCodeTools.map((t) => [t.toLowerCase(), t]));
|
||||
|
||||
// Convert tool name to CC canonical casing if it matches (case-insensitive)
|
||||
const toClaudeCodeName = (name: string) => ccToolLookup.get(name.toLowerCase()) ?? name;
|
||||
const fromClaudeCodeName = (name: string, tools?: Tool[]) => {
|
||||
if (tools && tools.length > 0) {
|
||||
const lowerName = name.toLowerCase();
|
||||
const matchedTool = tools.find((tool) => tool.name.toLowerCase() === lowerName);
|
||||
if (matchedTool) return matchedTool.name;
|
||||
}
|
||||
return name;
|
||||
};
|
||||
|
||||
/**
|
||||
* Convert content blocks to Anthropic API format
|
||||
*/
|
||||
function convertContentBlocks(content: (TextContent | ImageContent)[]):
|
||||
| string
|
||||
| Array<
|
||||
| { type: "text"; text: string }
|
||||
| {
|
||||
type: "image";
|
||||
source: {
|
||||
type: "base64";
|
||||
media_type: "image/jpeg" | "image/png" | "image/gif" | "image/webp";
|
||||
data: string;
|
||||
};
|
||||
}
|
||||
> {
|
||||
// If only text blocks, return as concatenated string for simplicity
|
||||
const hasImages = content.some((c) => c.type === "image");
|
||||
if (!hasImages) {
|
||||
return sanitizeSurrogates(content.map((c) => (c as TextContent).text).join("\n"));
|
||||
}
|
||||
|
||||
// If we have images, convert to content block array
|
||||
const blocks = content.map((block) => {
|
||||
if (block.type === "text") {
|
||||
return {
|
||||
type: "text" as const,
|
||||
text: sanitizeSurrogates(block.text),
|
||||
};
|
||||
}
|
||||
return {
|
||||
type: "image" as const,
|
||||
source: {
|
||||
type: "base64" as const,
|
||||
media_type: block.mimeType as "image/jpeg" | "image/png" | "image/gif" | "image/webp",
|
||||
data: block.data,
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
// If only images (no text), add placeholder text block
|
||||
const hasText = blocks.some((b) => b.type === "text");
|
||||
if (!hasText) {
|
||||
blocks.unshift({
|
||||
type: "text" as const,
|
||||
text: "(see attached image)",
|
||||
});
|
||||
}
|
||||
|
||||
return blocks;
|
||||
}
|
||||
|
||||
export type AnthropicEffort = "low" | "medium" | "high" | "xhigh" | "max";
|
||||
|
||||
export type AnthropicThinkingDisplay = "summarized" | "omitted";
|
||||
|
||||
export interface AnthropicOptions extends StreamOptions {
|
||||
/**
|
||||
* Enable extended thinking.
|
||||
* For Opus 4.6 and Sonnet 4.6: uses adaptive thinking (model decides when/how much to think).
|
||||
* For older models: uses budget-based thinking with thinkingBudgetTokens.
|
||||
*/
|
||||
thinkingEnabled?: boolean;
|
||||
/**
|
||||
* Token budget for extended thinking (older models only).
|
||||
* Ignored for Opus 4.6 and Sonnet 4.6, which use adaptive thinking.
|
||||
*/
|
||||
thinkingBudgetTokens?: number;
|
||||
/**
|
||||
* Effort level for adaptive thinking (Opus 4.6+ and Sonnet 4.6).
|
||||
* Controls how much thinking Claude allocates:
|
||||
* - "max": Always thinks with no constraints (Opus 4.6 only)
|
||||
* - "xhigh": Highest reasoning level (Opus 4.7)
|
||||
* - "high": Always thinks, deep reasoning (default)
|
||||
* - "medium": Moderate thinking, may skip for simple queries
|
||||
* - "low": Minimal thinking, skips for simple tasks
|
||||
* Ignored for older models.
|
||||
*/
|
||||
effort?: AnthropicEffort;
|
||||
/**
|
||||
* Controls how thinking content is returned in API responses.
|
||||
* - "summarized": Thinking blocks contain summarized thinking text (default here).
|
||||
* - "omitted": Thinking blocks return an empty thinking field; the encrypted
|
||||
* signature still travels back for multi-turn continuity. Use for faster
|
||||
* time-to-first-text-token when your UI does not surface thinking.
|
||||
*
|
||||
* Note: Anthropic's API default for Claude Opus 4.7 and Claude Mythos Preview
|
||||
* is "omitted". We default to "summarized" here to keep behavior consistent
|
||||
* with older Claude 4 models. Set this explicitly to "omitted" to opt in.
|
||||
*/
|
||||
thinkingDisplay?: AnthropicThinkingDisplay;
|
||||
interleavedThinking?: boolean;
|
||||
toolChoice?: "auto" | "any" | "none" | { type: "tool"; name: string };
|
||||
/**
|
||||
* Pre-built Anthropic client instance. When provided, skips internal client
|
||||
* construction entirely. Use this to inject alternative SDK clients such as
|
||||
* `AnthropicVertex` that shares the same messaging API.
|
||||
*/
|
||||
client?: Anthropic;
|
||||
}
|
||||
|
||||
function mergeHeaders(...headerSources: (Record<string, string> | undefined)[]): Record<string, string> {
|
||||
const merged: Record<string, string> = {};
|
||||
for (const headers of headerSources) {
|
||||
if (headers) {
|
||||
Object.assign(merged, headers);
|
||||
}
|
||||
}
|
||||
return merged;
|
||||
}
|
||||
|
||||
export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOptions> = (
|
||||
model: Model<"anthropic-messages">,
|
||||
context: Context,
|
||||
options?: AnthropicOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: model.api as Api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
try {
|
||||
let client: Anthropic;
|
||||
let isOAuth: boolean;
|
||||
|
||||
if (options?.client) {
|
||||
client = options.client;
|
||||
isOAuth = false;
|
||||
} else {
|
||||
const apiKey = options?.apiKey ?? getEnvApiKey(model.provider) ?? "";
|
||||
|
||||
let copilotDynamicHeaders: Record<string, string> | undefined;
|
||||
if (model.provider === "github-copilot") {
|
||||
const hasImages = hasCopilotVisionInput(context.messages);
|
||||
copilotDynamicHeaders = buildCopilotDynamicHeaders({
|
||||
messages: context.messages,
|
||||
hasImages,
|
||||
});
|
||||
}
|
||||
|
||||
const created = createClient(
|
||||
model,
|
||||
apiKey,
|
||||
options?.interleavedThinking ?? true,
|
||||
options?.headers,
|
||||
copilotDynamicHeaders,
|
||||
);
|
||||
client = created.client;
|
||||
isOAuth = created.isOAuthToken;
|
||||
}
|
||||
let params = buildParams(model, context, isOAuth, options);
|
||||
const nextParams = await options?.onPayload?.(params, model);
|
||||
if (nextParams !== undefined) {
|
||||
params = nextParams as MessageCreateParamsStreaming;
|
||||
}
|
||||
const { data: anthropicStream, response } = await client.messages
|
||||
.stream({ ...params, stream: true }, { signal: options?.signal })
|
||||
.withResponse();
|
||||
await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model);
|
||||
stream.push({ type: "start", partial: output });
|
||||
|
||||
type Block = (ThinkingContent | TextContent | (ToolCall & { partialJson: string })) & { index: number };
|
||||
const blocks = output.content as Block[];
|
||||
|
||||
for await (const event of anthropicStream) {
|
||||
if (event.type === "message_start") {
|
||||
output.responseId = event.message.id;
|
||||
// Capture initial token usage from message_start event
|
||||
// This ensures we have input token counts even if the stream is aborted early
|
||||
output.usage.input = event.message.usage.input_tokens || 0;
|
||||
output.usage.output = event.message.usage.output_tokens || 0;
|
||||
output.usage.cacheRead = event.message.usage.cache_read_input_tokens || 0;
|
||||
output.usage.cacheWrite = event.message.usage.cache_creation_input_tokens || 0;
|
||||
// Anthropic doesn't provide total_tokens, compute from components
|
||||
output.usage.totalTokens =
|
||||
output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite;
|
||||
calculateCost(model, output.usage);
|
||||
} else if (event.type === "content_block_start") {
|
||||
if (event.content_block.type === "text") {
|
||||
const block: Block = {
|
||||
type: "text",
|
||||
text: "",
|
||||
index: event.index,
|
||||
};
|
||||
output.content.push(block);
|
||||
stream.push({ type: "text_start", contentIndex: output.content.length - 1, partial: output });
|
||||
} else if (event.content_block.type === "thinking") {
|
||||
const block: Block = {
|
||||
type: "thinking",
|
||||
thinking: "",
|
||||
thinkingSignature: "",
|
||||
index: event.index,
|
||||
};
|
||||
output.content.push(block);
|
||||
stream.push({ type: "thinking_start", contentIndex: output.content.length - 1, partial: output });
|
||||
} else if (event.content_block.type === "redacted_thinking") {
|
||||
const block: Block = {
|
||||
type: "thinking",
|
||||
thinking: "[Reasoning redacted]",
|
||||
thinkingSignature: event.content_block.data,
|
||||
redacted: true,
|
||||
index: event.index,
|
||||
};
|
||||
output.content.push(block);
|
||||
stream.push({ type: "thinking_start", contentIndex: output.content.length - 1, partial: output });
|
||||
} else if (event.content_block.type === "tool_use") {
|
||||
const block: Block = {
|
||||
type: "toolCall",
|
||||
id: event.content_block.id,
|
||||
name: isOAuth
|
||||
? fromClaudeCodeName(event.content_block.name, context.tools)
|
||||
: event.content_block.name,
|
||||
arguments: (event.content_block.input as Record<string, any>) ?? {},
|
||||
partialJson: "",
|
||||
index: event.index,
|
||||
};
|
||||
output.content.push(block);
|
||||
stream.push({ type: "toolcall_start", contentIndex: output.content.length - 1, partial: output });
|
||||
}
|
||||
} else if (event.type === "content_block_delta") {
|
||||
if (event.delta.type === "text_delta") {
|
||||
const index = blocks.findIndex((b) => b.index === event.index);
|
||||
const block = blocks[index];
|
||||
if (block && block.type === "text") {
|
||||
block.text += event.delta.text;
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: index,
|
||||
delta: event.delta.text,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
} else if (event.delta.type === "thinking_delta") {
|
||||
const index = blocks.findIndex((b) => b.index === event.index);
|
||||
const block = blocks[index];
|
||||
if (block && block.type === "thinking") {
|
||||
block.thinking += event.delta.thinking;
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: index,
|
||||
delta: event.delta.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
} else if (event.delta.type === "input_json_delta") {
|
||||
const index = blocks.findIndex((b) => b.index === event.index);
|
||||
const block = blocks[index];
|
||||
if (block && block.type === "toolCall") {
|
||||
block.partialJson += event.delta.partial_json;
|
||||
block.arguments = parseStreamingJson(block.partialJson);
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: index,
|
||||
delta: event.delta.partial_json,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
} else if (event.delta.type === "signature_delta") {
|
||||
const index = blocks.findIndex((b) => b.index === event.index);
|
||||
const block = blocks[index];
|
||||
if (block && block.type === "thinking") {
|
||||
block.thinkingSignature = block.thinkingSignature || "";
|
||||
block.thinkingSignature += event.delta.signature;
|
||||
}
|
||||
}
|
||||
} else if (event.type === "content_block_stop") {
|
||||
const index = blocks.findIndex((b) => b.index === event.index);
|
||||
const block = blocks[index];
|
||||
if (block) {
|
||||
delete (block as any).index;
|
||||
if (block.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: index,
|
||||
content: block.text,
|
||||
partial: output,
|
||||
});
|
||||
} else if (block.type === "thinking") {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: index,
|
||||
content: block.thinking,
|
||||
partial: output,
|
||||
});
|
||||
} else if (block.type === "toolCall") {
|
||||
block.arguments = parseStreamingJson(block.partialJson);
|
||||
// Finalize in-place and strip the scratch buffer so replay only
|
||||
// carries parsed arguments.
|
||||
delete (block as { partialJson?: string }).partialJson;
|
||||
stream.push({
|
||||
type: "toolcall_end",
|
||||
contentIndex: index,
|
||||
toolCall: block,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
} else if (event.type === "message_delta") {
|
||||
if (event.delta.stop_reason) {
|
||||
output.stopReason = mapStopReason(event.delta.stop_reason);
|
||||
}
|
||||
// Only update usage fields if present (not null).
|
||||
// Preserves input_tokens from message_start when proxies omit it in message_delta.
|
||||
if (event.usage.input_tokens != null) {
|
||||
output.usage.input = event.usage.input_tokens;
|
||||
}
|
||||
if (event.usage.output_tokens != null) {
|
||||
output.usage.output = event.usage.output_tokens;
|
||||
}
|
||||
if (event.usage.cache_read_input_tokens != null) {
|
||||
output.usage.cacheRead = event.usage.cache_read_input_tokens;
|
||||
}
|
||||
if (event.usage.cache_creation_input_tokens != null) {
|
||||
output.usage.cacheWrite = event.usage.cache_creation_input_tokens;
|
||||
}
|
||||
// Anthropic doesn't provide total_tokens, compute from components
|
||||
output.usage.totalTokens =
|
||||
output.usage.input + output.usage.output + output.usage.cacheRead + output.usage.cacheWrite;
|
||||
calculateCost(model, output.usage);
|
||||
}
|
||||
}
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "aborted" || output.stopReason === "error") {
|
||||
throw new Error("An unknown error occurred");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
for (const block of output.content) {
|
||||
delete (block as { index?: number }).index;
|
||||
// partialJson is only a streaming scratch buffer; never persist it.
|
||||
delete (block as { partialJson?: string }).partialJson;
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
/**
|
||||
* Check if a model supports adaptive thinking (Opus 4.6+, Sonnet 4.6)
|
||||
*/
|
||||
function supportsAdaptiveThinking(modelId: string): boolean {
|
||||
// Adaptive-thinking model IDs (with or without date suffix)
|
||||
return (
|
||||
modelId.includes("opus-4-6") ||
|
||||
modelId.includes("opus-4.6") ||
|
||||
modelId.includes("opus-4-7") ||
|
||||
modelId.includes("opus-4.7") ||
|
||||
modelId.includes("sonnet-4-6") ||
|
||||
modelId.includes("sonnet-4.6")
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Map ThinkingLevel to Anthropic effort levels for adaptive thinking.
|
||||
* Note: effort "max" is only valid on Opus 4.6, while Opus 4.7 supports "xhigh".
|
||||
*/
|
||||
function mapThinkingLevelToEffort(level: SimpleStreamOptions["reasoning"], modelId: string): AnthropicEffort {
|
||||
switch (level) {
|
||||
case "minimal":
|
||||
return "low";
|
||||
case "low":
|
||||
return "low";
|
||||
case "medium":
|
||||
return "medium";
|
||||
case "high":
|
||||
return "high";
|
||||
case "xhigh":
|
||||
if (modelId.includes("opus-4-6") || modelId.includes("opus-4.6")) {
|
||||
return "max";
|
||||
}
|
||||
if (modelId.includes("opus-4-7") || modelId.includes("opus-4.7")) {
|
||||
return "xhigh";
|
||||
}
|
||||
return "high";
|
||||
default:
|
||||
return "high";
|
||||
}
|
||||
}
|
||||
|
||||
export const streamSimpleAnthropic: StreamFunction<"anthropic-messages", SimpleStreamOptions> = (
|
||||
model: Model<"anthropic-messages">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
|
||||
if (!apiKey) {
|
||||
throw new Error(`No API key for provider: ${model.provider}`);
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
if (!options?.reasoning) {
|
||||
return streamAnthropic(model, context, { ...base, thinkingEnabled: false } satisfies AnthropicOptions);
|
||||
}
|
||||
|
||||
// For Opus 4.6 and Sonnet 4.6: use adaptive thinking with effort level
|
||||
// For older models: use budget-based thinking
|
||||
if (supportsAdaptiveThinking(model.id)) {
|
||||
const effort = mapThinkingLevelToEffort(options.reasoning, model.id);
|
||||
return streamAnthropic(model, context, {
|
||||
...base,
|
||||
thinkingEnabled: true,
|
||||
effort,
|
||||
} satisfies AnthropicOptions);
|
||||
}
|
||||
|
||||
const adjusted = adjustMaxTokensForThinking(
|
||||
base.maxTokens || 0,
|
||||
model.maxTokens,
|
||||
options.reasoning,
|
||||
options.thinkingBudgets,
|
||||
);
|
||||
|
||||
return streamAnthropic(model, context, {
|
||||
...base,
|
||||
maxTokens: adjusted.maxTokens,
|
||||
thinkingEnabled: true,
|
||||
thinkingBudgetTokens: adjusted.thinkingBudget,
|
||||
} satisfies AnthropicOptions);
|
||||
};
|
||||
|
||||
function isOAuthToken(apiKey: string): boolean {
|
||||
return apiKey.includes("sk-ant-oat");
|
||||
}
|
||||
|
||||
function createClient(
|
||||
model: Model<"anthropic-messages">,
|
||||
apiKey: string,
|
||||
interleavedThinking: boolean,
|
||||
optionsHeaders?: Record<string, string>,
|
||||
dynamicHeaders?: Record<string, string>,
|
||||
): { client: Anthropic; isOAuthToken: boolean } {
|
||||
// Adaptive thinking models (Opus 4.6, Sonnet 4.6) have interleaved thinking built-in.
|
||||
// The beta header is deprecated on Opus 4.6 and redundant on Sonnet 4.6, so skip it.
|
||||
const needsInterleavedBeta = interleavedThinking && !supportsAdaptiveThinking(model.id);
|
||||
|
||||
// Copilot: Bearer auth, selective betas (no fine-grained-tool-streaming)
|
||||
if (model.provider === "github-copilot") {
|
||||
const betaFeatures: string[] = [];
|
||||
if (needsInterleavedBeta) {
|
||||
betaFeatures.push("interleaved-thinking-2025-05-14");
|
||||
}
|
||||
|
||||
const client = new Anthropic({
|
||||
apiKey: null,
|
||||
authToken: apiKey,
|
||||
baseURL: model.baseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: mergeHeaders(
|
||||
{
|
||||
accept: "application/json",
|
||||
"anthropic-dangerous-direct-browser-access": "true",
|
||||
...(betaFeatures.length > 0 ? { "anthropic-beta": betaFeatures.join(",") } : {}),
|
||||
},
|
||||
model.headers,
|
||||
dynamicHeaders,
|
||||
optionsHeaders,
|
||||
),
|
||||
});
|
||||
|
||||
return { client, isOAuthToken: false };
|
||||
}
|
||||
|
||||
const betaFeatures = ["fine-grained-tool-streaming-2025-05-14"];
|
||||
if (needsInterleavedBeta) {
|
||||
betaFeatures.push("interleaved-thinking-2025-05-14");
|
||||
}
|
||||
|
||||
// OAuth: Bearer auth, Claude Code identity headers
|
||||
if (isOAuthToken(apiKey)) {
|
||||
const client = new Anthropic({
|
||||
apiKey: null,
|
||||
authToken: apiKey,
|
||||
baseURL: model.baseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: mergeHeaders(
|
||||
{
|
||||
accept: "application/json",
|
||||
"anthropic-dangerous-direct-browser-access": "true",
|
||||
"anthropic-beta": `claude-code-20250219,oauth-2025-04-20,${betaFeatures.join(",")}`,
|
||||
"user-agent": `claude-cli/${claudeCodeVersion}`,
|
||||
"x-app": "cli",
|
||||
},
|
||||
model.headers,
|
||||
optionsHeaders,
|
||||
),
|
||||
});
|
||||
|
||||
return { client, isOAuthToken: true };
|
||||
}
|
||||
|
||||
// API key auth
|
||||
const client = new Anthropic({
|
||||
apiKey,
|
||||
baseURL: model.baseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: mergeHeaders(
|
||||
{
|
||||
accept: "application/json",
|
||||
"anthropic-dangerous-direct-browser-access": "true",
|
||||
"anthropic-beta": betaFeatures.join(","),
|
||||
},
|
||||
model.headers,
|
||||
optionsHeaders,
|
||||
),
|
||||
});
|
||||
|
||||
return { client, isOAuthToken: false };
|
||||
}
|
||||
|
||||
function buildParams(
|
||||
model: Model<"anthropic-messages">,
|
||||
context: Context,
|
||||
isOAuthToken: boolean,
|
||||
options?: AnthropicOptions,
|
||||
): MessageCreateParamsStreaming {
|
||||
const { cacheControl } = getCacheControl(model.baseUrl, options?.cacheRetention);
|
||||
const params: MessageCreateParamsStreaming = {
|
||||
model: model.id,
|
||||
messages: convertMessages(context.messages, model, isOAuthToken, cacheControl),
|
||||
max_tokens: options?.maxTokens || (model.maxTokens / 3) | 0,
|
||||
stream: true,
|
||||
};
|
||||
|
||||
// For OAuth tokens, we MUST include Claude Code identity
|
||||
if (isOAuthToken) {
|
||||
params.system = [
|
||||
{
|
||||
type: "text",
|
||||
text: "You are Claude Code, Anthropic's official CLI for Claude.",
|
||||
...(cacheControl ? { cache_control: cacheControl } : {}),
|
||||
},
|
||||
];
|
||||
if (context.systemPrompt) {
|
||||
params.system.push({
|
||||
type: "text",
|
||||
text: sanitizeSurrogates(context.systemPrompt),
|
||||
...(cacheControl ? { cache_control: cacheControl } : {}),
|
||||
});
|
||||
}
|
||||
} else if (context.systemPrompt) {
|
||||
// Add cache control to system prompt for non-OAuth tokens
|
||||
params.system = [
|
||||
{
|
||||
type: "text",
|
||||
text: sanitizeSurrogates(context.systemPrompt),
|
||||
...(cacheControl ? { cache_control: cacheControl } : {}),
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
// Temperature is incompatible with extended thinking (adaptive or budget-based).
|
||||
if (options?.temperature !== undefined && !options?.thinkingEnabled) {
|
||||
params.temperature = options.temperature;
|
||||
}
|
||||
|
||||
if (context.tools) {
|
||||
params.tools = convertTools(context.tools, isOAuthToken, cacheControl);
|
||||
}
|
||||
|
||||
// Configure thinking mode: adaptive (Opus 4.6+ and Sonnet 4.6),
|
||||
// budget-based (older models), or explicitly disabled.
|
||||
if (model.reasoning) {
|
||||
if (options?.thinkingEnabled) {
|
||||
// Default to "summarized" so Opus 4.7 and Mythos Preview behave like
|
||||
// older Claude 4 models (whose API default is also "summarized").
|
||||
const display: AnthropicThinkingDisplay = options.thinkingDisplay ?? "summarized";
|
||||
if (supportsAdaptiveThinking(model.id)) {
|
||||
// Adaptive thinking: Claude decides when and how much to think.
|
||||
params.thinking = { type: "adaptive", display };
|
||||
if (options.effort) {
|
||||
// The Anthropic SDK types can lag newly supported effort values such as "xhigh".
|
||||
params.output_config =
|
||||
options.effort === "xhigh"
|
||||
? ({ effort: options.effort } as unknown as NonNullable<
|
||||
MessageCreateParamsStreaming["output_config"]
|
||||
>)
|
||||
: { effort: options.effort };
|
||||
}
|
||||
} else {
|
||||
// Budget-based thinking for older models
|
||||
params.thinking = {
|
||||
type: "enabled",
|
||||
budget_tokens: options.thinkingBudgetTokens || 1024,
|
||||
display,
|
||||
};
|
||||
}
|
||||
} else if (options?.thinkingEnabled === false) {
|
||||
params.thinking = { type: "disabled" };
|
||||
}
|
||||
}
|
||||
|
||||
if (options?.metadata) {
|
||||
const userId = options.metadata.user_id;
|
||||
if (typeof userId === "string") {
|
||||
params.metadata = { user_id: userId };
|
||||
}
|
||||
}
|
||||
|
||||
if (options?.toolChoice) {
|
||||
if (typeof options.toolChoice === "string") {
|
||||
params.tool_choice = { type: options.toolChoice };
|
||||
} else {
|
||||
params.tool_choice = options.toolChoice;
|
||||
}
|
||||
}
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
// Normalize tool call IDs to match Anthropic's required pattern and length
|
||||
function normalizeToolCallId(id: string): string {
|
||||
return id.replace(/[^a-zA-Z0-9_-]/g, "_").slice(0, 64);
|
||||
}
|
||||
|
||||
function convertMessages(
|
||||
messages: Message[],
|
||||
model: Model<"anthropic-messages">,
|
||||
isOAuthToken: boolean,
|
||||
cacheControl?: CacheControlEphemeral,
|
||||
): MessageParam[] {
|
||||
const params: MessageParam[] = [];
|
||||
|
||||
// Transform messages for cross-provider compatibility
|
||||
const transformedMessages = transformMessages(messages, model, normalizeToolCallId);
|
||||
|
||||
for (let i = 0; i < transformedMessages.length; i++) {
|
||||
const msg = transformedMessages[i];
|
||||
|
||||
if (msg.role === "user") {
|
||||
if (typeof msg.content === "string") {
|
||||
if (msg.content.trim().length > 0) {
|
||||
params.push({
|
||||
role: "user",
|
||||
content: sanitizeSurrogates(msg.content),
|
||||
});
|
||||
}
|
||||
} else {
|
||||
const blocks: ContentBlockParam[] = msg.content.map((item) => {
|
||||
if (item.type === "text") {
|
||||
return {
|
||||
type: "text",
|
||||
text: sanitizeSurrogates(item.text),
|
||||
};
|
||||
} else {
|
||||
return {
|
||||
type: "image",
|
||||
source: {
|
||||
type: "base64",
|
||||
media_type: item.mimeType as "image/jpeg" | "image/png" | "image/gif" | "image/webp",
|
||||
data: item.data,
|
||||
},
|
||||
};
|
||||
}
|
||||
});
|
||||
let filteredBlocks = !model?.input.includes("image") ? blocks.filter((b) => b.type !== "image") : blocks;
|
||||
filteredBlocks = filteredBlocks.filter((b) => {
|
||||
if (b.type === "text") {
|
||||
return b.text.trim().length > 0;
|
||||
}
|
||||
return true;
|
||||
});
|
||||
if (filteredBlocks.length === 0) continue;
|
||||
params.push({
|
||||
role: "user",
|
||||
content: filteredBlocks,
|
||||
});
|
||||
}
|
||||
} else if (msg.role === "assistant") {
|
||||
const blocks: ContentBlockParam[] = [];
|
||||
|
||||
for (const block of msg.content) {
|
||||
if (block.type === "text") {
|
||||
if (block.text.trim().length === 0) continue;
|
||||
blocks.push({
|
||||
type: "text",
|
||||
text: sanitizeSurrogates(block.text),
|
||||
});
|
||||
} else if (block.type === "thinking") {
|
||||
// Redacted thinking: pass the opaque payload back as redacted_thinking
|
||||
if (block.redacted) {
|
||||
blocks.push({
|
||||
type: "redacted_thinking",
|
||||
data: block.thinkingSignature!,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
if (block.thinking.trim().length === 0) continue;
|
||||
// If thinking signature is missing/empty (e.g., from aborted stream),
|
||||
// convert to plain text block without <thinking> tags to avoid API rejection
|
||||
// and prevent Claude from mimicking the tags in responses
|
||||
if (!block.thinkingSignature || block.thinkingSignature.trim().length === 0) {
|
||||
blocks.push({
|
||||
type: "text",
|
||||
text: sanitizeSurrogates(block.thinking),
|
||||
});
|
||||
} else {
|
||||
blocks.push({
|
||||
type: "thinking",
|
||||
thinking: sanitizeSurrogates(block.thinking),
|
||||
signature: block.thinkingSignature,
|
||||
});
|
||||
}
|
||||
} else if (block.type === "toolCall") {
|
||||
blocks.push({
|
||||
type: "tool_use",
|
||||
id: block.id,
|
||||
name: isOAuthToken ? toClaudeCodeName(block.name) : block.name,
|
||||
input: block.arguments ?? {},
|
||||
});
|
||||
}
|
||||
}
|
||||
if (blocks.length === 0) continue;
|
||||
params.push({
|
||||
role: "assistant",
|
||||
content: blocks,
|
||||
});
|
||||
} else if (msg.role === "toolResult") {
|
||||
// Collect all consecutive toolResult messages, needed for z.ai Anthropic endpoint
|
||||
const toolResults: ContentBlockParam[] = [];
|
||||
|
||||
// Add the current tool result
|
||||
toolResults.push({
|
||||
type: "tool_result",
|
||||
tool_use_id: msg.toolCallId,
|
||||
content: convertContentBlocks(msg.content),
|
||||
is_error: msg.isError,
|
||||
});
|
||||
|
||||
// Look ahead for consecutive toolResult messages
|
||||
let j = i + 1;
|
||||
while (j < transformedMessages.length && transformedMessages[j].role === "toolResult") {
|
||||
const nextMsg = transformedMessages[j] as ToolResultMessage; // We know it's a toolResult
|
||||
toolResults.push({
|
||||
type: "tool_result",
|
||||
tool_use_id: nextMsg.toolCallId,
|
||||
content: convertContentBlocks(nextMsg.content),
|
||||
is_error: nextMsg.isError,
|
||||
});
|
||||
j++;
|
||||
}
|
||||
|
||||
// Skip the messages we've already processed
|
||||
i = j - 1;
|
||||
|
||||
// Add a single user message with all tool results
|
||||
params.push({
|
||||
role: "user",
|
||||
content: toolResults,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Add cache_control to the last user message to cache conversation history
|
||||
if (cacheControl && params.length > 0) {
|
||||
const lastMessage = params[params.length - 1];
|
||||
if (lastMessage.role === "user") {
|
||||
if (Array.isArray(lastMessage.content)) {
|
||||
const lastBlock = lastMessage.content[lastMessage.content.length - 1];
|
||||
if (
|
||||
lastBlock &&
|
||||
(lastBlock.type === "text" || lastBlock.type === "image" || lastBlock.type === "tool_result")
|
||||
) {
|
||||
(lastBlock as any).cache_control = cacheControl;
|
||||
}
|
||||
} else if (typeof lastMessage.content === "string") {
|
||||
lastMessage.content = [
|
||||
{
|
||||
type: "text",
|
||||
text: lastMessage.content,
|
||||
cache_control: cacheControl,
|
||||
},
|
||||
] as any;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
function convertTools(
|
||||
tools: Tool[],
|
||||
isOAuthToken: boolean,
|
||||
cacheControl?: CacheControlEphemeral,
|
||||
): Anthropic.Messages.Tool[] {
|
||||
if (!tools) return [];
|
||||
|
||||
return tools.map((tool, index) => {
|
||||
const schema = tool.parameters as { properties?: unknown; required?: string[] };
|
||||
|
||||
return {
|
||||
name: isOAuthToken ? toClaudeCodeName(tool.name) : tool.name,
|
||||
description: tool.description,
|
||||
input_schema: {
|
||||
type: "object",
|
||||
properties: schema.properties ?? {},
|
||||
required: schema.required ?? [],
|
||||
},
|
||||
...(cacheControl && index === tools.length - 1 ? { cache_control: cacheControl } : {}),
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
function mapStopReason(reason: Anthropic.Messages.StopReason | string): StopReason {
|
||||
switch (reason) {
|
||||
case "end_turn":
|
||||
return "stop";
|
||||
case "max_tokens":
|
||||
return "length";
|
||||
case "tool_use":
|
||||
return "toolUse";
|
||||
case "refusal":
|
||||
return "error";
|
||||
case "pause_turn": // Stop is good enough -> resubmit
|
||||
return "stop";
|
||||
case "stop_sequence":
|
||||
return "stop"; // We don't supply stop sequences, so this should never happen
|
||||
case "sensitive": // Content flagged by safety filters (not yet in SDK types)
|
||||
return "error";
|
||||
default:
|
||||
// Handle unknown stop reasons gracefully (API may add new values)
|
||||
throw new Error(`Unhandled stop reason: ${reason}`);
|
||||
}
|
||||
}
|
||||
@@ -1,253 +0,0 @@
|
||||
import { AzureOpenAI } from "openai";
|
||||
import type { ResponseCreateParamsStreaming } from "openai/resources/responses/responses.js";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { supportsXhigh } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
Context,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
|
||||
const DEFAULT_AZURE_API_VERSION = "v1";
|
||||
const AZURE_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode", "azure-openai-responses"]);
|
||||
|
||||
function parseDeploymentNameMap(value: string | undefined): Map<string, string> {
|
||||
const map = new Map<string, string>();
|
||||
if (!value) return map;
|
||||
for (const entry of value.split(",")) {
|
||||
const trimmed = entry.trim();
|
||||
if (!trimmed) continue;
|
||||
const [modelId, deploymentName] = trimmed.split("=", 2);
|
||||
if (!modelId || !deploymentName) continue;
|
||||
map.set(modelId.trim(), deploymentName.trim());
|
||||
}
|
||||
return map;
|
||||
}
|
||||
|
||||
function resolveDeploymentName(model: Model<"azure-openai-responses">, options?: AzureOpenAIResponsesOptions): string {
|
||||
if (options?.azureDeploymentName) {
|
||||
return options.azureDeploymentName;
|
||||
}
|
||||
const mappedDeployment = parseDeploymentNameMap(process.env.AZURE_OPENAI_DEPLOYMENT_NAME_MAP).get(model.id);
|
||||
return mappedDeployment || model.id;
|
||||
}
|
||||
|
||||
// Azure OpenAI Responses-specific options
|
||||
export interface AzureOpenAIResponsesOptions extends StreamOptions {
|
||||
reasoningEffort?: "minimal" | "low" | "medium" | "high" | "xhigh";
|
||||
reasoningSummary?: "auto" | "detailed" | "concise" | null;
|
||||
azureApiVersion?: string;
|
||||
azureResourceName?: string;
|
||||
azureBaseUrl?: string;
|
||||
azureDeploymentName?: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate function for Azure OpenAI Responses API
|
||||
*/
|
||||
export const streamAzureOpenAIResponses: StreamFunction<"azure-openai-responses", AzureOpenAIResponsesOptions> = (
|
||||
model: Model<"azure-openai-responses">,
|
||||
context: Context,
|
||||
options?: AzureOpenAIResponsesOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
// Start async processing
|
||||
(async () => {
|
||||
const deploymentName = resolveDeploymentName(model, options);
|
||||
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: "azure-openai-responses" as Api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
try {
|
||||
// Create Azure OpenAI client
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider) || "";
|
||||
const client = createClient(model, apiKey, options);
|
||||
let params = buildParams(model, context, options, deploymentName);
|
||||
const nextParams = await options?.onPayload?.(params, model);
|
||||
if (nextParams !== undefined) {
|
||||
params = nextParams as ResponseCreateParamsStreaming;
|
||||
}
|
||||
const { data: openaiStream, response } = await client.responses
|
||||
.create(params, options?.signal ? { signal: options.signal } : undefined)
|
||||
.withResponse();
|
||||
await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model);
|
||||
stream.push({ type: "start", partial: output });
|
||||
|
||||
await processResponsesStream(openaiStream, output, stream, model);
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "aborted" || output.stopReason === "error") {
|
||||
throw new Error("An unknown error occurred");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
for (const block of output.content) {
|
||||
delete (block as { index?: number }).index;
|
||||
// partialJson is only a streaming scratch buffer; never persist it.
|
||||
delete (block as { partialJson?: string }).partialJson;
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
export const streamSimpleAzureOpenAIResponses: StreamFunction<"azure-openai-responses", SimpleStreamOptions> = (
|
||||
model: Model<"azure-openai-responses">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
|
||||
if (!apiKey) {
|
||||
throw new Error(`No API key for provider: ${model.provider}`);
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoningEffort = supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning);
|
||||
|
||||
return streamAzureOpenAIResponses(model, context, {
|
||||
...base,
|
||||
reasoningEffort,
|
||||
} satisfies AzureOpenAIResponsesOptions);
|
||||
};
|
||||
|
||||
function normalizeAzureBaseUrl(baseUrl: string): string {
|
||||
return baseUrl.replace(/\/+$/, "");
|
||||
}
|
||||
|
||||
function buildDefaultBaseUrl(resourceName: string): string {
|
||||
return `https://${resourceName}.openai.azure.com/openai/v1`;
|
||||
}
|
||||
|
||||
function resolveAzureConfig(
|
||||
model: Model<"azure-openai-responses">,
|
||||
options?: AzureOpenAIResponsesOptions,
|
||||
): { baseUrl: string; apiVersion: string } {
|
||||
const apiVersion = options?.azureApiVersion || process.env.AZURE_OPENAI_API_VERSION || DEFAULT_AZURE_API_VERSION;
|
||||
|
||||
const baseUrl = options?.azureBaseUrl?.trim() || process.env.AZURE_OPENAI_BASE_URL?.trim() || undefined;
|
||||
const resourceName = options?.azureResourceName || process.env.AZURE_OPENAI_RESOURCE_NAME;
|
||||
|
||||
let resolvedBaseUrl = baseUrl;
|
||||
|
||||
if (!resolvedBaseUrl && resourceName) {
|
||||
resolvedBaseUrl = buildDefaultBaseUrl(resourceName);
|
||||
}
|
||||
|
||||
if (!resolvedBaseUrl && model.baseUrl) {
|
||||
resolvedBaseUrl = model.baseUrl;
|
||||
}
|
||||
|
||||
if (!resolvedBaseUrl) {
|
||||
throw new Error(
|
||||
"Azure OpenAI base URL is required. Set AZURE_OPENAI_BASE_URL or AZURE_OPENAI_RESOURCE_NAME, or pass azureBaseUrl, azureResourceName, or model.baseUrl.",
|
||||
);
|
||||
}
|
||||
|
||||
return {
|
||||
baseUrl: normalizeAzureBaseUrl(resolvedBaseUrl),
|
||||
apiVersion,
|
||||
};
|
||||
}
|
||||
|
||||
function createClient(model: Model<"azure-openai-responses">, apiKey: string, options?: AzureOpenAIResponsesOptions) {
|
||||
if (!apiKey) {
|
||||
if (!process.env.AZURE_OPENAI_API_KEY) {
|
||||
throw new Error(
|
||||
"Azure OpenAI API key is required. Set AZURE_OPENAI_API_KEY environment variable or pass it as an argument.",
|
||||
);
|
||||
}
|
||||
apiKey = process.env.AZURE_OPENAI_API_KEY;
|
||||
}
|
||||
|
||||
const headers = { ...model.headers };
|
||||
|
||||
if (options?.headers) {
|
||||
Object.assign(headers, options.headers);
|
||||
}
|
||||
|
||||
const { baseUrl, apiVersion } = resolveAzureConfig(model, options);
|
||||
|
||||
return new AzureOpenAI({
|
||||
apiKey,
|
||||
apiVersion,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: headers,
|
||||
baseURL: baseUrl,
|
||||
});
|
||||
}
|
||||
|
||||
function buildParams(
|
||||
model: Model<"azure-openai-responses">,
|
||||
context: Context,
|
||||
options: AzureOpenAIResponsesOptions | undefined,
|
||||
deploymentName: string,
|
||||
) {
|
||||
const messages = convertResponsesMessages(model, context, AZURE_TOOL_CALL_PROVIDERS);
|
||||
|
||||
const params: ResponseCreateParamsStreaming = {
|
||||
model: deploymentName,
|
||||
input: messages,
|
||||
stream: true,
|
||||
prompt_cache_key: options?.sessionId,
|
||||
};
|
||||
|
||||
if (options?.maxTokens) {
|
||||
params.max_output_tokens = options?.maxTokens;
|
||||
}
|
||||
|
||||
if (options?.temperature !== undefined) {
|
||||
params.temperature = options?.temperature;
|
||||
}
|
||||
|
||||
if (context.tools) {
|
||||
params.tools = convertResponsesTools(context.tools);
|
||||
}
|
||||
|
||||
if (model.reasoning) {
|
||||
if (options?.reasoningEffort || options?.reasoningSummary) {
|
||||
params.reasoning = {
|
||||
effort: options?.reasoningEffort || "medium",
|
||||
summary: options?.reasoningSummary || "auto",
|
||||
};
|
||||
params.include = ["reasoning.encrypted_content"];
|
||||
} else {
|
||||
params.reasoning = { effort: "none" };
|
||||
}
|
||||
}
|
||||
|
||||
return params;
|
||||
}
|
||||
@@ -1,499 +0,0 @@
|
||||
import { registerApiProvider, unregisterApiProviders } from "../api-registry.js";
|
||||
import type {
|
||||
AssistantMessage,
|
||||
AssistantMessageEventStream,
|
||||
Context,
|
||||
ImageContent,
|
||||
Message,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
TextContent,
|
||||
ThinkingContent,
|
||||
ToolCall,
|
||||
ToolResultMessage,
|
||||
Usage,
|
||||
} from "../types.js";
|
||||
import { createAssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
|
||||
const DEFAULT_API = "faux";
|
||||
const DEFAULT_PROVIDER = "faux";
|
||||
const DEFAULT_MODEL_ID = "faux-1";
|
||||
const DEFAULT_MODEL_NAME = "Faux Model";
|
||||
const DEFAULT_BASE_URL = "http://localhost:0";
|
||||
const DEFAULT_MIN_TOKEN_SIZE = 3;
|
||||
const DEFAULT_MAX_TOKEN_SIZE = 5;
|
||||
|
||||
const DEFAULT_USAGE: Usage = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
|
||||
export interface FauxModelDefinition {
|
||||
id: string;
|
||||
name?: string;
|
||||
reasoning?: boolean;
|
||||
input?: ("text" | "image")[];
|
||||
cost?: { input: number; output: number; cacheRead: number; cacheWrite: number };
|
||||
contextWindow?: number;
|
||||
maxTokens?: number;
|
||||
}
|
||||
|
||||
export type FauxContentBlock = TextContent | ThinkingContent | ToolCall;
|
||||
|
||||
export function fauxText(text: string): TextContent {
|
||||
return { type: "text", text };
|
||||
}
|
||||
|
||||
export function fauxThinking(thinking: string): ThinkingContent {
|
||||
return { type: "thinking", thinking };
|
||||
}
|
||||
|
||||
export function fauxToolCall(name: string, arguments_: ToolCall["arguments"], options: { id?: string } = {}): ToolCall {
|
||||
return {
|
||||
type: "toolCall",
|
||||
id: options.id ?? randomId("tool"),
|
||||
name,
|
||||
arguments: arguments_,
|
||||
};
|
||||
}
|
||||
|
||||
function normalizeFauxAssistantContent(content: string | FauxContentBlock | FauxContentBlock[]): FauxContentBlock[] {
|
||||
if (typeof content === "string") {
|
||||
return [fauxText(content)];
|
||||
}
|
||||
return Array.isArray(content) ? content : [content];
|
||||
}
|
||||
|
||||
export function fauxAssistantMessage(
|
||||
content: string | FauxContentBlock | FauxContentBlock[],
|
||||
options: {
|
||||
stopReason?: AssistantMessage["stopReason"];
|
||||
errorMessage?: string;
|
||||
responseId?: string;
|
||||
timestamp?: number;
|
||||
} = {},
|
||||
): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: normalizeFauxAssistantContent(content),
|
||||
api: DEFAULT_API,
|
||||
provider: DEFAULT_PROVIDER,
|
||||
model: DEFAULT_MODEL_ID,
|
||||
usage: DEFAULT_USAGE,
|
||||
stopReason: options.stopReason ?? "stop",
|
||||
errorMessage: options.errorMessage,
|
||||
responseId: options.responseId,
|
||||
timestamp: options.timestamp ?? Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
export type FauxResponseFactory = (
|
||||
context: Context,
|
||||
options: StreamOptions | undefined,
|
||||
state: { callCount: number },
|
||||
model: Model<string>,
|
||||
) => AssistantMessage | Promise<AssistantMessage>;
|
||||
|
||||
export type FauxResponseStep = AssistantMessage | FauxResponseFactory;
|
||||
|
||||
export interface RegisterFauxProviderOptions {
|
||||
api?: string;
|
||||
provider?: string;
|
||||
models?: FauxModelDefinition[];
|
||||
tokensPerSecond?: number;
|
||||
tokenSize?: {
|
||||
min?: number;
|
||||
max?: number;
|
||||
};
|
||||
}
|
||||
|
||||
export interface FauxProviderRegistration {
|
||||
api: string;
|
||||
models: [Model<string>, ...Model<string>[]];
|
||||
getModel(): Model<string>;
|
||||
getModel(modelId: string): Model<string> | undefined;
|
||||
state: { callCount: number };
|
||||
setResponses: (responses: FauxResponseStep[]) => void;
|
||||
appendResponses: (responses: FauxResponseStep[]) => void;
|
||||
getPendingResponseCount: () => number;
|
||||
unregister: () => void;
|
||||
}
|
||||
|
||||
function estimateTokens(text: string): number {
|
||||
return Math.ceil(text.length / 4);
|
||||
}
|
||||
|
||||
function randomId(prefix: string): string {
|
||||
return `${prefix}:${Date.now()}:${Math.random().toString(36).slice(2)}`;
|
||||
}
|
||||
|
||||
function contentToText(content: string | Array<TextContent | ImageContent>): string {
|
||||
if (typeof content === "string") {
|
||||
return content;
|
||||
}
|
||||
return content
|
||||
.map((block) => {
|
||||
if (block.type === "text") {
|
||||
return block.text;
|
||||
}
|
||||
return `[image:${block.mimeType}:${block.data.length}]`;
|
||||
})
|
||||
.join("\n");
|
||||
}
|
||||
|
||||
function assistantContentToText(content: Array<TextContent | ThinkingContent | ToolCall>): string {
|
||||
return content
|
||||
.map((block) => {
|
||||
if (block.type === "text") {
|
||||
return block.text;
|
||||
}
|
||||
if (block.type === "thinking") {
|
||||
return block.thinking;
|
||||
}
|
||||
return `${block.name}:${JSON.stringify(block.arguments)}`;
|
||||
})
|
||||
.join("\n");
|
||||
}
|
||||
|
||||
function toolResultToText(message: ToolResultMessage): string {
|
||||
return [message.toolName, ...message.content.map((block) => contentToText([block]))].join("\n");
|
||||
}
|
||||
|
||||
function messageToText(message: Message): string {
|
||||
if (message.role === "user") {
|
||||
return contentToText(message.content);
|
||||
}
|
||||
if (message.role === "assistant") {
|
||||
return assistantContentToText(message.content);
|
||||
}
|
||||
return toolResultToText(message);
|
||||
}
|
||||
|
||||
function serializeContext(context: Context): string {
|
||||
const parts: string[] = [];
|
||||
if (context.systemPrompt) {
|
||||
parts.push(`system:${context.systemPrompt}`);
|
||||
}
|
||||
for (const message of context.messages) {
|
||||
parts.push(`${message.role}:${messageToText(message)}`);
|
||||
}
|
||||
if (context.tools?.length) {
|
||||
parts.push(`tools:${JSON.stringify(context.tools)}`);
|
||||
}
|
||||
return parts.join("\n\n");
|
||||
}
|
||||
|
||||
function commonPrefixLength(a: string, b: string): number {
|
||||
const length = Math.min(a.length, b.length);
|
||||
let index = 0;
|
||||
while (index < length && a[index] === b[index]) {
|
||||
index++;
|
||||
}
|
||||
return index;
|
||||
}
|
||||
|
||||
function withUsageEstimate(
|
||||
message: AssistantMessage,
|
||||
context: Context,
|
||||
options: StreamOptions | undefined,
|
||||
promptCache: Map<string, string>,
|
||||
): AssistantMessage {
|
||||
const promptText = serializeContext(context);
|
||||
const promptTokens = estimateTokens(promptText);
|
||||
const outputTokens = estimateTokens(assistantContentToText(message.content));
|
||||
let input = promptTokens;
|
||||
let cacheRead = 0;
|
||||
let cacheWrite = 0;
|
||||
const sessionId = options?.sessionId;
|
||||
|
||||
if (sessionId && options?.cacheRetention !== "none") {
|
||||
const previousPrompt = promptCache.get(sessionId);
|
||||
if (previousPrompt) {
|
||||
const cachedChars = commonPrefixLength(previousPrompt, promptText);
|
||||
cacheRead = estimateTokens(previousPrompt.slice(0, cachedChars));
|
||||
cacheWrite = estimateTokens(promptText.slice(cachedChars));
|
||||
input = Math.max(0, promptTokens - cacheRead);
|
||||
} else {
|
||||
cacheWrite = promptTokens;
|
||||
}
|
||||
promptCache.set(sessionId, promptText);
|
||||
}
|
||||
|
||||
return {
|
||||
...message,
|
||||
usage: {
|
||||
input,
|
||||
output: outputTokens,
|
||||
cacheRead,
|
||||
cacheWrite,
|
||||
totalTokens: input + outputTokens + cacheRead + cacheWrite,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function splitStringByTokenSize(text: string, minTokenSize: number, maxTokenSize: number): string[] {
|
||||
const chunks: string[] = [];
|
||||
let index = 0;
|
||||
while (index < text.length) {
|
||||
const tokenSize = minTokenSize + Math.floor(Math.random() * (maxTokenSize - minTokenSize + 1));
|
||||
const charSize = Math.max(1, tokenSize * 4);
|
||||
chunks.push(text.slice(index, index + charSize));
|
||||
index += charSize;
|
||||
}
|
||||
return chunks.length > 0 ? chunks : [""];
|
||||
}
|
||||
|
||||
function cloneMessage(message: AssistantMessage, api: string, provider: string, modelId: string): AssistantMessage {
|
||||
const cloned = structuredClone(message);
|
||||
return {
|
||||
...cloned,
|
||||
api,
|
||||
provider,
|
||||
model: modelId,
|
||||
timestamp: cloned.timestamp ?? Date.now(),
|
||||
usage: cloned.usage ?? DEFAULT_USAGE,
|
||||
};
|
||||
}
|
||||
|
||||
function createErrorMessage(error: unknown, api: string, provider: string, modelId: string): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api,
|
||||
provider,
|
||||
model: modelId,
|
||||
usage: DEFAULT_USAGE,
|
||||
stopReason: "error",
|
||||
errorMessage: error instanceof Error ? error.message : String(error),
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createAbortedMessage(partial: AssistantMessage): AssistantMessage {
|
||||
return {
|
||||
...partial,
|
||||
stopReason: "aborted",
|
||||
errorMessage: "Request was aborted",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function scheduleChunk(chunk: string, tokensPerSecond: number | undefined): Promise<void> {
|
||||
if (!tokensPerSecond || tokensPerSecond <= 0) {
|
||||
return new Promise((resolve) => queueMicrotask(resolve));
|
||||
}
|
||||
const delayMs = (estimateTokens(chunk) / tokensPerSecond) * 1000;
|
||||
return new Promise((resolve) => setTimeout(resolve, delayMs));
|
||||
}
|
||||
|
||||
async function streamWithDeltas(
|
||||
stream: AssistantMessageEventStream,
|
||||
message: AssistantMessage,
|
||||
minTokenSize: number,
|
||||
maxTokenSize: number,
|
||||
tokensPerSecond: number | undefined,
|
||||
signal: AbortSignal | undefined,
|
||||
): Promise<void> {
|
||||
const partial: AssistantMessage = { ...message, content: [] };
|
||||
if (signal?.aborted) {
|
||||
const aborted = createAbortedMessage(partial);
|
||||
stream.push({ type: "error", reason: "aborted", error: aborted });
|
||||
stream.end(aborted);
|
||||
return;
|
||||
}
|
||||
|
||||
stream.push({ type: "start", partial: { ...partial } });
|
||||
|
||||
for (let index = 0; index < message.content.length; index++) {
|
||||
if (signal?.aborted) {
|
||||
const aborted = createAbortedMessage(partial);
|
||||
stream.push({ type: "error", reason: "aborted", error: aborted });
|
||||
stream.end(aborted);
|
||||
return;
|
||||
}
|
||||
|
||||
const block = message.content[index];
|
||||
|
||||
if (block.type === "thinking") {
|
||||
partial.content = [...partial.content, { type: "thinking", thinking: "" }];
|
||||
stream.push({ type: "thinking_start", contentIndex: index, partial: { ...partial } });
|
||||
for (const chunk of splitStringByTokenSize(block.thinking, minTokenSize, maxTokenSize)) {
|
||||
await scheduleChunk(chunk, tokensPerSecond);
|
||||
if (signal?.aborted) {
|
||||
const aborted = createAbortedMessage(partial);
|
||||
stream.push({ type: "error", reason: "aborted", error: aborted });
|
||||
stream.end(aborted);
|
||||
return;
|
||||
}
|
||||
(partial.content[index] as ThinkingContent).thinking += chunk;
|
||||
stream.push({ type: "thinking_delta", contentIndex: index, delta: chunk, partial: { ...partial } });
|
||||
}
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: index,
|
||||
content: block.thinking,
|
||||
partial: { ...partial },
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
if (block.type === "text") {
|
||||
partial.content = [...partial.content, { type: "text", text: "" }];
|
||||
stream.push({ type: "text_start", contentIndex: index, partial: { ...partial } });
|
||||
for (const chunk of splitStringByTokenSize(block.text, minTokenSize, maxTokenSize)) {
|
||||
await scheduleChunk(chunk, tokensPerSecond);
|
||||
if (signal?.aborted) {
|
||||
const aborted = createAbortedMessage(partial);
|
||||
stream.push({ type: "error", reason: "aborted", error: aborted });
|
||||
stream.end(aborted);
|
||||
return;
|
||||
}
|
||||
(partial.content[index] as TextContent).text += chunk;
|
||||
stream.push({ type: "text_delta", contentIndex: index, delta: chunk, partial: { ...partial } });
|
||||
}
|
||||
stream.push({ type: "text_end", contentIndex: index, content: block.text, partial: { ...partial } });
|
||||
continue;
|
||||
}
|
||||
|
||||
partial.content = [...partial.content, { type: "toolCall", id: block.id, name: block.name, arguments: {} }];
|
||||
stream.push({ type: "toolcall_start", contentIndex: index, partial: { ...partial } });
|
||||
for (const chunk of splitStringByTokenSize(JSON.stringify(block.arguments), minTokenSize, maxTokenSize)) {
|
||||
await scheduleChunk(chunk, tokensPerSecond);
|
||||
if (signal?.aborted) {
|
||||
const aborted = createAbortedMessage(partial);
|
||||
stream.push({ type: "error", reason: "aborted", error: aborted });
|
||||
stream.end(aborted);
|
||||
return;
|
||||
}
|
||||
stream.push({ type: "toolcall_delta", contentIndex: index, delta: chunk, partial: { ...partial } });
|
||||
}
|
||||
(partial.content[index] as ToolCall).arguments = block.arguments;
|
||||
stream.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: { ...partial } });
|
||||
}
|
||||
|
||||
if (message.stopReason === "error" || message.stopReason === "aborted") {
|
||||
stream.push({ type: "error", reason: message.stopReason, error: message });
|
||||
stream.end(message);
|
||||
return;
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: message.stopReason, message });
|
||||
stream.end(message);
|
||||
}
|
||||
|
||||
export function registerFauxProvider(options: RegisterFauxProviderOptions = {}): FauxProviderRegistration {
|
||||
const api = options.api ?? randomId(DEFAULT_API);
|
||||
const provider = options.provider ?? DEFAULT_PROVIDER;
|
||||
const sourceId = randomId("faux-provider");
|
||||
const minTokenSize = Math.max(
|
||||
1,
|
||||
Math.min(options.tokenSize?.min ?? DEFAULT_MIN_TOKEN_SIZE, options.tokenSize?.max ?? DEFAULT_MAX_TOKEN_SIZE),
|
||||
);
|
||||
const maxTokenSize = Math.max(minTokenSize, options.tokenSize?.max ?? DEFAULT_MAX_TOKEN_SIZE);
|
||||
let pendingResponses: FauxResponseStep[] = [];
|
||||
const tokensPerSecond = options.tokensPerSecond;
|
||||
const state = { callCount: 0 };
|
||||
const promptCache = new Map<string, string>();
|
||||
|
||||
const modelDefinitions = options.models?.length
|
||||
? options.models
|
||||
: [
|
||||
{
|
||||
id: DEFAULT_MODEL_ID,
|
||||
name: DEFAULT_MODEL_NAME,
|
||||
reasoning: false,
|
||||
input: ["text", "image"] as ("text" | "image")[],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 16384,
|
||||
},
|
||||
];
|
||||
const models = modelDefinitions.map((definition) => ({
|
||||
id: definition.id,
|
||||
name: definition.name ?? definition.id,
|
||||
api,
|
||||
provider,
|
||||
baseUrl: DEFAULT_BASE_URL,
|
||||
reasoning: definition.reasoning ?? false,
|
||||
input: definition.input ?? ["text", "image"],
|
||||
cost: definition.cost ?? { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: definition.contextWindow ?? 128000,
|
||||
maxTokens: definition.maxTokens ?? 16384,
|
||||
})) as [Model<string>, ...Model<string>[]];
|
||||
|
||||
const stream: StreamFunction<string, StreamOptions> = (requestModel, context, streamOptions) => {
|
||||
const outer = createAssistantMessageEventStream();
|
||||
const step = pendingResponses.shift();
|
||||
state.callCount++;
|
||||
|
||||
queueMicrotask(async () => {
|
||||
try {
|
||||
await streamOptions?.onResponse?.({ status: 200, headers: {} }, requestModel);
|
||||
if (!step) {
|
||||
let message = createErrorMessage(
|
||||
new Error("No more faux responses queued"),
|
||||
api,
|
||||
provider,
|
||||
requestModel.id,
|
||||
);
|
||||
message = withUsageEstimate(message, context, streamOptions, promptCache);
|
||||
outer.push({ type: "error", reason: "error", error: message });
|
||||
outer.end(message);
|
||||
return;
|
||||
}
|
||||
|
||||
const resolved =
|
||||
typeof step === "function" ? await step(context, streamOptions, state, requestModel) : step;
|
||||
let message = cloneMessage(resolved, api, provider, requestModel.id);
|
||||
message = withUsageEstimate(message, context, streamOptions, promptCache);
|
||||
await streamWithDeltas(outer, message, minTokenSize, maxTokenSize, tokensPerSecond, streamOptions?.signal);
|
||||
} catch (error) {
|
||||
const message = createErrorMessage(error, api, provider, requestModel.id);
|
||||
outer.push({ type: "error", reason: "error", error: message });
|
||||
outer.end(message);
|
||||
}
|
||||
});
|
||||
|
||||
return outer;
|
||||
};
|
||||
|
||||
const streamSimple: StreamFunction<string, SimpleStreamOptions> = (streamModel, context, streamOptions) =>
|
||||
stream(streamModel, context, streamOptions);
|
||||
|
||||
registerApiProvider({ api, stream, streamSimple }, sourceId);
|
||||
|
||||
function getModel(): Model<string>;
|
||||
function getModel(requestedModelId: string): Model<string> | undefined;
|
||||
function getModel(requestedModelId?: string): Model<string> | undefined {
|
||||
if (!requestedModelId) {
|
||||
return models[0];
|
||||
}
|
||||
return models.find((candidate) => candidate.id === requestedModelId);
|
||||
}
|
||||
|
||||
return {
|
||||
api,
|
||||
models,
|
||||
getModel,
|
||||
state,
|
||||
setResponses(responses) {
|
||||
pendingResponses = [...responses];
|
||||
},
|
||||
appendResponses(responses) {
|
||||
pendingResponses.push(...responses);
|
||||
},
|
||||
getPendingResponseCount() {
|
||||
return pendingResponses.length;
|
||||
},
|
||||
unregister() {
|
||||
unregisterApiProviders(sourceId);
|
||||
},
|
||||
};
|
||||
}
|
||||
@@ -1,37 +0,0 @@
|
||||
import type { Message } from "../types.js";
|
||||
|
||||
// Copilot expects X-Initiator to indicate whether the request is user-initiated
|
||||
// or agent-initiated (e.g. follow-up after assistant/tool messages).
|
||||
export function inferCopilotInitiator(messages: Message[]): "user" | "agent" {
|
||||
const last = messages[messages.length - 1];
|
||||
return last && last.role !== "user" ? "agent" : "user";
|
||||
}
|
||||
|
||||
// Copilot requires Copilot-Vision-Request header when sending images
|
||||
export function hasCopilotVisionInput(messages: Message[]): boolean {
|
||||
return messages.some((msg) => {
|
||||
if (msg.role === "user" && Array.isArray(msg.content)) {
|
||||
return msg.content.some((c) => c.type === "image");
|
||||
}
|
||||
if (msg.role === "toolResult" && Array.isArray(msg.content)) {
|
||||
return msg.content.some((c) => c.type === "image");
|
||||
}
|
||||
return false;
|
||||
});
|
||||
}
|
||||
|
||||
export function buildCopilotDynamicHeaders(params: {
|
||||
messages: Message[];
|
||||
hasImages: boolean;
|
||||
}): Record<string, string> {
|
||||
const headers: Record<string, string> = {
|
||||
"X-Initiator": inferCopilotInitiator(params.messages),
|
||||
"Openai-Intent": "conversation-edits",
|
||||
};
|
||||
|
||||
if (params.hasImages) {
|
||||
headers["Copilot-Vision-Request"] = "true";
|
||||
}
|
||||
|
||||
return headers;
|
||||
}
|
||||
@@ -1,996 +0,0 @@
|
||||
/**
|
||||
* Google Gemini CLI / Antigravity provider.
|
||||
* Shared implementation for both google-gemini-cli and google-antigravity providers.
|
||||
* Uses the Cloud Code Assist API endpoint to access Gemini and Claude models.
|
||||
*/
|
||||
|
||||
import type { Content, ThinkingConfig } from "@google/genai";
|
||||
import { calculateCost } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
Context,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
TextContent,
|
||||
ThinkingBudgets,
|
||||
ThinkingContent,
|
||||
ThinkingLevel,
|
||||
ToolCall,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import {
|
||||
convertMessages,
|
||||
convertTools,
|
||||
isThinkingPart,
|
||||
mapStopReasonString,
|
||||
mapToolChoice,
|
||||
retainThoughtSignature,
|
||||
} from "./google-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
|
||||
/**
|
||||
* Thinking level for Gemini 3 models.
|
||||
* Mirrors Google's ThinkingLevel enum values.
|
||||
*/
|
||||
export type GoogleThinkingLevel = "THINKING_LEVEL_UNSPECIFIED" | "MINIMAL" | "LOW" | "MEDIUM" | "HIGH";
|
||||
|
||||
export interface GoogleGeminiCliOptions extends StreamOptions {
|
||||
toolChoice?: "auto" | "none" | "any";
|
||||
/**
|
||||
* Thinking/reasoning configuration.
|
||||
* - Gemini 2.x models: use `budgetTokens` to set the thinking budget
|
||||
* - Gemini 3 models (gemini-3-pro-*, gemini-3-flash-*): use `level` instead
|
||||
*
|
||||
* When using `streamSimple`, this is handled automatically based on the model.
|
||||
*/
|
||||
thinking?: {
|
||||
enabled: boolean;
|
||||
/** Thinking budget in tokens. Use for Gemini 2.x models. */
|
||||
budgetTokens?: number;
|
||||
/** Thinking level. Use for Gemini 3 models (LOW/HIGH for Pro, MINIMAL/LOW/MEDIUM/HIGH for Flash). */
|
||||
level?: GoogleThinkingLevel;
|
||||
};
|
||||
projectId?: string;
|
||||
}
|
||||
|
||||
const DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com";
|
||||
const ANTIGRAVITY_DAILY_ENDPOINT = "https://daily-cloudcode-pa.sandbox.googleapis.com";
|
||||
const ANTIGRAVITY_AUTOPUSH_ENDPOINT = "https://autopush-cloudcode-pa.sandbox.googleapis.com";
|
||||
const ANTIGRAVITY_ENDPOINT_FALLBACKS = [
|
||||
ANTIGRAVITY_DAILY_ENDPOINT,
|
||||
ANTIGRAVITY_AUTOPUSH_ENDPOINT,
|
||||
DEFAULT_ENDPOINT,
|
||||
] as const;
|
||||
// Headers for Gemini CLI (prod endpoint)
|
||||
const GEMINI_CLI_HEADERS = {
|
||||
"User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1",
|
||||
"X-Goog-Api-Client": "gl-node/22.17.0",
|
||||
"Client-Metadata": JSON.stringify({
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
}),
|
||||
};
|
||||
|
||||
// Headers for Antigravity (sandbox endpoint) - requires specific User-Agent
|
||||
const DEFAULT_ANTIGRAVITY_VERSION = "1.21.9";
|
||||
|
||||
function getAntigravityHeaders() {
|
||||
const version = process.env.PI_AI_ANTIGRAVITY_VERSION || DEFAULT_ANTIGRAVITY_VERSION;
|
||||
return {
|
||||
"User-Agent": `antigravity/${version} darwin/arm64`,
|
||||
};
|
||||
}
|
||||
|
||||
// Antigravity system instruction (compact version from CLIProxyAPI).
|
||||
const ANTIGRAVITY_SYSTEM_INSTRUCTION =
|
||||
"You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding." +
|
||||
"You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question." +
|
||||
"**Absolute paths only**" +
|
||||
"**Proactiveness**";
|
||||
|
||||
// Counter for generating unique tool call IDs
|
||||
let toolCallCounter = 0;
|
||||
|
||||
// Retry configuration
|
||||
const MAX_RETRIES = 3;
|
||||
const BASE_DELAY_MS = 1000;
|
||||
const MAX_EMPTY_STREAM_RETRIES = 2;
|
||||
const EMPTY_STREAM_BASE_DELAY_MS = 500;
|
||||
const CLAUDE_THINKING_BETA_HEADER = "interleaved-thinking-2025-05-14";
|
||||
|
||||
/**
|
||||
* Extract retry delay from Gemini error response (in milliseconds).
|
||||
* Checks headers first (Retry-After, x-ratelimit-reset, x-ratelimit-reset-after),
|
||||
* then parses body patterns like:
|
||||
* - "Your quota will reset after 39s"
|
||||
* - "Your quota will reset after 18h31m10s"
|
||||
* - "Please retry in Xs" or "Please retry in Xms"
|
||||
* - "retryDelay": "34.074824224s" (JSON field)
|
||||
*/
|
||||
export function extractRetryDelay(errorText: string, response?: Response | Headers): number | undefined {
|
||||
const normalizeDelay = (ms: number): number | undefined => (ms > 0 ? Math.ceil(ms + 1000) : undefined);
|
||||
|
||||
const headers = response instanceof Headers ? response : response?.headers;
|
||||
if (headers) {
|
||||
const retryAfter = headers.get("retry-after");
|
||||
if (retryAfter) {
|
||||
const retryAfterSeconds = Number(retryAfter);
|
||||
if (Number.isFinite(retryAfterSeconds)) {
|
||||
const delay = normalizeDelay(retryAfterSeconds * 1000);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
const retryAfterDate = new Date(retryAfter);
|
||||
const retryAfterMs = retryAfterDate.getTime();
|
||||
if (!Number.isNaN(retryAfterMs)) {
|
||||
const delay = normalizeDelay(retryAfterMs - Date.now());
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const rateLimitReset = headers.get("x-ratelimit-reset");
|
||||
if (rateLimitReset) {
|
||||
const resetSeconds = Number.parseInt(rateLimitReset, 10);
|
||||
if (!Number.isNaN(resetSeconds)) {
|
||||
const delay = normalizeDelay(resetSeconds * 1000 - Date.now());
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const rateLimitResetAfter = headers.get("x-ratelimit-reset-after");
|
||||
if (rateLimitResetAfter) {
|
||||
const resetAfterSeconds = Number(rateLimitResetAfter);
|
||||
if (Number.isFinite(resetAfterSeconds)) {
|
||||
const delay = normalizeDelay(resetAfterSeconds * 1000);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Pattern 1: "Your quota will reset after ..." (formats: "18h31m10s", "10m15s", "6s", "39s")
|
||||
const durationMatch = errorText.match(/reset after (?:(\d+)h)?(?:(\d+)m)?(\d+(?:\.\d+)?)s/i);
|
||||
if (durationMatch) {
|
||||
const hours = durationMatch[1] ? parseInt(durationMatch[1], 10) : 0;
|
||||
const minutes = durationMatch[2] ? parseInt(durationMatch[2], 10) : 0;
|
||||
const seconds = parseFloat(durationMatch[3]);
|
||||
if (!Number.isNaN(seconds)) {
|
||||
const totalMs = ((hours * 60 + minutes) * 60 + seconds) * 1000;
|
||||
const delay = normalizeDelay(totalMs);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Pattern 2: "Please retry in X[ms|s]"
|
||||
const retryInMatch = errorText.match(/Please retry in ([0-9.]+)(ms|s)/i);
|
||||
if (retryInMatch?.[1]) {
|
||||
const value = parseFloat(retryInMatch[1]);
|
||||
if (!Number.isNaN(value) && value > 0) {
|
||||
const ms = retryInMatch[2].toLowerCase() === "ms" ? value : value * 1000;
|
||||
const delay = normalizeDelay(ms);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Pattern 3: "retryDelay": "34.074824224s" (JSON field in error details)
|
||||
const retryDelayMatch = errorText.match(/"retryDelay":\s*"([0-9.]+)(ms|s)"/i);
|
||||
if (retryDelayMatch?.[1]) {
|
||||
const value = parseFloat(retryDelayMatch[1]);
|
||||
if (!Number.isNaN(value) && value > 0) {
|
||||
const ms = retryDelayMatch[2].toLowerCase() === "ms" ? value : value * 1000;
|
||||
const delay = normalizeDelay(ms);
|
||||
if (delay !== undefined) {
|
||||
return delay;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function needsClaudeThinkingBetaHeader(model: Model<"google-gemini-cli">): boolean {
|
||||
return model.provider === "google-antigravity" && model.id.startsWith("claude-") && model.reasoning;
|
||||
}
|
||||
|
||||
function isGemini3ProModel(modelId: string): boolean {
|
||||
return /gemini-3(?:\.1)?-pro/.test(modelId.toLowerCase());
|
||||
}
|
||||
|
||||
function isGemini3FlashModel(modelId: string): boolean {
|
||||
return /gemini-3(?:\.1)?-flash/.test(modelId.toLowerCase());
|
||||
}
|
||||
|
||||
function isGemini3Model(modelId: string): boolean {
|
||||
return isGemini3ProModel(modelId) || isGemini3FlashModel(modelId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if an error is retryable (rate limit, server error, network error, etc.)
|
||||
*/
|
||||
function isRetryableError(status: number, errorText: string): boolean {
|
||||
if (status === 429 || status === 500 || status === 502 || status === 503 || status === 504) {
|
||||
return true;
|
||||
}
|
||||
return /resource.?exhausted|rate.?limit|overloaded|service.?unavailable|other.?side.?closed/i.test(errorText);
|
||||
}
|
||||
|
||||
/**
|
||||
* Extract a clean, user-friendly error message from Google API error response.
|
||||
* Parses JSON error responses and returns just the message field.
|
||||
*/
|
||||
function extractErrorMessage(errorText: string): string {
|
||||
try {
|
||||
const parsed = JSON.parse(errorText) as { error?: { message?: string } };
|
||||
if (parsed.error?.message) {
|
||||
return parsed.error.message;
|
||||
}
|
||||
} catch {
|
||||
// Not JSON, return as-is
|
||||
}
|
||||
return errorText;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sleep for a given number of milliseconds, respecting abort signal.
|
||||
*/
|
||||
function sleep(ms: number, signal?: AbortSignal): Promise<void> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (signal?.aborted) {
|
||||
reject(new Error("Request was aborted"));
|
||||
return;
|
||||
}
|
||||
const timeout = setTimeout(resolve, ms);
|
||||
signal?.addEventListener("abort", () => {
|
||||
clearTimeout(timeout);
|
||||
reject(new Error("Request was aborted"));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
interface CloudCodeAssistRequest {
|
||||
project: string;
|
||||
model: string;
|
||||
request: {
|
||||
contents: Content[];
|
||||
sessionId?: string;
|
||||
systemInstruction?: { role?: string; parts: { text: string }[] };
|
||||
generationConfig?: {
|
||||
maxOutputTokens?: number;
|
||||
temperature?: number;
|
||||
thinkingConfig?: ThinkingConfig;
|
||||
};
|
||||
tools?: ReturnType<typeof convertTools>;
|
||||
toolConfig?: {
|
||||
functionCallingConfig: {
|
||||
mode: ReturnType<typeof mapToolChoice>;
|
||||
};
|
||||
};
|
||||
};
|
||||
requestType?: string;
|
||||
userAgent?: string;
|
||||
requestId?: string;
|
||||
}
|
||||
|
||||
interface CloudCodeAssistResponseChunk {
|
||||
response?: {
|
||||
candidates?: Array<{
|
||||
content?: {
|
||||
role: string;
|
||||
parts?: Array<{
|
||||
text?: string;
|
||||
thought?: boolean;
|
||||
thoughtSignature?: string;
|
||||
functionCall?: {
|
||||
name: string;
|
||||
args: Record<string, unknown>;
|
||||
id?: string;
|
||||
};
|
||||
}>;
|
||||
};
|
||||
finishReason?: string;
|
||||
}>;
|
||||
usageMetadata?: {
|
||||
promptTokenCount?: number;
|
||||
candidatesTokenCount?: number;
|
||||
thoughtsTokenCount?: number;
|
||||
totalTokenCount?: number;
|
||||
cachedContentTokenCount?: number;
|
||||
};
|
||||
modelVersion?: string;
|
||||
responseId?: string;
|
||||
};
|
||||
traceId?: string;
|
||||
}
|
||||
|
||||
export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli", GoogleGeminiCliOptions> = (
|
||||
model: Model<"google-gemini-cli">,
|
||||
context: Context,
|
||||
options?: GoogleGeminiCliOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: "google-gemini-cli" as Api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
try {
|
||||
// apiKey is JSON-encoded: { token, projectId }
|
||||
const apiKeyRaw = options?.apiKey;
|
||||
if (!apiKeyRaw) {
|
||||
throw new Error("Google Cloud Code Assist requires OAuth authentication. Use /login to authenticate.");
|
||||
}
|
||||
|
||||
let accessToken: string;
|
||||
let projectId: string;
|
||||
|
||||
try {
|
||||
const parsed = JSON.parse(apiKeyRaw) as { token: string; projectId: string };
|
||||
accessToken = parsed.token;
|
||||
projectId = parsed.projectId;
|
||||
} catch {
|
||||
throw new Error("Invalid Google Cloud Code Assist credentials. Use /login to re-authenticate.");
|
||||
}
|
||||
|
||||
if (!accessToken || !projectId) {
|
||||
throw new Error("Missing token or projectId in Google Cloud credentials. Use /login to re-authenticate.");
|
||||
}
|
||||
|
||||
const isAntigravity = model.provider === "google-antigravity";
|
||||
const baseUrl = model.baseUrl?.trim();
|
||||
const endpoints = baseUrl ? [baseUrl] : isAntigravity ? ANTIGRAVITY_ENDPOINT_FALLBACKS : [DEFAULT_ENDPOINT];
|
||||
|
||||
let requestBody = buildRequest(model, context, projectId, options, isAntigravity);
|
||||
const nextRequestBody = await options?.onPayload?.(requestBody, model);
|
||||
if (nextRequestBody !== undefined) {
|
||||
requestBody = nextRequestBody as CloudCodeAssistRequest;
|
||||
}
|
||||
const headers = isAntigravity ? getAntigravityHeaders() : GEMINI_CLI_HEADERS;
|
||||
|
||||
const requestHeaders = {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
Accept: "text/event-stream",
|
||||
...headers,
|
||||
...(needsClaudeThinkingBetaHeader(model) ? { "anthropic-beta": CLAUDE_THINKING_BETA_HEADER } : {}),
|
||||
...options?.headers,
|
||||
};
|
||||
const requestBodyJson = JSON.stringify(requestBody);
|
||||
|
||||
// Fetch with retry logic for rate limits, transient errors, and endpoint fallbacks.
|
||||
// On 403/404, immediately try the next endpoint (no delay).
|
||||
// On 429/5xx, retry with backoff on the same or next endpoint.
|
||||
let response: Response | undefined;
|
||||
let lastError: Error | undefined;
|
||||
let requestUrl: string | undefined;
|
||||
let endpointIndex = 0;
|
||||
|
||||
for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
try {
|
||||
const endpoint = endpoints[endpointIndex];
|
||||
requestUrl = `${endpoint}/v1internal:streamGenerateContent?alt=sse`;
|
||||
response = await fetch(requestUrl, {
|
||||
method: "POST",
|
||||
headers: requestHeaders,
|
||||
body: requestBodyJson,
|
||||
signal: options?.signal,
|
||||
});
|
||||
await options?.onResponse?.(
|
||||
{ status: response.status, headers: headersToRecord(response.headers) },
|
||||
model,
|
||||
);
|
||||
|
||||
if (response.ok) {
|
||||
break; // Success, exit retry loop
|
||||
}
|
||||
|
||||
const errorText = await response.text();
|
||||
|
||||
// On 403/404, cascade to the next endpoint immediately (no delay)
|
||||
if ((response.status === 403 || response.status === 404) && endpointIndex < endpoints.length - 1) {
|
||||
endpointIndex++;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Check if retryable (429, 5xx, network patterns)
|
||||
if (attempt < MAX_RETRIES && isRetryableError(response.status, errorText)) {
|
||||
// Advance endpoint if possible
|
||||
if (endpointIndex < endpoints.length - 1) {
|
||||
endpointIndex++;
|
||||
}
|
||||
|
||||
// Use server-provided delay or exponential backoff
|
||||
const serverDelay = extractRetryDelay(errorText, response);
|
||||
const delayMs = serverDelay ?? BASE_DELAY_MS * 2 ** attempt;
|
||||
|
||||
// Check if server delay exceeds max allowed (default: 60s)
|
||||
const maxDelayMs = options?.maxRetryDelayMs ?? 60000;
|
||||
if (maxDelayMs > 0 && serverDelay && serverDelay > maxDelayMs) {
|
||||
const delaySeconds = Math.ceil(serverDelay / 1000);
|
||||
throw new Error(
|
||||
`Server requested ${delaySeconds}s retry delay (max: ${Math.ceil(maxDelayMs / 1000)}s). ${extractErrorMessage(errorText)}`,
|
||||
);
|
||||
}
|
||||
|
||||
await sleep(delayMs, options?.signal);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Not retryable or max retries exceeded
|
||||
throw new Error(`Cloud Code Assist API error (${response.status}): ${extractErrorMessage(errorText)}`);
|
||||
} catch (error) {
|
||||
// Check for abort - fetch throws AbortError, our code throws "Request was aborted"
|
||||
if (error instanceof Error) {
|
||||
if (error.name === "AbortError" || error.message === "Request was aborted") {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
}
|
||||
// Extract detailed error message from fetch errors (Node includes cause)
|
||||
lastError = error instanceof Error ? error : new Error(String(error));
|
||||
if (lastError.message === "fetch failed" && lastError.cause instanceof Error) {
|
||||
lastError = new Error(`Network error: ${lastError.cause.message}`);
|
||||
}
|
||||
// Network errors are retryable
|
||||
if (attempt < MAX_RETRIES) {
|
||||
const delayMs = BASE_DELAY_MS * 2 ** attempt;
|
||||
await sleep(delayMs, options?.signal);
|
||||
continue;
|
||||
}
|
||||
throw lastError;
|
||||
}
|
||||
}
|
||||
|
||||
if (!response || !response.ok) {
|
||||
throw lastError ?? new Error("Failed to get response after retries");
|
||||
}
|
||||
|
||||
let started = false;
|
||||
const ensureStarted = () => {
|
||||
if (!started) {
|
||||
stream.push({ type: "start", partial: output });
|
||||
started = true;
|
||||
}
|
||||
};
|
||||
|
||||
const resetOutput = () => {
|
||||
output.content = [];
|
||||
output.usage = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
output.stopReason = "stop";
|
||||
output.errorMessage = undefined;
|
||||
output.timestamp = Date.now();
|
||||
started = false;
|
||||
};
|
||||
|
||||
const streamResponse = async (activeResponse: Response): Promise<boolean> => {
|
||||
if (!activeResponse.body) {
|
||||
throw new Error("No response body");
|
||||
}
|
||||
|
||||
let hasContent = false;
|
||||
let currentBlock: TextContent | ThinkingContent | null = null;
|
||||
const blocks = output.content;
|
||||
const blockIndex = () => blocks.length - 1;
|
||||
|
||||
// Read SSE stream
|
||||
const reader = activeResponse.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
|
||||
// Set up abort handler to cancel reader when signal fires
|
||||
const abortHandler = () => {
|
||||
void reader.cancel().catch(() => {});
|
||||
};
|
||||
options?.signal?.addEventListener("abort", abortHandler);
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
// Check abort signal before each read
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
const lines = buffer.split("\n");
|
||||
buffer = lines.pop() || "";
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.startsWith("data:")) continue;
|
||||
|
||||
const jsonStr = line.slice(5).trim();
|
||||
if (!jsonStr) continue;
|
||||
|
||||
let chunk: CloudCodeAssistResponseChunk;
|
||||
try {
|
||||
chunk = JSON.parse(jsonStr);
|
||||
} catch {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Unwrap the response
|
||||
const responseData = chunk.response;
|
||||
if (!responseData) continue;
|
||||
// Cloud Code Assist mirrors Gemini's responseId field. Keep the first non-empty one.
|
||||
// A single streamed response should retain the same ID across chunks.
|
||||
output.responseId ||= responseData.responseId;
|
||||
|
||||
const candidate = responseData.candidates?.[0];
|
||||
if (candidate?.content?.parts) {
|
||||
for (const part of candidate.content.parts) {
|
||||
if (part.text !== undefined) {
|
||||
hasContent = true;
|
||||
const isThinking = isThinkingPart(part);
|
||||
if (
|
||||
!currentBlock ||
|
||||
(isThinking && currentBlock.type !== "thinking") ||
|
||||
(!isThinking && currentBlock.type !== "text")
|
||||
) {
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blocks.length - 1,
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
if (isThinking) {
|
||||
currentBlock = { type: "thinking", thinking: "", thinkingSignature: undefined };
|
||||
output.content.push(currentBlock);
|
||||
ensureStarted();
|
||||
stream.push({
|
||||
type: "thinking_start",
|
||||
contentIndex: blockIndex(),
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
currentBlock = { type: "text", text: "" };
|
||||
output.content.push(currentBlock);
|
||||
ensureStarted();
|
||||
stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
}
|
||||
if (currentBlock.type === "thinking") {
|
||||
currentBlock.thinking += part.text;
|
||||
currentBlock.thinkingSignature = retainThoughtSignature(
|
||||
currentBlock.thinkingSignature,
|
||||
part.thoughtSignature,
|
||||
);
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: part.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
currentBlock.text += part.text;
|
||||
currentBlock.textSignature = retainThoughtSignature(
|
||||
currentBlock.textSignature,
|
||||
part.thoughtSignature,
|
||||
);
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: part.text,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (part.functionCall) {
|
||||
hasContent = true;
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
currentBlock = null;
|
||||
}
|
||||
|
||||
const providedId = part.functionCall.id;
|
||||
const needsNewId =
|
||||
!providedId ||
|
||||
output.content.some((b) => b.type === "toolCall" && b.id === providedId);
|
||||
const toolCallId = needsNewId
|
||||
? `${part.functionCall.name}_${Date.now()}_${++toolCallCounter}`
|
||||
: providedId;
|
||||
|
||||
const toolCall: ToolCall = {
|
||||
type: "toolCall",
|
||||
id: toolCallId,
|
||||
name: part.functionCall.name || "",
|
||||
arguments: (part.functionCall.args as Record<string, unknown>) ?? {},
|
||||
...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }),
|
||||
};
|
||||
|
||||
output.content.push(toolCall);
|
||||
ensureStarted();
|
||||
stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: JSON.stringify(toolCall.arguments),
|
||||
partial: output,
|
||||
});
|
||||
stream.push({
|
||||
type: "toolcall_end",
|
||||
contentIndex: blockIndex(),
|
||||
toolCall,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (candidate?.finishReason) {
|
||||
output.stopReason = mapStopReasonString(candidate.finishReason);
|
||||
if (output.content.some((b) => b.type === "toolCall")) {
|
||||
output.stopReason = "toolUse";
|
||||
}
|
||||
}
|
||||
|
||||
if (responseData.usageMetadata) {
|
||||
// promptTokenCount includes cachedContentTokenCount, so subtract to get fresh input
|
||||
const promptTokens = responseData.usageMetadata.promptTokenCount || 0;
|
||||
const cacheReadTokens = responseData.usageMetadata.cachedContentTokenCount || 0;
|
||||
output.usage = {
|
||||
input: promptTokens - cacheReadTokens,
|
||||
output:
|
||||
(responseData.usageMetadata.candidatesTokenCount || 0) +
|
||||
(responseData.usageMetadata.thoughtsTokenCount || 0),
|
||||
cacheRead: cacheReadTokens,
|
||||
cacheWrite: 0,
|
||||
totalTokens: responseData.usageMetadata.totalTokenCount || 0,
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
total: 0,
|
||||
},
|
||||
};
|
||||
calculateCost(model, output.usage);
|
||||
}
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
options?.signal?.removeEventListener("abort", abortHandler);
|
||||
}
|
||||
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
return hasContent;
|
||||
};
|
||||
|
||||
let receivedContent = false;
|
||||
let currentResponse = response;
|
||||
|
||||
for (let emptyAttempt = 0; emptyAttempt <= MAX_EMPTY_STREAM_RETRIES; emptyAttempt++) {
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (emptyAttempt > 0) {
|
||||
const backoffMs = EMPTY_STREAM_BASE_DELAY_MS * 2 ** (emptyAttempt - 1);
|
||||
await sleep(backoffMs, options?.signal);
|
||||
|
||||
if (!requestUrl) {
|
||||
throw new Error("Missing request URL");
|
||||
}
|
||||
|
||||
currentResponse = await fetch(requestUrl, {
|
||||
method: "POST",
|
||||
headers: requestHeaders,
|
||||
body: requestBodyJson,
|
||||
signal: options?.signal,
|
||||
});
|
||||
await options?.onResponse?.(
|
||||
{ status: currentResponse.status, headers: headersToRecord(currentResponse.headers) },
|
||||
model,
|
||||
);
|
||||
|
||||
if (!currentResponse.ok) {
|
||||
const retryErrorText = await currentResponse.text();
|
||||
throw new Error(`Cloud Code Assist API error (${currentResponse.status}): ${retryErrorText}`);
|
||||
}
|
||||
}
|
||||
|
||||
const streamed = await streamResponse(currentResponse);
|
||||
if (streamed) {
|
||||
receivedContent = true;
|
||||
break;
|
||||
}
|
||||
|
||||
if (emptyAttempt < MAX_EMPTY_STREAM_RETRIES) {
|
||||
resetOutput();
|
||||
}
|
||||
}
|
||||
|
||||
if (!receivedContent) {
|
||||
throw new Error("Cloud Code Assist API returned an empty response");
|
||||
}
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "aborted" || output.stopReason === "error") {
|
||||
throw new Error("An unknown error occurred");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
for (const block of output.content) {
|
||||
if ("index" in block) {
|
||||
delete (block as { index?: number }).index;
|
||||
}
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
export const streamSimpleGoogleGeminiCli: StreamFunction<"google-gemini-cli", SimpleStreamOptions> = (
|
||||
model: Model<"google-gemini-cli">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const apiKey = options?.apiKey;
|
||||
if (!apiKey) {
|
||||
throw new Error("Google Cloud Code Assist requires OAuth authentication. Use /login to authenticate.");
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
if (!options?.reasoning) {
|
||||
return streamGoogleGeminiCli(model, context, {
|
||||
...base,
|
||||
thinking: { enabled: false },
|
||||
} satisfies GoogleGeminiCliOptions);
|
||||
}
|
||||
|
||||
const effort = clampReasoning(options.reasoning)!;
|
||||
if (isGemini3Model(model.id)) {
|
||||
return streamGoogleGeminiCli(model, context, {
|
||||
...base,
|
||||
thinking: {
|
||||
enabled: true,
|
||||
level: getGeminiCliThinkingLevel(effort, model.id),
|
||||
},
|
||||
} satisfies GoogleGeminiCliOptions);
|
||||
}
|
||||
|
||||
const defaultBudgets: ThinkingBudgets = {
|
||||
minimal: 1024,
|
||||
low: 2048,
|
||||
medium: 8192,
|
||||
high: 16384,
|
||||
};
|
||||
const budgets = { ...defaultBudgets, ...options.thinkingBudgets };
|
||||
|
||||
const minOutputTokens = 1024;
|
||||
let thinkingBudget = budgets[effort]!;
|
||||
const maxTokens = Math.min((base.maxTokens || 0) + thinkingBudget, model.maxTokens);
|
||||
|
||||
if (maxTokens <= thinkingBudget) {
|
||||
thinkingBudget = Math.max(0, maxTokens - minOutputTokens);
|
||||
}
|
||||
|
||||
return streamGoogleGeminiCli(model, context, {
|
||||
...base,
|
||||
maxTokens,
|
||||
thinking: {
|
||||
enabled: true,
|
||||
budgetTokens: thinkingBudget,
|
||||
},
|
||||
} satisfies GoogleGeminiCliOptions);
|
||||
};
|
||||
|
||||
export function buildRequest(
|
||||
model: Model<"google-gemini-cli">,
|
||||
context: Context,
|
||||
projectId: string,
|
||||
options: GoogleGeminiCliOptions = {},
|
||||
isAntigravity = false,
|
||||
): CloudCodeAssistRequest {
|
||||
const contents = convertMessages(model, context);
|
||||
|
||||
const generationConfig: CloudCodeAssistRequest["request"]["generationConfig"] = {};
|
||||
if (options.temperature !== undefined) {
|
||||
generationConfig.temperature = options.temperature;
|
||||
}
|
||||
if (options.maxTokens !== undefined) {
|
||||
generationConfig.maxOutputTokens = options.maxTokens;
|
||||
}
|
||||
|
||||
// Thinking config
|
||||
if (options.thinking?.enabled && model.reasoning) {
|
||||
generationConfig.thinkingConfig = {
|
||||
includeThoughts: true,
|
||||
};
|
||||
// Gemini 3 models use thinkingLevel, older models use thinkingBudget
|
||||
if (options.thinking.level !== undefined) {
|
||||
// Cast to any since our GoogleThinkingLevel mirrors Google's ThinkingLevel enum values
|
||||
generationConfig.thinkingConfig.thinkingLevel = options.thinking.level as any;
|
||||
} else if (options.thinking.budgetTokens !== undefined) {
|
||||
generationConfig.thinkingConfig.thinkingBudget = options.thinking.budgetTokens;
|
||||
}
|
||||
} else if (model.reasoning && options.thinking && !options.thinking.enabled) {
|
||||
generationConfig.thinkingConfig = getDisabledThinkingConfig(model.id);
|
||||
}
|
||||
|
||||
const request: CloudCodeAssistRequest["request"] = {
|
||||
contents,
|
||||
};
|
||||
|
||||
request.sessionId = options.sessionId;
|
||||
|
||||
// System instruction must be object with parts, not plain string
|
||||
if (context.systemPrompt) {
|
||||
request.systemInstruction = {
|
||||
parts: [{ text: sanitizeSurrogates(context.systemPrompt) }],
|
||||
};
|
||||
}
|
||||
|
||||
if (Object.keys(generationConfig).length > 0) {
|
||||
request.generationConfig = generationConfig;
|
||||
}
|
||||
|
||||
if (context.tools && context.tools.length > 0) {
|
||||
// Claude models on Cloud Code Assist need the legacy `parameters` field;
|
||||
// the API translates it into Anthropic's `input_schema`.
|
||||
const useParameters = model.id.startsWith("claude-");
|
||||
request.tools = convertTools(context.tools, useParameters);
|
||||
if (options.toolChoice) {
|
||||
request.toolConfig = {
|
||||
functionCallingConfig: {
|
||||
mode: mapToolChoice(options.toolChoice),
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
if (isAntigravity) {
|
||||
const existingParts = request.systemInstruction?.parts ?? [];
|
||||
request.systemInstruction = {
|
||||
role: "user",
|
||||
parts: [
|
||||
{ text: ANTIGRAVITY_SYSTEM_INSTRUCTION },
|
||||
{ text: `Please ignore following [ignore]${ANTIGRAVITY_SYSTEM_INSTRUCTION}[/ignore]` },
|
||||
...existingParts,
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
return {
|
||||
project: projectId,
|
||||
model: model.id,
|
||||
request,
|
||||
...(isAntigravity ? { requestType: "agent" } : {}),
|
||||
userAgent: isAntigravity ? "antigravity" : "pi-coding-agent",
|
||||
requestId: `${isAntigravity ? "agent" : "pi"}-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`,
|
||||
};
|
||||
}
|
||||
|
||||
type ClampedThinkingLevel = Exclude<ThinkingLevel, "xhigh">;
|
||||
|
||||
function getDisabledThinkingConfig(modelId: string): ThinkingConfig {
|
||||
// Google docs: Gemini 3.1 Pro cannot disable thinking, and Gemini 3 Flash / Flash-Lite
|
||||
// do not support full thinking-off either. For Gemini 3 models, use the lowest supported
|
||||
// thinkingLevel without includeThoughts so hidden thinking remains invisible to pi.
|
||||
if (isGemini3ProModel(modelId)) {
|
||||
return { thinkingLevel: "LOW" as any };
|
||||
}
|
||||
if (isGemini3FlashModel(modelId)) {
|
||||
return { thinkingLevel: "MINIMAL" as any };
|
||||
}
|
||||
|
||||
// Gemini 2.x supports disabling via thinkingBudget = 0.
|
||||
return { thinkingBudget: 0 };
|
||||
}
|
||||
|
||||
function getGeminiCliThinkingLevel(effort: ClampedThinkingLevel, modelId: string): GoogleThinkingLevel {
|
||||
if (isGemini3ProModel(modelId)) {
|
||||
switch (effort) {
|
||||
case "minimal":
|
||||
case "low":
|
||||
return "LOW";
|
||||
case "medium":
|
||||
case "high":
|
||||
return "HIGH";
|
||||
}
|
||||
}
|
||||
switch (effort) {
|
||||
case "minimal":
|
||||
return "MINIMAL";
|
||||
case "low":
|
||||
return "LOW";
|
||||
case "medium":
|
||||
return "MEDIUM";
|
||||
case "high":
|
||||
return "HIGH";
|
||||
}
|
||||
}
|
||||
@@ -1,326 +0,0 @@
|
||||
/**
|
||||
* Shared utilities for Google Generative AI and Google Cloud Code Assist providers.
|
||||
*/
|
||||
|
||||
import { type Content, FinishReason, FunctionCallingConfigMode, type Part } from "@google/genai";
|
||||
import type { Context, ImageContent, Model, StopReason, TextContent, Tool } from "../types.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
|
||||
type GoogleApiType = "google-generative-ai" | "google-gemini-cli" | "google-vertex";
|
||||
|
||||
/**
|
||||
* Determines whether a streamed Gemini `Part` should be treated as "thinking".
|
||||
*
|
||||
* Protocol note (Gemini / Vertex AI thought signatures):
|
||||
* - `thought: true` is the definitive marker for thinking content (thought summaries).
|
||||
* - `thoughtSignature` is an encrypted representation of the model's internal thought process
|
||||
* used to preserve reasoning context across multi-turn interactions.
|
||||
* - `thoughtSignature` can appear on ANY part type (text, functionCall, etc.) - it does NOT
|
||||
* indicate the part itself is thinking content.
|
||||
* - For non-functionCall responses, the signature appears on the last part for context replay.
|
||||
* - When persisting/replaying model outputs, signature-bearing parts must be preserved as-is;
|
||||
* do not merge/move signatures across parts.
|
||||
*
|
||||
* See: https://ai.google.dev/gemini-api/docs/thought-signatures
|
||||
*/
|
||||
export function isThinkingPart(part: Pick<Part, "thought" | "thoughtSignature">): boolean {
|
||||
return part.thought === true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Retain thought signatures during streaming.
|
||||
*
|
||||
* Some backends only send `thoughtSignature` on the first delta for a given part/block; later deltas may omit it.
|
||||
* This helper preserves the last non-empty signature for the current block.
|
||||
*
|
||||
* Note: this does NOT merge or move signatures across distinct response parts. It only prevents
|
||||
* a signature from being overwritten with `undefined` within the same streamed block.
|
||||
*/
|
||||
export function retainThoughtSignature(existing: string | undefined, incoming: string | undefined): string | undefined {
|
||||
if (typeof incoming === "string" && incoming.length > 0) return incoming;
|
||||
return existing;
|
||||
}
|
||||
|
||||
// Thought signatures must be base64 for Google APIs (TYPE_BYTES).
|
||||
const base64SignaturePattern = /^[A-Za-z0-9+/]+={0,2}$/;
|
||||
|
||||
// Sentinel value that tells the Gemini API to skip thought signature validation.
|
||||
// Used for unsigned function call parts (e.g. replayed from providers without thought signatures).
|
||||
// See: https://ai.google.dev/gemini-api/docs/thought-signatures
|
||||
const SKIP_THOUGHT_SIGNATURE = "skip_thought_signature_validator";
|
||||
|
||||
function isValidThoughtSignature(signature: string | undefined): boolean {
|
||||
if (!signature) return false;
|
||||
if (signature.length % 4 !== 0) return false;
|
||||
return base64SignaturePattern.test(signature);
|
||||
}
|
||||
|
||||
/**
|
||||
* Only keep signatures from the same provider/model and with valid base64.
|
||||
*/
|
||||
function resolveThoughtSignature(isSameProviderAndModel: boolean, signature: string | undefined): string | undefined {
|
||||
return isSameProviderAndModel && isValidThoughtSignature(signature) ? signature : undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Models via Google APIs that require explicit tool call IDs in function calls/responses.
|
||||
*/
|
||||
export function requiresToolCallId(modelId: string): boolean {
|
||||
return modelId.startsWith("claude-") || modelId.startsWith("gpt-oss-");
|
||||
}
|
||||
|
||||
function getGeminiMajorVersion(modelId: string): number | undefined {
|
||||
const match = modelId.toLowerCase().match(/^gemini(?:-live)?-(\d+)/);
|
||||
if (!match) return undefined;
|
||||
return Number.parseInt(match[1], 10);
|
||||
}
|
||||
|
||||
function supportsMultimodalFunctionResponse(modelId: string): boolean {
|
||||
const geminiMajorVersion = getGeminiMajorVersion(modelId);
|
||||
if (geminiMajorVersion !== undefined) {
|
||||
return geminiMajorVersion >= 3;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert internal messages to Gemini Content[] format.
|
||||
*/
|
||||
export function convertMessages<T extends GoogleApiType>(model: Model<T>, context: Context): Content[] {
|
||||
const contents: Content[] = [];
|
||||
const normalizeToolCallId = (id: string): string => {
|
||||
if (!requiresToolCallId(model.id)) return id;
|
||||
return id.replace(/[^a-zA-Z0-9_-]/g, "_").slice(0, 64);
|
||||
};
|
||||
|
||||
const transformedMessages = transformMessages(context.messages, model, normalizeToolCallId);
|
||||
|
||||
for (const msg of transformedMessages) {
|
||||
if (msg.role === "user") {
|
||||
if (typeof msg.content === "string") {
|
||||
contents.push({
|
||||
role: "user",
|
||||
parts: [{ text: sanitizeSurrogates(msg.content) }],
|
||||
});
|
||||
} else {
|
||||
const parts: Part[] = msg.content.map((item) => {
|
||||
if (item.type === "text") {
|
||||
return { text: sanitizeSurrogates(item.text) };
|
||||
} else {
|
||||
return {
|
||||
inlineData: {
|
||||
mimeType: item.mimeType,
|
||||
data: item.data,
|
||||
},
|
||||
};
|
||||
}
|
||||
});
|
||||
const filteredParts = !model.input.includes("image") ? parts.filter((p) => p.text !== undefined) : parts;
|
||||
if (filteredParts.length === 0) continue;
|
||||
contents.push({
|
||||
role: "user",
|
||||
parts: filteredParts,
|
||||
});
|
||||
}
|
||||
} else if (msg.role === "assistant") {
|
||||
const parts: Part[] = [];
|
||||
// Check if message is from same provider and model - only then keep thinking blocks
|
||||
const isSameProviderAndModel = msg.provider === model.provider && msg.model === model.id;
|
||||
|
||||
for (const block of msg.content) {
|
||||
if (block.type === "text") {
|
||||
// Skip empty text blocks - they can cause issues with some models (e.g. Claude via Antigravity)
|
||||
if (!block.text || block.text.trim() === "") continue;
|
||||
const thoughtSignature = resolveThoughtSignature(isSameProviderAndModel, block.textSignature);
|
||||
parts.push({
|
||||
text: sanitizeSurrogates(block.text),
|
||||
...(thoughtSignature && { thoughtSignature }),
|
||||
});
|
||||
} else if (block.type === "thinking") {
|
||||
// Skip empty thinking blocks
|
||||
if (!block.thinking || block.thinking.trim() === "") continue;
|
||||
// Only keep as thinking block if same provider AND same model
|
||||
// Otherwise convert to plain text (no tags to avoid model mimicking them)
|
||||
if (isSameProviderAndModel) {
|
||||
const thoughtSignature = resolveThoughtSignature(isSameProviderAndModel, block.thinkingSignature);
|
||||
parts.push({
|
||||
thought: true,
|
||||
text: sanitizeSurrogates(block.thinking),
|
||||
...(thoughtSignature && { thoughtSignature }),
|
||||
});
|
||||
} else {
|
||||
parts.push({
|
||||
text: sanitizeSurrogates(block.thinking),
|
||||
});
|
||||
}
|
||||
} else if (block.type === "toolCall") {
|
||||
const thoughtSignature = resolveThoughtSignature(isSameProviderAndModel, block.thoughtSignature);
|
||||
// Gemini 3 requires thoughtSignature on all function calls when thinking mode is enabled.
|
||||
// Use the skip_thought_signature_validator sentinel for unsigned function calls
|
||||
// (e.g. replayed from providers without thought signatures like Claude via Antigravity).
|
||||
const isGemini3 = model.id.toLowerCase().includes("gemini-3");
|
||||
const effectiveSignature = thoughtSignature || (isGemini3 ? SKIP_THOUGHT_SIGNATURE : undefined);
|
||||
const part: Part = {
|
||||
functionCall: {
|
||||
name: block.name,
|
||||
args: block.arguments ?? {},
|
||||
...(requiresToolCallId(model.id) ? { id: block.id } : {}),
|
||||
},
|
||||
...(effectiveSignature && { thoughtSignature: effectiveSignature }),
|
||||
};
|
||||
parts.push(part);
|
||||
}
|
||||
}
|
||||
|
||||
if (parts.length === 0) continue;
|
||||
contents.push({
|
||||
role: "model",
|
||||
parts,
|
||||
});
|
||||
} else if (msg.role === "toolResult") {
|
||||
// Extract text and image content
|
||||
const textContent = msg.content.filter((c): c is TextContent => c.type === "text");
|
||||
const textResult = textContent.map((c) => c.text).join("\n");
|
||||
const imageContent = model.input.includes("image")
|
||||
? msg.content.filter((c): c is ImageContent => c.type === "image")
|
||||
: [];
|
||||
|
||||
const hasText = textResult.length > 0;
|
||||
const hasImages = imageContent.length > 0;
|
||||
|
||||
// Gemini 3+ models support multimodal function responses with images nested inside
|
||||
// functionResponse.parts. Claude and other non-Gemini models behind Cloud Code Assist /
|
||||
// Antigravity also accept this shape. Gemini < 3 still needs a separate user image turn.
|
||||
const modelSupportsMultimodalFunctionResponse = supportsMultimodalFunctionResponse(model.id);
|
||||
|
||||
// Use "output" key for success, "error" key for errors as per SDK documentation
|
||||
const responseValue = hasText ? sanitizeSurrogates(textResult) : hasImages ? "(see attached image)" : "";
|
||||
|
||||
const imageParts: Part[] = imageContent.map((imageBlock) => ({
|
||||
inlineData: {
|
||||
mimeType: imageBlock.mimeType,
|
||||
data: imageBlock.data,
|
||||
},
|
||||
}));
|
||||
|
||||
const includeId = requiresToolCallId(model.id);
|
||||
const functionResponsePart: Part = {
|
||||
functionResponse: {
|
||||
name: msg.toolName,
|
||||
response: msg.isError ? { error: responseValue } : { output: responseValue },
|
||||
...(hasImages && modelSupportsMultimodalFunctionResponse && { parts: imageParts }),
|
||||
...(includeId ? { id: msg.toolCallId } : {}),
|
||||
},
|
||||
};
|
||||
|
||||
// Cloud Code Assist API requires all function responses to be in a single user turn.
|
||||
// Check if the last content is already a user turn with function responses and merge.
|
||||
const lastContent = contents[contents.length - 1];
|
||||
if (lastContent?.role === "user" && lastContent.parts?.some((p) => p.functionResponse)) {
|
||||
lastContent.parts.push(functionResponsePart);
|
||||
} else {
|
||||
contents.push({
|
||||
role: "user",
|
||||
parts: [functionResponsePart],
|
||||
});
|
||||
}
|
||||
|
||||
// For Gemini < 3, add images in a separate user message
|
||||
if (hasImages && !modelSupportsMultimodalFunctionResponse) {
|
||||
contents.push({
|
||||
role: "user",
|
||||
parts: [{ text: "Tool result image:" }, ...imageParts],
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return contents;
|
||||
}
|
||||
|
||||
/**
|
||||
* Convert tools to Gemini function declarations format.
|
||||
*
|
||||
* By default uses `parametersJsonSchema` which supports full JSON Schema (including
|
||||
* anyOf, oneOf, const, etc.). Set `useParameters` to true to use the legacy `parameters`
|
||||
* field instead (OpenAPI 3.03 Schema). This is needed for Cloud Code Assist with Claude
|
||||
* models, where the API translates `parameters` into Anthropic's `input_schema`.
|
||||
*/
|
||||
export function convertTools(
|
||||
tools: Tool[],
|
||||
useParameters = false,
|
||||
): { functionDeclarations: Record<string, unknown>[] }[] | undefined {
|
||||
if (tools.length === 0) return undefined;
|
||||
return [
|
||||
{
|
||||
functionDeclarations: tools.map((tool) => ({
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
...(useParameters ? { parameters: tool.parameters } : { parametersJsonSchema: tool.parameters }),
|
||||
})),
|
||||
},
|
||||
];
|
||||
}
|
||||
|
||||
/**
|
||||
* Map tool choice string to Gemini FunctionCallingConfigMode.
|
||||
*/
|
||||
export function mapToolChoice(choice: string): FunctionCallingConfigMode {
|
||||
switch (choice) {
|
||||
case "auto":
|
||||
return FunctionCallingConfigMode.AUTO;
|
||||
case "none":
|
||||
return FunctionCallingConfigMode.NONE;
|
||||
case "any":
|
||||
return FunctionCallingConfigMode.ANY;
|
||||
default:
|
||||
return FunctionCallingConfigMode.AUTO;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Map Gemini FinishReason to our StopReason.
|
||||
*/
|
||||
export function mapStopReason(reason: FinishReason): StopReason {
|
||||
switch (reason) {
|
||||
case FinishReason.STOP:
|
||||
return "stop";
|
||||
case FinishReason.MAX_TOKENS:
|
||||
return "length";
|
||||
case FinishReason.BLOCKLIST:
|
||||
case FinishReason.PROHIBITED_CONTENT:
|
||||
case FinishReason.SPII:
|
||||
case FinishReason.SAFETY:
|
||||
case FinishReason.IMAGE_SAFETY:
|
||||
case FinishReason.IMAGE_PROHIBITED_CONTENT:
|
||||
case FinishReason.IMAGE_RECITATION:
|
||||
case FinishReason.IMAGE_OTHER:
|
||||
case FinishReason.RECITATION:
|
||||
case FinishReason.FINISH_REASON_UNSPECIFIED:
|
||||
case FinishReason.OTHER:
|
||||
case FinishReason.LANGUAGE:
|
||||
case FinishReason.MALFORMED_FUNCTION_CALL:
|
||||
case FinishReason.UNEXPECTED_TOOL_CALL:
|
||||
case FinishReason.NO_IMAGE:
|
||||
return "error";
|
||||
default: {
|
||||
const _exhaustive: never = reason;
|
||||
throw new Error(`Unhandled stop reason: ${_exhaustive}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Map string finish reason to our StopReason (for raw API responses).
|
||||
*/
|
||||
export function mapStopReasonString(reason: string): StopReason {
|
||||
switch (reason) {
|
||||
case "STOP":
|
||||
return "stop";
|
||||
case "MAX_TOKENS":
|
||||
return "length";
|
||||
default:
|
||||
return "error";
|
||||
}
|
||||
}
|
||||
@@ -1,543 +0,0 @@
|
||||
import {
|
||||
type GenerateContentConfig,
|
||||
type GenerateContentParameters,
|
||||
GoogleGenAI,
|
||||
type ThinkingConfig,
|
||||
ThinkingLevel,
|
||||
} from "@google/genai";
|
||||
import { calculateCost } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
Context,
|
||||
Model,
|
||||
ThinkingLevel as PiThinkingLevel,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
TextContent,
|
||||
ThinkingBudgets,
|
||||
ThinkingContent,
|
||||
ToolCall,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import type { GoogleThinkingLevel } from "./google-gemini-cli.js";
|
||||
import {
|
||||
convertMessages,
|
||||
convertTools,
|
||||
isThinkingPart,
|
||||
mapStopReason,
|
||||
mapToolChoice,
|
||||
retainThoughtSignature,
|
||||
} from "./google-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
|
||||
export interface GoogleVertexOptions extends StreamOptions {
|
||||
toolChoice?: "auto" | "none" | "any";
|
||||
thinking?: {
|
||||
enabled: boolean;
|
||||
budgetTokens?: number; // -1 for dynamic, 0 to disable
|
||||
level?: GoogleThinkingLevel;
|
||||
};
|
||||
project?: string;
|
||||
location?: string;
|
||||
}
|
||||
|
||||
const API_VERSION = "v1";
|
||||
const GCP_VERTEX_CREDENTIALS_MARKER = "gcp-vertex-credentials";
|
||||
|
||||
const THINKING_LEVEL_MAP: Record<GoogleThinkingLevel, ThinkingLevel> = {
|
||||
THINKING_LEVEL_UNSPECIFIED: ThinkingLevel.THINKING_LEVEL_UNSPECIFIED,
|
||||
MINIMAL: ThinkingLevel.MINIMAL,
|
||||
LOW: ThinkingLevel.LOW,
|
||||
MEDIUM: ThinkingLevel.MEDIUM,
|
||||
HIGH: ThinkingLevel.HIGH,
|
||||
};
|
||||
|
||||
// Counter for generating unique tool call IDs
|
||||
let toolCallCounter = 0;
|
||||
|
||||
export const streamGoogleVertex: StreamFunction<"google-vertex", GoogleVertexOptions> = (
|
||||
model: Model<"google-vertex">,
|
||||
context: Context,
|
||||
options?: GoogleVertexOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: "google-vertex" as Api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
try {
|
||||
const apiKey = resolveApiKey(options);
|
||||
// Create the client using either a Vertex API key, if provided, or ADC with project and location
|
||||
const client = apiKey
|
||||
? createClientWithApiKey(model, apiKey, options?.headers)
|
||||
: createClient(model, resolveProject(options), resolveLocation(options), options?.headers);
|
||||
let params = buildParams(model, context, options);
|
||||
const nextParams = await options?.onPayload?.(params, model);
|
||||
if (nextParams !== undefined) {
|
||||
params = nextParams as GenerateContentParameters;
|
||||
}
|
||||
const googleStream = await client.models.generateContentStream(params);
|
||||
|
||||
stream.push({ type: "start", partial: output });
|
||||
let currentBlock: TextContent | ThinkingContent | null = null;
|
||||
const blocks = output.content;
|
||||
const blockIndex = () => blocks.length - 1;
|
||||
for await (const chunk of googleStream) {
|
||||
// Vertex uses the same @google/genai GenerateContentResponse type as Gemini.
|
||||
// responseId is documented there as an output-only identifier for each response.
|
||||
output.responseId ||= chunk.responseId;
|
||||
const candidate = chunk.candidates?.[0];
|
||||
if (candidate?.content?.parts) {
|
||||
for (const part of candidate.content.parts) {
|
||||
if (part.text !== undefined) {
|
||||
const isThinking = isThinkingPart(part);
|
||||
if (
|
||||
!currentBlock ||
|
||||
(isThinking && currentBlock.type !== "thinking") ||
|
||||
(!isThinking && currentBlock.type !== "text")
|
||||
) {
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blocks.length - 1,
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
if (isThinking) {
|
||||
currentBlock = { type: "thinking", thinking: "", thinkingSignature: undefined };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output });
|
||||
} else {
|
||||
currentBlock = { type: "text", text: "" };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
}
|
||||
if (currentBlock.type === "thinking") {
|
||||
currentBlock.thinking += part.text;
|
||||
currentBlock.thinkingSignature = retainThoughtSignature(
|
||||
currentBlock.thinkingSignature,
|
||||
part.thoughtSignature,
|
||||
);
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: part.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
currentBlock.text += part.text;
|
||||
currentBlock.textSignature = retainThoughtSignature(
|
||||
currentBlock.textSignature,
|
||||
part.thoughtSignature,
|
||||
);
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: part.text,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (part.functionCall) {
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
currentBlock = null;
|
||||
}
|
||||
|
||||
const providedId = part.functionCall.id;
|
||||
const needsNewId =
|
||||
!providedId || output.content.some((b) => b.type === "toolCall" && b.id === providedId);
|
||||
const toolCallId = needsNewId
|
||||
? `${part.functionCall.name}_${Date.now()}_${++toolCallCounter}`
|
||||
: providedId;
|
||||
|
||||
const toolCall: ToolCall = {
|
||||
type: "toolCall",
|
||||
id: toolCallId,
|
||||
name: part.functionCall.name || "",
|
||||
arguments: (part.functionCall.args as Record<string, any>) ?? {},
|
||||
...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }),
|
||||
};
|
||||
|
||||
output.content.push(toolCall);
|
||||
stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: JSON.stringify(toolCall.arguments),
|
||||
partial: output,
|
||||
});
|
||||
stream.push({ type: "toolcall_end", contentIndex: blockIndex(), toolCall, partial: output });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (candidate?.finishReason) {
|
||||
output.stopReason = mapStopReason(candidate.finishReason);
|
||||
if (output.content.some((b) => b.type === "toolCall")) {
|
||||
output.stopReason = "toolUse";
|
||||
}
|
||||
}
|
||||
|
||||
if (chunk.usageMetadata) {
|
||||
output.usage = {
|
||||
input:
|
||||
(chunk.usageMetadata.promptTokenCount || 0) - (chunk.usageMetadata.cachedContentTokenCount || 0),
|
||||
output:
|
||||
(chunk.usageMetadata.candidatesTokenCount || 0) + (chunk.usageMetadata.thoughtsTokenCount || 0),
|
||||
cacheRead: chunk.usageMetadata.cachedContentTokenCount || 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: chunk.usageMetadata.totalTokenCount || 0,
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
total: 0,
|
||||
},
|
||||
};
|
||||
calculateCost(model, output.usage);
|
||||
}
|
||||
}
|
||||
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "aborted" || output.stopReason === "error") {
|
||||
throw new Error("An unknown error occurred");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
// Remove internal index property used during streaming
|
||||
for (const block of output.content) {
|
||||
if ("index" in block) {
|
||||
delete (block as { index?: number }).index;
|
||||
}
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
export const streamSimpleGoogleVertex: StreamFunction<"google-vertex", SimpleStreamOptions> = (
|
||||
model: Model<"google-vertex">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const base = buildBaseOptions(model, options, undefined);
|
||||
if (!options?.reasoning) {
|
||||
return streamGoogleVertex(model, context, {
|
||||
...base,
|
||||
thinking: { enabled: false },
|
||||
} satisfies GoogleVertexOptions);
|
||||
}
|
||||
|
||||
const effort = clampReasoning(options.reasoning)!;
|
||||
const geminiModel = model as unknown as Model<"google-generative-ai">;
|
||||
|
||||
if (isGemini3ProModel(geminiModel) || isGemini3FlashModel(geminiModel)) {
|
||||
return streamGoogleVertex(model, context, {
|
||||
...base,
|
||||
thinking: {
|
||||
enabled: true,
|
||||
level: getGemini3ThinkingLevel(effort, geminiModel),
|
||||
},
|
||||
} satisfies GoogleVertexOptions);
|
||||
}
|
||||
|
||||
return streamGoogleVertex(model, context, {
|
||||
...base,
|
||||
thinking: {
|
||||
enabled: true,
|
||||
budgetTokens: getGoogleBudget(geminiModel, effort, options.thinkingBudgets),
|
||||
},
|
||||
} satisfies GoogleVertexOptions);
|
||||
};
|
||||
|
||||
function createClient(
|
||||
model: Model<"google-vertex">,
|
||||
project: string,
|
||||
location: string,
|
||||
optionsHeaders?: Record<string, string>,
|
||||
): GoogleGenAI {
|
||||
const httpOptions: { headers?: Record<string, string> } = {};
|
||||
|
||||
if (model.headers || optionsHeaders) {
|
||||
httpOptions.headers = { ...model.headers, ...optionsHeaders };
|
||||
}
|
||||
|
||||
const hasHttpOptions = Object.values(httpOptions).some(Boolean);
|
||||
|
||||
return new GoogleGenAI({
|
||||
vertexai: true,
|
||||
project,
|
||||
location,
|
||||
apiVersion: API_VERSION,
|
||||
httpOptions: hasHttpOptions ? httpOptions : undefined,
|
||||
});
|
||||
}
|
||||
|
||||
function createClientWithApiKey(
|
||||
model: Model<"google-vertex">,
|
||||
apiKey: string,
|
||||
optionsHeaders?: Record<string, string>,
|
||||
): GoogleGenAI {
|
||||
const httpOptions: { headers?: Record<string, string> } = {};
|
||||
|
||||
if (model.headers || optionsHeaders) {
|
||||
httpOptions.headers = { ...model.headers, ...optionsHeaders };
|
||||
}
|
||||
|
||||
const hasHttpOptions = Object.values(httpOptions).some(Boolean);
|
||||
|
||||
return new GoogleGenAI({
|
||||
vertexai: true,
|
||||
apiKey,
|
||||
apiVersion: API_VERSION,
|
||||
httpOptions: hasHttpOptions ? httpOptions : undefined,
|
||||
});
|
||||
}
|
||||
|
||||
function resolveApiKey(options?: GoogleVertexOptions): string | undefined {
|
||||
const apiKey = options?.apiKey?.trim() || process.env.GOOGLE_CLOUD_API_KEY?.trim();
|
||||
if (!apiKey || apiKey === GCP_VERTEX_CREDENTIALS_MARKER || isPlaceholderApiKey(apiKey)) {
|
||||
return undefined;
|
||||
}
|
||||
return apiKey;
|
||||
}
|
||||
|
||||
function isPlaceholderApiKey(apiKey: string): boolean {
|
||||
return /^<[^>]+>$/.test(apiKey);
|
||||
}
|
||||
|
||||
function resolveProject(options?: GoogleVertexOptions): string {
|
||||
const project = options?.project || process.env.GOOGLE_CLOUD_PROJECT || process.env.GCLOUD_PROJECT;
|
||||
if (!project) {
|
||||
throw new Error(
|
||||
"Vertex AI requires a project ID. Set GOOGLE_CLOUD_PROJECT/GCLOUD_PROJECT or pass project in options.",
|
||||
);
|
||||
}
|
||||
return project;
|
||||
}
|
||||
|
||||
function resolveLocation(options?: GoogleVertexOptions): string {
|
||||
const location = options?.location || process.env.GOOGLE_CLOUD_LOCATION;
|
||||
if (!location) {
|
||||
throw new Error("Vertex AI requires a location. Set GOOGLE_CLOUD_LOCATION or pass location in options.");
|
||||
}
|
||||
return location;
|
||||
}
|
||||
|
||||
function buildParams(
|
||||
model: Model<"google-vertex">,
|
||||
context: Context,
|
||||
options: GoogleVertexOptions = {},
|
||||
): GenerateContentParameters {
|
||||
const contents = convertMessages(model, context);
|
||||
|
||||
const generationConfig: GenerateContentConfig = {};
|
||||
if (options.temperature !== undefined) {
|
||||
generationConfig.temperature = options.temperature;
|
||||
}
|
||||
if (options.maxTokens !== undefined) {
|
||||
generationConfig.maxOutputTokens = options.maxTokens;
|
||||
}
|
||||
|
||||
const config: GenerateContentConfig = {
|
||||
...(Object.keys(generationConfig).length > 0 && generationConfig),
|
||||
...(context.systemPrompt && { systemInstruction: sanitizeSurrogates(context.systemPrompt) }),
|
||||
...(context.tools && context.tools.length > 0 && { tools: convertTools(context.tools) }),
|
||||
};
|
||||
|
||||
if (context.tools && context.tools.length > 0 && options.toolChoice) {
|
||||
config.toolConfig = {
|
||||
functionCallingConfig: {
|
||||
mode: mapToolChoice(options.toolChoice),
|
||||
},
|
||||
};
|
||||
} else {
|
||||
config.toolConfig = undefined;
|
||||
}
|
||||
|
||||
if (options.thinking?.enabled && model.reasoning) {
|
||||
const thinkingConfig: ThinkingConfig = { includeThoughts: true };
|
||||
if (options.thinking.level !== undefined) {
|
||||
thinkingConfig.thinkingLevel = THINKING_LEVEL_MAP[options.thinking.level];
|
||||
} else if (options.thinking.budgetTokens !== undefined) {
|
||||
thinkingConfig.thinkingBudget = options.thinking.budgetTokens;
|
||||
}
|
||||
config.thinkingConfig = thinkingConfig;
|
||||
} else if (model.reasoning && options.thinking && !options.thinking.enabled) {
|
||||
config.thinkingConfig = getDisabledThinkingConfig(model);
|
||||
}
|
||||
|
||||
if (options.signal) {
|
||||
if (options.signal.aborted) {
|
||||
throw new Error("Request aborted");
|
||||
}
|
||||
config.abortSignal = options.signal;
|
||||
}
|
||||
|
||||
const params: GenerateContentParameters = {
|
||||
model: model.id,
|
||||
contents,
|
||||
config,
|
||||
};
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
type ClampedThinkingLevel = Exclude<PiThinkingLevel, "xhigh">;
|
||||
|
||||
function isGemini3ProModel(model: Model<"google-generative-ai">): boolean {
|
||||
return /gemini-3(?:\.\d+)?-pro/.test(model.id.toLowerCase());
|
||||
}
|
||||
|
||||
function isGemini3FlashModel(model: Model<"google-generative-ai">): boolean {
|
||||
return /gemini-3(?:\.\d+)?-flash/.test(model.id.toLowerCase());
|
||||
}
|
||||
|
||||
function getDisabledThinkingConfig(model: Model<"google-vertex">): ThinkingConfig {
|
||||
// Google docs: Gemini 3.1 Pro cannot disable thinking, and Gemini 3 Flash / Flash-Lite
|
||||
// do not support full thinking-off either. For Gemini 3 models, use the lowest supported
|
||||
// thinkingLevel without includeThoughts so hidden thinking remains invisible to pi.
|
||||
const geminiModel = model as unknown as Model<"google-generative-ai">;
|
||||
if (isGemini3ProModel(geminiModel)) {
|
||||
return { thinkingLevel: ThinkingLevel.LOW };
|
||||
}
|
||||
if (isGemini3FlashModel(geminiModel)) {
|
||||
return { thinkingLevel: ThinkingLevel.MINIMAL };
|
||||
}
|
||||
|
||||
// Gemini 2.x supports disabling via thinkingBudget = 0.
|
||||
return { thinkingBudget: 0 };
|
||||
}
|
||||
|
||||
function getGemini3ThinkingLevel(
|
||||
effort: ClampedThinkingLevel,
|
||||
model: Model<"google-generative-ai">,
|
||||
): GoogleThinkingLevel {
|
||||
if (isGemini3ProModel(model)) {
|
||||
switch (effort) {
|
||||
case "minimal":
|
||||
case "low":
|
||||
return "LOW";
|
||||
case "medium":
|
||||
case "high":
|
||||
return "HIGH";
|
||||
}
|
||||
}
|
||||
switch (effort) {
|
||||
case "minimal":
|
||||
return "MINIMAL";
|
||||
case "low":
|
||||
return "LOW";
|
||||
case "medium":
|
||||
return "MEDIUM";
|
||||
case "high":
|
||||
return "HIGH";
|
||||
}
|
||||
}
|
||||
|
||||
function getGoogleBudget(
|
||||
model: Model<"google-generative-ai">,
|
||||
effort: ClampedThinkingLevel,
|
||||
customBudgets?: ThinkingBudgets,
|
||||
): number {
|
||||
if (customBudgets?.[effort] !== undefined) {
|
||||
return customBudgets[effort]!;
|
||||
}
|
||||
|
||||
if (model.id.includes("2.5-pro")) {
|
||||
const budgets: Record<ClampedThinkingLevel, number> = {
|
||||
minimal: 128,
|
||||
low: 2048,
|
||||
medium: 8192,
|
||||
high: 32768,
|
||||
};
|
||||
return budgets[effort];
|
||||
}
|
||||
|
||||
if (model.id.includes("2.5-flash")) {
|
||||
const budgets: Record<ClampedThinkingLevel, number> = {
|
||||
minimal: 128,
|
||||
low: 2048,
|
||||
medium: 8192,
|
||||
high: 24576,
|
||||
};
|
||||
return budgets[effort];
|
||||
}
|
||||
|
||||
return -1;
|
||||
}
|
||||
@@ -1,500 +0,0 @@
|
||||
import {
|
||||
type GenerateContentConfig,
|
||||
type GenerateContentParameters,
|
||||
GoogleGenAI,
|
||||
type ThinkingConfig,
|
||||
} from "@google/genai";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { calculateCost } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
Context,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
TextContent,
|
||||
ThinkingBudgets,
|
||||
ThinkingContent,
|
||||
ThinkingLevel,
|
||||
ToolCall,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import type { GoogleThinkingLevel } from "./google-gemini-cli.js";
|
||||
import {
|
||||
convertMessages,
|
||||
convertTools,
|
||||
isThinkingPart,
|
||||
mapStopReason,
|
||||
mapToolChoice,
|
||||
retainThoughtSignature,
|
||||
} from "./google-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
|
||||
export interface GoogleOptions extends StreamOptions {
|
||||
toolChoice?: "auto" | "none" | "any";
|
||||
thinking?: {
|
||||
enabled: boolean;
|
||||
budgetTokens?: number; // -1 for dynamic, 0 to disable
|
||||
level?: GoogleThinkingLevel;
|
||||
};
|
||||
}
|
||||
|
||||
// Counter for generating unique tool call IDs
|
||||
let toolCallCounter = 0;
|
||||
|
||||
export const streamGoogle: StreamFunction<"google-generative-ai", GoogleOptions> = (
|
||||
model: Model<"google-generative-ai">,
|
||||
context: Context,
|
||||
options?: GoogleOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: "google-generative-ai" as Api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
try {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider) || "";
|
||||
const client = createClient(model, apiKey, options?.headers);
|
||||
let params = buildParams(model, context, options);
|
||||
const nextParams = await options?.onPayload?.(params, model);
|
||||
if (nextParams !== undefined) {
|
||||
params = nextParams as GenerateContentParameters;
|
||||
}
|
||||
const googleStream = await client.models.generateContentStream(params);
|
||||
|
||||
stream.push({ type: "start", partial: output });
|
||||
let currentBlock: TextContent | ThinkingContent | null = null;
|
||||
const blocks = output.content;
|
||||
const blockIndex = () => blocks.length - 1;
|
||||
for await (const chunk of googleStream) {
|
||||
// @google/genai documents GenerateContentResponse.responseId as an output-only field
|
||||
// used to identify each response. Keep the first non-empty one from the stream.
|
||||
output.responseId ||= chunk.responseId;
|
||||
const candidate = chunk.candidates?.[0];
|
||||
if (candidate?.content?.parts) {
|
||||
for (const part of candidate.content.parts) {
|
||||
if (part.text !== undefined) {
|
||||
const isThinking = isThinkingPart(part);
|
||||
if (
|
||||
!currentBlock ||
|
||||
(isThinking && currentBlock.type !== "thinking") ||
|
||||
(!isThinking && currentBlock.type !== "text")
|
||||
) {
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blocks.length - 1,
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
if (isThinking) {
|
||||
currentBlock = { type: "thinking", thinking: "", thinkingSignature: undefined };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output });
|
||||
} else {
|
||||
currentBlock = { type: "text", text: "" };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
}
|
||||
if (currentBlock.type === "thinking") {
|
||||
currentBlock.thinking += part.text;
|
||||
currentBlock.thinkingSignature = retainThoughtSignature(
|
||||
currentBlock.thinkingSignature,
|
||||
part.thoughtSignature,
|
||||
);
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: part.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
currentBlock.text += part.text;
|
||||
currentBlock.textSignature = retainThoughtSignature(
|
||||
currentBlock.textSignature,
|
||||
part.thoughtSignature,
|
||||
);
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: part.text,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (part.functionCall) {
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
currentBlock = null;
|
||||
}
|
||||
|
||||
// Generate unique ID if not provided or if it's a duplicate
|
||||
const providedId = part.functionCall.id;
|
||||
const needsNewId =
|
||||
!providedId || output.content.some((b) => b.type === "toolCall" && b.id === providedId);
|
||||
const toolCallId = needsNewId
|
||||
? `${part.functionCall.name}_${Date.now()}_${++toolCallCounter}`
|
||||
: providedId;
|
||||
|
||||
const toolCall: ToolCall = {
|
||||
type: "toolCall",
|
||||
id: toolCallId,
|
||||
name: part.functionCall.name || "",
|
||||
arguments: (part.functionCall.args as Record<string, any>) ?? {},
|
||||
...(part.thoughtSignature && { thoughtSignature: part.thoughtSignature }),
|
||||
};
|
||||
|
||||
output.content.push(toolCall);
|
||||
stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: JSON.stringify(toolCall.arguments),
|
||||
partial: output,
|
||||
});
|
||||
stream.push({ type: "toolcall_end", contentIndex: blockIndex(), toolCall, partial: output });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (candidate?.finishReason) {
|
||||
output.stopReason = mapStopReason(candidate.finishReason);
|
||||
if (output.content.some((b) => b.type === "toolCall")) {
|
||||
output.stopReason = "toolUse";
|
||||
}
|
||||
}
|
||||
|
||||
if (chunk.usageMetadata) {
|
||||
output.usage = {
|
||||
input:
|
||||
(chunk.usageMetadata.promptTokenCount || 0) - (chunk.usageMetadata.cachedContentTokenCount || 0),
|
||||
output:
|
||||
(chunk.usageMetadata.candidatesTokenCount || 0) + (chunk.usageMetadata.thoughtsTokenCount || 0),
|
||||
cacheRead: chunk.usageMetadata.cachedContentTokenCount || 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: chunk.usageMetadata.totalTokenCount || 0,
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
total: 0,
|
||||
},
|
||||
};
|
||||
calculateCost(model, output.usage);
|
||||
}
|
||||
}
|
||||
|
||||
if (currentBlock) {
|
||||
if (currentBlock.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
} else {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "aborted" || output.stopReason === "error") {
|
||||
throw new Error("An unknown error occurred");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
// Remove internal index property used during streaming
|
||||
for (const block of output.content) {
|
||||
if ("index" in block) {
|
||||
delete (block as { index?: number }).index;
|
||||
}
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
export const streamSimpleGoogle: StreamFunction<"google-generative-ai", SimpleStreamOptions> = (
|
||||
model: Model<"google-generative-ai">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
|
||||
if (!apiKey) {
|
||||
throw new Error(`No API key for provider: ${model.provider}`);
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
if (!options?.reasoning) {
|
||||
return streamGoogle(model, context, { ...base, thinking: { enabled: false } } satisfies GoogleOptions);
|
||||
}
|
||||
|
||||
const effort = clampReasoning(options.reasoning)!;
|
||||
const googleModel = model as Model<"google-generative-ai">;
|
||||
|
||||
if (isGemini3ProModel(googleModel) || isGemini3FlashModel(googleModel) || isGemma4Model(googleModel)) {
|
||||
return streamGoogle(model, context, {
|
||||
...base,
|
||||
thinking: {
|
||||
enabled: true,
|
||||
level: getThinkingLevel(effort, googleModel),
|
||||
},
|
||||
} satisfies GoogleOptions);
|
||||
}
|
||||
|
||||
return streamGoogle(model, context, {
|
||||
...base,
|
||||
thinking: {
|
||||
enabled: true,
|
||||
budgetTokens: getGoogleBudget(googleModel, effort, options.thinkingBudgets),
|
||||
},
|
||||
} satisfies GoogleOptions);
|
||||
};
|
||||
|
||||
function createClient(
|
||||
model: Model<"google-generative-ai">,
|
||||
apiKey?: string,
|
||||
optionsHeaders?: Record<string, string>,
|
||||
): GoogleGenAI {
|
||||
const httpOptions: { baseUrl?: string; apiVersion?: string; headers?: Record<string, string> } = {};
|
||||
if (model.baseUrl) {
|
||||
httpOptions.baseUrl = model.baseUrl;
|
||||
httpOptions.apiVersion = ""; // baseUrl already includes version path, don't append
|
||||
}
|
||||
if (model.headers || optionsHeaders) {
|
||||
httpOptions.headers = { ...model.headers, ...optionsHeaders };
|
||||
}
|
||||
|
||||
return new GoogleGenAI({
|
||||
apiKey,
|
||||
httpOptions: Object.keys(httpOptions).length > 0 ? httpOptions : undefined,
|
||||
});
|
||||
}
|
||||
|
||||
function buildParams(
|
||||
model: Model<"google-generative-ai">,
|
||||
context: Context,
|
||||
options: GoogleOptions = {},
|
||||
): GenerateContentParameters {
|
||||
const contents = convertMessages(model, context);
|
||||
|
||||
const generationConfig: GenerateContentConfig = {};
|
||||
if (options.temperature !== undefined) {
|
||||
generationConfig.temperature = options.temperature;
|
||||
}
|
||||
if (options.maxTokens !== undefined) {
|
||||
generationConfig.maxOutputTokens = options.maxTokens;
|
||||
}
|
||||
|
||||
const config: GenerateContentConfig = {
|
||||
...(Object.keys(generationConfig).length > 0 && generationConfig),
|
||||
...(context.systemPrompt && { systemInstruction: sanitizeSurrogates(context.systemPrompt) }),
|
||||
...(context.tools && context.tools.length > 0 && { tools: convertTools(context.tools) }),
|
||||
};
|
||||
|
||||
if (context.tools && context.tools.length > 0 && options.toolChoice) {
|
||||
config.toolConfig = {
|
||||
functionCallingConfig: {
|
||||
mode: mapToolChoice(options.toolChoice),
|
||||
},
|
||||
};
|
||||
} else {
|
||||
config.toolConfig = undefined;
|
||||
}
|
||||
|
||||
if (options.thinking?.enabled && model.reasoning) {
|
||||
const thinkingConfig: ThinkingConfig = { includeThoughts: true };
|
||||
if (options.thinking.level !== undefined) {
|
||||
// Cast to any since our GoogleThinkingLevel mirrors Google's ThinkingLevel enum values
|
||||
thinkingConfig.thinkingLevel = options.thinking.level as any;
|
||||
} else if (options.thinking.budgetTokens !== undefined) {
|
||||
thinkingConfig.thinkingBudget = options.thinking.budgetTokens;
|
||||
}
|
||||
config.thinkingConfig = thinkingConfig;
|
||||
} else if (model.reasoning && options.thinking && !options.thinking.enabled) {
|
||||
config.thinkingConfig = getDisabledThinkingConfig(model);
|
||||
}
|
||||
|
||||
if (options.signal) {
|
||||
if (options.signal.aborted) {
|
||||
throw new Error("Request aborted");
|
||||
}
|
||||
config.abortSignal = options.signal;
|
||||
}
|
||||
|
||||
const params: GenerateContentParameters = {
|
||||
model: model.id,
|
||||
contents,
|
||||
config,
|
||||
};
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
type ClampedThinkingLevel = Exclude<ThinkingLevel, "xhigh">;
|
||||
|
||||
function isGemma4Model(model: Model<"google-generative-ai">): boolean {
|
||||
return /gemma-?4/.test(model.id.toLowerCase());
|
||||
}
|
||||
|
||||
function isGemini3ProModel(model: Model<"google-generative-ai">): boolean {
|
||||
return /gemini-3(?:\.\d+)?-pro/.test(model.id.toLowerCase());
|
||||
}
|
||||
|
||||
function isGemini3FlashModel(model: Model<"google-generative-ai">): boolean {
|
||||
return /gemini-3(?:\.\d+)?-flash/.test(model.id.toLowerCase());
|
||||
}
|
||||
|
||||
function getDisabledThinkingConfig(model: Model<"google-generative-ai">): ThinkingConfig {
|
||||
// Google docs: Gemini 3.1 Pro cannot disable thinking, and Gemini 3 Flash / Flash-Lite
|
||||
// do not support full thinking-off either. For Gemini 3 models, use the lowest supported
|
||||
// thinkingLevel without includeThoughts so hidden thinking remains invisible to pi.
|
||||
if (isGemini3ProModel(model)) {
|
||||
return { thinkingLevel: "LOW" as any };
|
||||
}
|
||||
if (isGemini3FlashModel(model)) {
|
||||
return { thinkingLevel: "MINIMAL" as any };
|
||||
}
|
||||
if (isGemma4Model(model)) {
|
||||
return { thinkingLevel: "MINIMAL" as any };
|
||||
}
|
||||
|
||||
// Gemini 2.x supports disabling via thinkingBudget = 0.
|
||||
return { thinkingBudget: 0 };
|
||||
}
|
||||
|
||||
function getThinkingLevel(effort: ClampedThinkingLevel, model: Model<"google-generative-ai">): GoogleThinkingLevel {
|
||||
if (isGemini3ProModel(model)) {
|
||||
switch (effort) {
|
||||
case "minimal":
|
||||
case "low":
|
||||
return "LOW";
|
||||
case "medium":
|
||||
case "high":
|
||||
return "HIGH";
|
||||
}
|
||||
}
|
||||
if (isGemma4Model(model)) {
|
||||
switch (effort) {
|
||||
case "minimal":
|
||||
case "low":
|
||||
return "MINIMAL";
|
||||
case "medium":
|
||||
case "high":
|
||||
return "HIGH";
|
||||
}
|
||||
}
|
||||
switch (effort) {
|
||||
case "minimal":
|
||||
return "MINIMAL";
|
||||
case "low":
|
||||
return "LOW";
|
||||
case "medium":
|
||||
return "MEDIUM";
|
||||
case "high":
|
||||
return "HIGH";
|
||||
}
|
||||
}
|
||||
|
||||
function getGoogleBudget(
|
||||
model: Model<"google-generative-ai">,
|
||||
effort: ClampedThinkingLevel,
|
||||
customBudgets?: ThinkingBudgets,
|
||||
): number {
|
||||
if (customBudgets?.[effort] !== undefined) {
|
||||
return customBudgets[effort]!;
|
||||
}
|
||||
|
||||
if (model.id.includes("2.5-pro")) {
|
||||
const budgets: Record<ClampedThinkingLevel, number> = {
|
||||
minimal: 128,
|
||||
low: 2048,
|
||||
medium: 8192,
|
||||
high: 32768,
|
||||
};
|
||||
return budgets[effort];
|
||||
}
|
||||
|
||||
if (model.id.includes("2.5-flash-lite")) {
|
||||
const budgets: Record<ClampedThinkingLevel, number> = {
|
||||
minimal: 512,
|
||||
low: 2048,
|
||||
medium: 8192,
|
||||
high: 24576,
|
||||
};
|
||||
return budgets[effort];
|
||||
}
|
||||
|
||||
if (model.id.includes("2.5-flash")) {
|
||||
const budgets: Record<ClampedThinkingLevel, number> = {
|
||||
minimal: 128,
|
||||
low: 2048,
|
||||
medium: 8192,
|
||||
high: 24576,
|
||||
};
|
||||
return budgets[effort];
|
||||
}
|
||||
|
||||
return -1;
|
||||
}
|
||||
@@ -1,629 +0,0 @@
|
||||
import { Mistral } from "@mistralai/mistralai";
|
||||
import type {
|
||||
ChatCompletionStreamRequest,
|
||||
ChatCompletionStreamRequestMessage,
|
||||
CompletionEvent,
|
||||
ContentChunk,
|
||||
FunctionTool,
|
||||
} from "@mistralai/mistralai/models/components";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { calculateCost } from "../models.js";
|
||||
import type {
|
||||
AssistantMessage,
|
||||
Context,
|
||||
Message,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StopReason,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
TextContent,
|
||||
ThinkingContent,
|
||||
Tool,
|
||||
ToolCall,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { shortHash } from "../utils/hash.js";
|
||||
import { parseStreamingJson } from "../utils/json-parse.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
|
||||
const MISTRAL_TOOL_CALL_ID_LENGTH = 9;
|
||||
const MAX_MISTRAL_ERROR_BODY_CHARS = 4000;
|
||||
|
||||
/**
|
||||
* Provider-specific options for the Mistral API.
|
||||
*/
|
||||
type MistralReasoningEffort = "none" | "high";
|
||||
|
||||
export interface MistralOptions extends StreamOptions {
|
||||
toolChoice?: "auto" | "none" | "any" | "required" | { type: "function"; function: { name: string } };
|
||||
promptMode?: "reasoning";
|
||||
reasoningEffort?: MistralReasoningEffort;
|
||||
}
|
||||
|
||||
/**
|
||||
* Stream responses from Mistral using `chat.stream`.
|
||||
*/
|
||||
export const streamMistral: StreamFunction<"mistral-conversations", MistralOptions> = (
|
||||
model: Model<"mistral-conversations">,
|
||||
context: Context,
|
||||
options?: MistralOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
const output = createOutput(model);
|
||||
|
||||
try {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
|
||||
if (!apiKey) {
|
||||
throw new Error(`No API key for provider: ${model.provider}`);
|
||||
}
|
||||
|
||||
// Intentionally per-request: avoids shared SDK mutable state across concurrent consumers.
|
||||
const mistral = new Mistral({
|
||||
apiKey,
|
||||
serverURL: model.baseUrl,
|
||||
});
|
||||
|
||||
const normalizeMistralToolCallId = createMistralToolCallIdNormalizer();
|
||||
const transformedMessages = transformMessages(context.messages, model, (id) => normalizeMistralToolCallId(id));
|
||||
|
||||
let payload = buildChatPayload(model, context, transformedMessages, options);
|
||||
const nextPayload = await options?.onPayload?.(payload, model);
|
||||
if (nextPayload !== undefined) {
|
||||
payload = nextPayload as ChatCompletionStreamRequest;
|
||||
}
|
||||
const mistralStream = await mistral.chat.stream(payload, buildRequestOptions(model, options));
|
||||
stream.push({ type: "start", partial: output });
|
||||
await consumeChatStream(model, output, stream, mistralStream);
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "aborted" || output.stopReason === "error") {
|
||||
throw new Error("An unknown error occurred");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
for (const block of output.content) {
|
||||
// partialArgs is only a streaming scratch buffer; never persist it.
|
||||
delete (block as { partialArgs?: string }).partialArgs;
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = formatMistralError(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
/**
|
||||
* Maps provider-agnostic `SimpleStreamOptions` to Mistral options.
|
||||
*/
|
||||
export const streamSimpleMistral: StreamFunction<"mistral-conversations", SimpleStreamOptions> = (
|
||||
model: Model<"mistral-conversations">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
|
||||
if (!apiKey) {
|
||||
throw new Error(`No API key for provider: ${model.provider}`);
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoning = clampReasoning(options?.reasoning);
|
||||
const shouldUseReasoning = model.reasoning && reasoning !== undefined;
|
||||
|
||||
return streamMistral(model, context, {
|
||||
...base,
|
||||
promptMode: shouldUseReasoning && usesPromptModeReasoning(model) ? "reasoning" : undefined,
|
||||
reasoningEffort: shouldUseReasoning && usesReasoningEffort(model) ? mapReasoningEffort(reasoning) : undefined,
|
||||
} satisfies MistralOptions);
|
||||
};
|
||||
|
||||
function createOutput(model: Model<"mistral-conversations">): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createMistralToolCallIdNormalizer(): (id: string) => string {
|
||||
const idMap = new Map<string, string>();
|
||||
const reverseMap = new Map<string, string>();
|
||||
|
||||
return (id: string): string => {
|
||||
const existing = idMap.get(id);
|
||||
if (existing) return existing;
|
||||
|
||||
let attempt = 0;
|
||||
while (true) {
|
||||
const candidate = deriveMistralToolCallId(id, attempt);
|
||||
const owner = reverseMap.get(candidate);
|
||||
if (!owner || owner === id) {
|
||||
idMap.set(id, candidate);
|
||||
reverseMap.set(candidate, id);
|
||||
return candidate;
|
||||
}
|
||||
attempt++;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
function deriveMistralToolCallId(id: string, attempt: number): string {
|
||||
const normalized = id.replace(/[^a-zA-Z0-9]/g, "");
|
||||
if (attempt === 0 && normalized.length === MISTRAL_TOOL_CALL_ID_LENGTH) return normalized;
|
||||
const seedBase = normalized || id;
|
||||
const seed = attempt === 0 ? seedBase : `${seedBase}:${attempt}`;
|
||||
return shortHash(seed)
|
||||
.replace(/[^a-zA-Z0-9]/g, "")
|
||||
.slice(0, MISTRAL_TOOL_CALL_ID_LENGTH);
|
||||
}
|
||||
|
||||
function formatMistralError(error: unknown): string {
|
||||
if (error instanceof Error) {
|
||||
const sdkError = error as Error & { statusCode?: unknown; body?: unknown };
|
||||
const statusCode = typeof sdkError.statusCode === "number" ? sdkError.statusCode : undefined;
|
||||
const bodyText = typeof sdkError.body === "string" ? sdkError.body.trim() : undefined;
|
||||
if (statusCode !== undefined && bodyText) {
|
||||
return `Mistral API error (${statusCode}): ${truncateErrorText(bodyText, MAX_MISTRAL_ERROR_BODY_CHARS)}`;
|
||||
}
|
||||
if (statusCode !== undefined) return `Mistral API error (${statusCode}): ${error.message}`;
|
||||
return error.message;
|
||||
}
|
||||
return safeJsonStringify(error);
|
||||
}
|
||||
|
||||
function truncateErrorText(text: string, maxChars: number): string {
|
||||
if (text.length <= maxChars) return text;
|
||||
return `${text.slice(0, maxChars)}... [truncated ${text.length - maxChars} chars]`;
|
||||
}
|
||||
|
||||
function safeJsonStringify(value: unknown): string {
|
||||
try {
|
||||
const serialized = JSON.stringify(value);
|
||||
return serialized === undefined ? String(value) : serialized;
|
||||
} catch {
|
||||
return String(value);
|
||||
}
|
||||
}
|
||||
|
||||
function buildRequestOptions(model: Model<"mistral-conversations">, options?: MistralOptions) {
|
||||
const requestOptions: {
|
||||
signal?: AbortSignal;
|
||||
retries: { strategy: "none" };
|
||||
headers?: Record<string, string>;
|
||||
} = {
|
||||
retries: { strategy: "none" },
|
||||
};
|
||||
if (options?.signal) requestOptions.signal = options.signal;
|
||||
|
||||
const headers: Record<string, string> = {};
|
||||
if (model.headers) Object.assign(headers, model.headers);
|
||||
if (options?.headers) Object.assign(headers, options.headers);
|
||||
|
||||
// Mistral infrastructure uses `x-affinity` for KV-cache reuse (prefix caching).
|
||||
// Respect explicit caller-provided header values.
|
||||
if (options?.sessionId && !headers["x-affinity"]) {
|
||||
headers["x-affinity"] = options.sessionId;
|
||||
}
|
||||
|
||||
if (Object.keys(headers).length > 0) {
|
||||
requestOptions.headers = headers;
|
||||
}
|
||||
|
||||
return requestOptions;
|
||||
}
|
||||
|
||||
function buildChatPayload(
|
||||
model: Model<"mistral-conversations">,
|
||||
context: Context,
|
||||
messages: Message[],
|
||||
options?: MistralOptions,
|
||||
): ChatCompletionStreamRequest {
|
||||
const payload: ChatCompletionStreamRequest = {
|
||||
model: model.id,
|
||||
stream: true,
|
||||
messages: toChatMessages(messages, model.input.includes("image")),
|
||||
};
|
||||
|
||||
if (context.tools?.length) payload.tools = toFunctionTools(context.tools);
|
||||
if (options?.temperature !== undefined) payload.temperature = options.temperature;
|
||||
if (options?.maxTokens !== undefined) payload.maxTokens = options.maxTokens;
|
||||
if (options?.toolChoice) payload.toolChoice = mapToolChoice(options.toolChoice);
|
||||
if (options?.promptMode) payload.promptMode = options.promptMode;
|
||||
if (options?.reasoningEffort) payload.reasoningEffort = options.reasoningEffort;
|
||||
|
||||
if (context.systemPrompt) {
|
||||
payload.messages.unshift({
|
||||
role: "system",
|
||||
content: sanitizeSurrogates(context.systemPrompt),
|
||||
});
|
||||
}
|
||||
|
||||
return payload;
|
||||
}
|
||||
|
||||
async function consumeChatStream(
|
||||
model: Model<"mistral-conversations">,
|
||||
output: AssistantMessage,
|
||||
stream: AssistantMessageEventStream,
|
||||
mistralStream: AsyncIterable<CompletionEvent>,
|
||||
): Promise<void> {
|
||||
let currentBlock: TextContent | ThinkingContent | null = null;
|
||||
const blocks = output.content;
|
||||
const blockIndex = () => blocks.length - 1;
|
||||
const toolBlocksByKey = new Map<string, number>();
|
||||
|
||||
const finishCurrentBlock = (block?: typeof currentBlock) => {
|
||||
if (!block) return;
|
||||
if (block.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: block.text,
|
||||
partial: output,
|
||||
});
|
||||
return;
|
||||
}
|
||||
if (block.type === "thinking") {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: block.thinking,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
for await (const event of mistralStream) {
|
||||
const chunk = event.data;
|
||||
// Mistral's streamed CompletionChunk carries an id field. Keep the first non-empty one,
|
||||
// mirroring how OpenAI-style streaming exposes a stable response identifier per stream.
|
||||
output.responseId ||= chunk.id;
|
||||
|
||||
if (chunk.usage) {
|
||||
output.usage.input = chunk.usage.promptTokens || 0;
|
||||
output.usage.output = chunk.usage.completionTokens || 0;
|
||||
output.usage.cacheRead = 0;
|
||||
output.usage.cacheWrite = 0;
|
||||
output.usage.totalTokens = chunk.usage.totalTokens || output.usage.input + output.usage.output;
|
||||
calculateCost(model, output.usage);
|
||||
}
|
||||
|
||||
const choice = chunk.choices[0];
|
||||
if (!choice) continue;
|
||||
|
||||
if (choice.finishReason) {
|
||||
output.stopReason = mapChatStopReason(choice.finishReason);
|
||||
}
|
||||
|
||||
const delta = choice.delta;
|
||||
if (delta.content !== null && delta.content !== undefined) {
|
||||
const contentItems = typeof delta.content === "string" ? [delta.content] : delta.content;
|
||||
for (const item of contentItems) {
|
||||
if (typeof item === "string") {
|
||||
const textDelta = sanitizeSurrogates(item);
|
||||
if (!currentBlock || currentBlock.type !== "text") {
|
||||
finishCurrentBlock(currentBlock);
|
||||
currentBlock = { type: "text", text: "" };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
currentBlock.text += textDelta;
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: textDelta,
|
||||
partial: output,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
if (item.type === "thinking") {
|
||||
const deltaText = item.thinking
|
||||
.map((part) => ("text" in part ? part.text : ""))
|
||||
.filter((text) => text.length > 0)
|
||||
.join("");
|
||||
const thinkingDelta = sanitizeSurrogates(deltaText);
|
||||
if (!thinkingDelta) continue;
|
||||
if (!currentBlock || currentBlock.type !== "thinking") {
|
||||
finishCurrentBlock(currentBlock);
|
||||
currentBlock = { type: "thinking", thinking: "" };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
currentBlock.thinking += thinkingDelta;
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: thinkingDelta,
|
||||
partial: output,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
if (item.type === "text") {
|
||||
const textDelta = sanitizeSurrogates(item.text);
|
||||
if (!currentBlock || currentBlock.type !== "text") {
|
||||
finishCurrentBlock(currentBlock);
|
||||
currentBlock = { type: "text", text: "" };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
currentBlock.text += textDelta;
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: textDelta,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const toolCalls = delta.toolCalls || [];
|
||||
for (const toolCall of toolCalls) {
|
||||
if (currentBlock) {
|
||||
finishCurrentBlock(currentBlock);
|
||||
currentBlock = null;
|
||||
}
|
||||
const callId =
|
||||
toolCall.id && toolCall.id !== "null"
|
||||
? toolCall.id
|
||||
: deriveMistralToolCallId(`toolcall:${toolCall.index ?? 0}`, 0);
|
||||
const key = `${callId}:${toolCall.index || 0}`;
|
||||
const existingIndex = toolBlocksByKey.get(key);
|
||||
let block: (ToolCall & { partialArgs?: string }) | undefined;
|
||||
|
||||
if (existingIndex !== undefined) {
|
||||
const existing = output.content[existingIndex];
|
||||
if (existing?.type === "toolCall") {
|
||||
block = existing as ToolCall & { partialArgs?: string };
|
||||
}
|
||||
}
|
||||
|
||||
if (!block) {
|
||||
block = {
|
||||
type: "toolCall",
|
||||
id: callId,
|
||||
name: toolCall.function.name,
|
||||
arguments: {},
|
||||
partialArgs: "",
|
||||
};
|
||||
output.content.push(block);
|
||||
toolBlocksByKey.set(key, output.content.length - 1);
|
||||
stream.push({ type: "toolcall_start", contentIndex: output.content.length - 1, partial: output });
|
||||
}
|
||||
|
||||
const argsDelta =
|
||||
typeof toolCall.function.arguments === "string"
|
||||
? toolCall.function.arguments
|
||||
: JSON.stringify(toolCall.function.arguments || {});
|
||||
block.partialArgs = (block.partialArgs || "") + argsDelta;
|
||||
block.arguments = parseStreamingJson<Record<string, unknown>>(block.partialArgs);
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: toolBlocksByKey.get(key)!,
|
||||
delta: argsDelta,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
finishCurrentBlock(currentBlock);
|
||||
for (const index of toolBlocksByKey.values()) {
|
||||
const block = output.content[index];
|
||||
if (block.type !== "toolCall") continue;
|
||||
const toolBlock = block as ToolCall & { partialArgs?: string };
|
||||
toolBlock.arguments = parseStreamingJson<Record<string, unknown>>(toolBlock.partialArgs);
|
||||
// Finalize in-place and strip the scratch buffer so replay only
|
||||
// carries parsed arguments.
|
||||
delete toolBlock.partialArgs;
|
||||
stream.push({
|
||||
type: "toolcall_end",
|
||||
contentIndex: index,
|
||||
toolCall: toolBlock,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
function toFunctionTools(tools: Tool[]): Array<FunctionTool & { type: "function" }> {
|
||||
return tools.map((tool) => ({
|
||||
type: "function",
|
||||
function: {
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
parameters: stripSymbolKeys(tool.parameters) as Record<string, unknown>,
|
||||
strict: false,
|
||||
},
|
||||
}));
|
||||
}
|
||||
|
||||
function stripSymbolKeys(value: unknown): unknown {
|
||||
if (Array.isArray(value)) {
|
||||
return value.map((item) => stripSymbolKeys(item));
|
||||
}
|
||||
|
||||
if (value && typeof value === "object") {
|
||||
const result: Record<string, unknown> = {};
|
||||
for (const [key, entry] of Object.entries(value)) {
|
||||
result[key] = stripSymbolKeys(entry);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
return value;
|
||||
}
|
||||
|
||||
function toChatMessages(messages: Message[], supportsImages: boolean): ChatCompletionStreamRequestMessage[] {
|
||||
const result: ChatCompletionStreamRequestMessage[] = [];
|
||||
|
||||
for (const msg of messages) {
|
||||
if (msg.role === "user") {
|
||||
if (typeof msg.content === "string") {
|
||||
result.push({ role: "user", content: sanitizeSurrogates(msg.content) });
|
||||
continue;
|
||||
}
|
||||
const hadImages = msg.content.some((item) => item.type === "image");
|
||||
const content: ContentChunk[] = msg.content
|
||||
.filter((item) => item.type === "text" || supportsImages)
|
||||
.map((item) => {
|
||||
if (item.type === "text") return { type: "text", text: sanitizeSurrogates(item.text) };
|
||||
return { type: "image_url", imageUrl: `data:${item.mimeType};base64,${item.data}` };
|
||||
});
|
||||
if (content.length > 0) {
|
||||
result.push({ role: "user", content });
|
||||
continue;
|
||||
}
|
||||
if (hadImages && !supportsImages) {
|
||||
result.push({ role: "user", content: "(image omitted: model does not support images)" });
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (msg.role === "assistant") {
|
||||
const contentParts: ContentChunk[] = [];
|
||||
const toolCalls: Array<{ id: string; type: "function"; function: { name: string; arguments: string } }> = [];
|
||||
|
||||
for (const block of msg.content) {
|
||||
if (block.type === "text") {
|
||||
if (block.text.trim().length > 0) {
|
||||
contentParts.push({ type: "text", text: sanitizeSurrogates(block.text) });
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (block.type === "thinking") {
|
||||
if (block.thinking.trim().length > 0) {
|
||||
contentParts.push({
|
||||
type: "thinking",
|
||||
thinking: [{ type: "text", text: sanitizeSurrogates(block.thinking) }],
|
||||
});
|
||||
}
|
||||
continue;
|
||||
}
|
||||
toolCalls.push({
|
||||
id: block.id,
|
||||
type: "function",
|
||||
function: { name: block.name, arguments: JSON.stringify(block.arguments || {}) },
|
||||
});
|
||||
}
|
||||
|
||||
const assistantMessage: ChatCompletionStreamRequestMessage = { role: "assistant" };
|
||||
if (contentParts.length > 0) assistantMessage.content = contentParts;
|
||||
if (toolCalls.length > 0) assistantMessage.toolCalls = toolCalls;
|
||||
if (contentParts.length > 0 || toolCalls.length > 0) result.push(assistantMessage);
|
||||
continue;
|
||||
}
|
||||
|
||||
const toolContent: ContentChunk[] = [];
|
||||
const textResult = msg.content
|
||||
.filter((part) => part.type === "text")
|
||||
.map((part) => (part.type === "text" ? sanitizeSurrogates(part.text) : ""))
|
||||
.join("\n");
|
||||
const hasImages = msg.content.some((part) => part.type === "image");
|
||||
const toolText = buildToolResultText(textResult, hasImages, supportsImages, msg.isError);
|
||||
toolContent.push({ type: "text", text: toolText });
|
||||
for (const part of msg.content) {
|
||||
if (!supportsImages) continue;
|
||||
if (part.type !== "image") continue;
|
||||
toolContent.push({
|
||||
type: "image_url",
|
||||
imageUrl: `data:${part.mimeType};base64,${part.data}`,
|
||||
});
|
||||
}
|
||||
result.push({
|
||||
role: "tool",
|
||||
toolCallId: msg.toolCallId,
|
||||
name: msg.toolName,
|
||||
content: toolContent,
|
||||
});
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
function buildToolResultText(text: string, hasImages: boolean, supportsImages: boolean, isError: boolean): string {
|
||||
const trimmed = text.trim();
|
||||
const errorPrefix = isError ? "[tool error] " : "";
|
||||
|
||||
if (trimmed.length > 0) {
|
||||
const imageSuffix = hasImages && !supportsImages ? "\n[tool image omitted: model does not support images]" : "";
|
||||
return `${errorPrefix}${trimmed}${imageSuffix}`;
|
||||
}
|
||||
|
||||
if (hasImages) {
|
||||
if (supportsImages) {
|
||||
return isError ? "[tool error] (see attached image)" : "(see attached image)";
|
||||
}
|
||||
return isError
|
||||
? "[tool error] (image omitted: model does not support images)"
|
||||
: "(image omitted: model does not support images)";
|
||||
}
|
||||
|
||||
return isError ? "[tool error] (no tool output)" : "(no tool output)";
|
||||
}
|
||||
|
||||
function usesReasoningEffort(model: Model<"mistral-conversations">): boolean {
|
||||
return model.id === "mistral-small-2603" || model.id === "mistral-small-latest";
|
||||
}
|
||||
|
||||
function usesPromptModeReasoning(model: Model<"mistral-conversations">): boolean {
|
||||
return model.reasoning && !usesReasoningEffort(model);
|
||||
}
|
||||
|
||||
function mapReasoningEffort(_level: Exclude<SimpleStreamOptions["reasoning"], undefined>): MistralReasoningEffort {
|
||||
return "high";
|
||||
}
|
||||
|
||||
function mapToolChoice(
|
||||
choice: MistralOptions["toolChoice"],
|
||||
): "auto" | "none" | "any" | "required" | { type: "function"; function: { name: string } } | undefined {
|
||||
if (!choice) return undefined;
|
||||
if (choice === "auto" || choice === "none" || choice === "any" || choice === "required") {
|
||||
return choice as any;
|
||||
}
|
||||
return {
|
||||
type: "function",
|
||||
function: { name: choice.function.name },
|
||||
};
|
||||
}
|
||||
|
||||
function mapChatStopReason(reason: string | null): StopReason {
|
||||
if (reason === null) return "stop";
|
||||
switch (reason) {
|
||||
case "stop":
|
||||
return "stop";
|
||||
case "length":
|
||||
case "model_length":
|
||||
return "length";
|
||||
case "tool_calls":
|
||||
return "toolUse";
|
||||
case "error":
|
||||
return "error";
|
||||
default:
|
||||
return "stop";
|
||||
}
|
||||
}
|
||||
@@ -1,984 +0,0 @@
|
||||
import type * as NodeOs from "node:os";
|
||||
import type {
|
||||
Tool as OpenAITool,
|
||||
ResponseCreateParamsStreaming,
|
||||
ResponseInput,
|
||||
ResponseStreamEvent,
|
||||
} from "openai/resources/responses/responses.js";
|
||||
|
||||
// NEVER convert to top-level runtime imports - breaks browser/Vite builds (web-ui)
|
||||
let _os: typeof NodeOs | null = null;
|
||||
|
||||
type DynamicImport = (specifier: string) => Promise<unknown>;
|
||||
|
||||
const dynamicImport: DynamicImport = (specifier) => import(specifier);
|
||||
const NODE_OS_SPECIFIER = "node:" + "os";
|
||||
|
||||
if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) {
|
||||
dynamicImport(NODE_OS_SPECIFIER).then((m) => {
|
||||
_os = m as typeof NodeOs;
|
||||
});
|
||||
}
|
||||
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { supportsXhigh } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
Context,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
Usage,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
|
||||
// ============================================================================
|
||||
// Configuration
|
||||
// ============================================================================
|
||||
|
||||
const DEFAULT_CODEX_BASE_URL = "https://chatgpt.com/backend-api";
|
||||
const JWT_CLAIM_PATH = "https://api.openai.com/auth" as const;
|
||||
const MAX_RETRIES = 3;
|
||||
const BASE_DELAY_MS = 1000;
|
||||
const CODEX_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode"]);
|
||||
|
||||
const CODEX_RESPONSE_STATUSES = new Set<CodexResponseStatus>([
|
||||
"completed",
|
||||
"incomplete",
|
||||
"failed",
|
||||
"cancelled",
|
||||
"queued",
|
||||
"in_progress",
|
||||
]);
|
||||
|
||||
// ============================================================================
|
||||
// Types
|
||||
// ============================================================================
|
||||
|
||||
export interface OpenAICodexResponsesOptions extends StreamOptions {
|
||||
reasoningEffort?: "none" | "minimal" | "low" | "medium" | "high" | "xhigh";
|
||||
reasoningSummary?: "auto" | "concise" | "detailed" | "off" | "on" | null;
|
||||
serviceTier?: ResponseCreateParamsStreaming["service_tier"];
|
||||
textVerbosity?: "low" | "medium" | "high";
|
||||
}
|
||||
|
||||
type CodexResponseStatus = "completed" | "incomplete" | "failed" | "cancelled" | "queued" | "in_progress";
|
||||
|
||||
interface RequestBody {
|
||||
model: string;
|
||||
store?: boolean;
|
||||
stream?: boolean;
|
||||
instructions?: string;
|
||||
input?: ResponseInput;
|
||||
tools?: OpenAITool[];
|
||||
tool_choice?: "auto";
|
||||
parallel_tool_calls?: boolean;
|
||||
temperature?: number;
|
||||
reasoning?: { effort?: string; summary?: string };
|
||||
service_tier?: ResponseCreateParamsStreaming["service_tier"];
|
||||
text?: { verbosity?: string };
|
||||
include?: string[];
|
||||
prompt_cache_key?: string;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Retry Helpers
|
||||
// ============================================================================
|
||||
|
||||
function isRetryableError(status: number, errorText: string): boolean {
|
||||
if (status === 429 || status === 500 || status === 502 || status === 503 || status === 504) {
|
||||
return true;
|
||||
}
|
||||
return /rate.?limit|overloaded|service.?unavailable|upstream.?connect|connection.?refused/i.test(errorText);
|
||||
}
|
||||
|
||||
function sleep(ms: number, signal?: AbortSignal): Promise<void> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (signal?.aborted) {
|
||||
reject(new Error("Request was aborted"));
|
||||
return;
|
||||
}
|
||||
const timeout = setTimeout(resolve, ms);
|
||||
signal?.addEventListener("abort", () => {
|
||||
clearTimeout(timeout);
|
||||
reject(new Error("Request was aborted"));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Main Stream Function
|
||||
// ============================================================================
|
||||
|
||||
export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses", OpenAICodexResponsesOptions> = (
|
||||
model: Model<"openai-codex-responses">,
|
||||
context: Context,
|
||||
options?: OpenAICodexResponsesOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: "openai-codex-responses" as Api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
try {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider) || "";
|
||||
if (!apiKey) {
|
||||
throw new Error(`No API key for provider: ${model.provider}`);
|
||||
}
|
||||
|
||||
const accountId = extractAccountId(apiKey);
|
||||
let body = buildRequestBody(model, context, options);
|
||||
const nextBody = await options?.onPayload?.(body, model);
|
||||
if (nextBody !== undefined) {
|
||||
body = nextBody as RequestBody;
|
||||
}
|
||||
const websocketRequestId = options?.sessionId || createCodexRequestId();
|
||||
const sseHeaders = buildSSEHeaders(model.headers, options?.headers, accountId, apiKey, options?.sessionId);
|
||||
const websocketHeaders = buildWebSocketHeaders(
|
||||
model.headers,
|
||||
options?.headers,
|
||||
accountId,
|
||||
apiKey,
|
||||
websocketRequestId,
|
||||
);
|
||||
const bodyJson = JSON.stringify(body);
|
||||
const transport = options?.transport || "sse";
|
||||
|
||||
if (transport !== "sse") {
|
||||
let websocketStarted = false;
|
||||
try {
|
||||
await processWebSocketStream(
|
||||
resolveCodexWebSocketUrl(model.baseUrl),
|
||||
body,
|
||||
websocketHeaders,
|
||||
output,
|
||||
stream,
|
||||
model,
|
||||
() => {
|
||||
websocketStarted = true;
|
||||
},
|
||||
options,
|
||||
);
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
stream.push({
|
||||
type: "done",
|
||||
reason: output.stopReason as "stop" | "length" | "toolUse",
|
||||
message: output,
|
||||
});
|
||||
stream.end();
|
||||
return;
|
||||
} catch (error) {
|
||||
if (transport === "websocket" || websocketStarted) {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Fetch with retry logic for rate limits and transient errors
|
||||
let response: Response | undefined;
|
||||
let lastError: Error | undefined;
|
||||
|
||||
for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
try {
|
||||
response = await fetch(resolveCodexUrl(model.baseUrl), {
|
||||
method: "POST",
|
||||
headers: sseHeaders,
|
||||
body: bodyJson,
|
||||
signal: options?.signal,
|
||||
});
|
||||
await options?.onResponse?.(
|
||||
{ status: response.status, headers: headersToRecord(response.headers) },
|
||||
model,
|
||||
);
|
||||
|
||||
if (response.ok) {
|
||||
break;
|
||||
}
|
||||
|
||||
const errorText = await response.text();
|
||||
if (attempt < MAX_RETRIES && isRetryableError(response.status, errorText)) {
|
||||
const delayMs = BASE_DELAY_MS * 2 ** attempt;
|
||||
await sleep(delayMs, options?.signal);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Parse error for friendly message on final attempt or non-retryable error
|
||||
const fakeResponse = new Response(errorText, {
|
||||
status: response.status,
|
||||
statusText: response.statusText,
|
||||
});
|
||||
const info = await parseErrorResponse(fakeResponse);
|
||||
throw new Error(info.friendlyMessage || info.message);
|
||||
} catch (error) {
|
||||
if (error instanceof Error) {
|
||||
if (error.name === "AbortError" || error.message === "Request was aborted") {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
}
|
||||
lastError = error instanceof Error ? error : new Error(String(error));
|
||||
// Network errors are retryable
|
||||
if (attempt < MAX_RETRIES && !lastError.message.includes("usage limit")) {
|
||||
const delayMs = BASE_DELAY_MS * 2 ** attempt;
|
||||
await sleep(delayMs, options?.signal);
|
||||
continue;
|
||||
}
|
||||
throw lastError;
|
||||
}
|
||||
}
|
||||
|
||||
if (!response?.ok) {
|
||||
throw lastError ?? new Error("Failed after retries");
|
||||
}
|
||||
|
||||
if (!response.body) {
|
||||
throw new Error("No response body");
|
||||
}
|
||||
|
||||
stream.push({ type: "start", partial: output });
|
||||
await processStream(response, output, stream, model, options);
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason as "stop" | "length" | "toolUse", message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
for (const block of output.content) {
|
||||
// partialJson is only a streaming scratch buffer; never persist it.
|
||||
delete (block as { partialJson?: string }).partialJson;
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = error instanceof Error ? error.message : String(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
export const streamSimpleOpenAICodexResponses: StreamFunction<"openai-codex-responses", SimpleStreamOptions> = (
|
||||
model: Model<"openai-codex-responses">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
|
||||
if (!apiKey) {
|
||||
throw new Error(`No API key for provider: ${model.provider}`);
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoningEffort = supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning);
|
||||
|
||||
return streamOpenAICodexResponses(model, context, {
|
||||
...base,
|
||||
reasoningEffort,
|
||||
} satisfies OpenAICodexResponsesOptions);
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// Request Building
|
||||
// ============================================================================
|
||||
|
||||
function buildRequestBody(
|
||||
model: Model<"openai-codex-responses">,
|
||||
context: Context,
|
||||
options?: OpenAICodexResponsesOptions,
|
||||
): RequestBody {
|
||||
const messages = convertResponsesMessages(model, context, CODEX_TOOL_CALL_PROVIDERS, {
|
||||
includeSystemPrompt: false,
|
||||
});
|
||||
|
||||
const body: RequestBody = {
|
||||
model: model.id,
|
||||
store: false,
|
||||
stream: true,
|
||||
instructions: context.systemPrompt,
|
||||
input: messages,
|
||||
text: { verbosity: options?.textVerbosity || "medium" },
|
||||
include: ["reasoning.encrypted_content"],
|
||||
prompt_cache_key: options?.sessionId,
|
||||
tool_choice: "auto",
|
||||
parallel_tool_calls: true,
|
||||
};
|
||||
|
||||
if (options?.temperature !== undefined) {
|
||||
body.temperature = options.temperature;
|
||||
}
|
||||
|
||||
if (options?.serviceTier !== undefined) {
|
||||
body.service_tier = options.serviceTier;
|
||||
}
|
||||
|
||||
if (context.tools) {
|
||||
body.tools = convertResponsesTools(context.tools, { strict: null });
|
||||
}
|
||||
|
||||
if (options?.reasoningEffort !== undefined) {
|
||||
body.reasoning = {
|
||||
effort: clampReasoningEffort(model.id, options.reasoningEffort),
|
||||
summary: options.reasoningSummary ?? "auto",
|
||||
};
|
||||
}
|
||||
|
||||
return body;
|
||||
}
|
||||
|
||||
function clampReasoningEffort(modelId: string, effort: string): string {
|
||||
const id = modelId.includes("/") ? modelId.split("/").pop()! : modelId;
|
||||
if ((id.startsWith("gpt-5.2") || id.startsWith("gpt-5.3") || id.startsWith("gpt-5.4")) && effort === "minimal")
|
||||
return "low";
|
||||
if (id === "gpt-5.1" && effort === "xhigh") return "high";
|
||||
if (id === "gpt-5.1-codex-mini") return effort === "high" || effort === "xhigh" ? "high" : "medium";
|
||||
return effort;
|
||||
}
|
||||
|
||||
function getServiceTierCostMultiplier(serviceTier: ResponseCreateParamsStreaming["service_tier"] | undefined): number {
|
||||
switch (serviceTier) {
|
||||
case "flex":
|
||||
return 0.5;
|
||||
case "priority":
|
||||
return 2;
|
||||
default:
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
function applyServiceTierPricing(usage: Usage, serviceTier: ResponseCreateParamsStreaming["service_tier"] | undefined) {
|
||||
const multiplier = getServiceTierCostMultiplier(serviceTier);
|
||||
if (multiplier === 1) return;
|
||||
|
||||
usage.cost.input *= multiplier;
|
||||
usage.cost.output *= multiplier;
|
||||
usage.cost.cacheRead *= multiplier;
|
||||
usage.cost.cacheWrite *= multiplier;
|
||||
usage.cost.total = usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite;
|
||||
}
|
||||
|
||||
function resolveCodexServiceTier(
|
||||
responseServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
|
||||
requestServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
|
||||
): ResponseCreateParamsStreaming["service_tier"] | undefined {
|
||||
if (responseServiceTier === "default" && (requestServiceTier === "flex" || requestServiceTier === "priority")) {
|
||||
return requestServiceTier;
|
||||
}
|
||||
return responseServiceTier ?? requestServiceTier;
|
||||
}
|
||||
|
||||
function resolveCodexUrl(baseUrl?: string): string {
|
||||
const raw = baseUrl && baseUrl.trim().length > 0 ? baseUrl : DEFAULT_CODEX_BASE_URL;
|
||||
const normalized = raw.replace(/\/+$/, "");
|
||||
if (normalized.endsWith("/codex/responses")) return normalized;
|
||||
if (normalized.endsWith("/codex")) return `${normalized}/responses`;
|
||||
return `${normalized}/codex/responses`;
|
||||
}
|
||||
|
||||
function resolveCodexWebSocketUrl(baseUrl?: string): string {
|
||||
const url = new URL(resolveCodexUrl(baseUrl));
|
||||
if (url.protocol === "https:") url.protocol = "wss:";
|
||||
if (url.protocol === "http:") url.protocol = "ws:";
|
||||
return url.toString();
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Response Processing
|
||||
// ============================================================================
|
||||
|
||||
async function processStream(
|
||||
response: Response,
|
||||
output: AssistantMessage,
|
||||
stream: AssistantMessageEventStream,
|
||||
model: Model<"openai-codex-responses">,
|
||||
options?: OpenAICodexResponsesOptions,
|
||||
): Promise<void> {
|
||||
await processResponsesStream(mapCodexEvents(parseSSE(response)), output, stream, model, {
|
||||
serviceTier: options?.serviceTier,
|
||||
resolveServiceTier: resolveCodexServiceTier,
|
||||
applyServiceTierPricing,
|
||||
});
|
||||
}
|
||||
|
||||
async function* mapCodexEvents(events: AsyncIterable<Record<string, unknown>>): AsyncGenerator<ResponseStreamEvent> {
|
||||
for await (const event of events) {
|
||||
const type = typeof event.type === "string" ? event.type : undefined;
|
||||
if (!type) continue;
|
||||
|
||||
if (type === "error") {
|
||||
const code = (event as { code?: string }).code || "";
|
||||
const message = (event as { message?: string }).message || "";
|
||||
throw new Error(`Codex error: ${message || code || JSON.stringify(event)}`);
|
||||
}
|
||||
|
||||
if (type === "response.failed") {
|
||||
const msg = (event as { response?: { error?: { message?: string } } }).response?.error?.message;
|
||||
throw new Error(msg || "Codex response failed");
|
||||
}
|
||||
|
||||
if (type === "response.done" || type === "response.completed" || type === "response.incomplete") {
|
||||
const response = (event as { response?: { status?: unknown } }).response;
|
||||
const normalizedResponse = response
|
||||
? { ...response, status: normalizeCodexStatus(response.status) }
|
||||
: response;
|
||||
yield { ...event, type: "response.completed", response: normalizedResponse } as ResponseStreamEvent;
|
||||
return;
|
||||
}
|
||||
|
||||
yield event as unknown as ResponseStreamEvent;
|
||||
}
|
||||
}
|
||||
|
||||
function normalizeCodexStatus(status: unknown): CodexResponseStatus | undefined {
|
||||
if (typeof status !== "string") return undefined;
|
||||
return CODEX_RESPONSE_STATUSES.has(status as CodexResponseStatus) ? (status as CodexResponseStatus) : undefined;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// SSE Parsing
|
||||
// ============================================================================
|
||||
|
||||
async function* parseSSE(response: Response): AsyncGenerator<Record<string, unknown>> {
|
||||
if (!response.body) return;
|
||||
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
|
||||
let idx = buffer.indexOf("\n\n");
|
||||
while (idx !== -1) {
|
||||
const chunk = buffer.slice(0, idx);
|
||||
buffer = buffer.slice(idx + 2);
|
||||
|
||||
const dataLines = chunk
|
||||
.split("\n")
|
||||
.filter((l) => l.startsWith("data:"))
|
||||
.map((l) => l.slice(5).trim());
|
||||
if (dataLines.length > 0) {
|
||||
const data = dataLines.join("\n").trim();
|
||||
if (data && data !== "[DONE]") {
|
||||
try {
|
||||
yield JSON.parse(data);
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
idx = buffer.indexOf("\n\n");
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
try {
|
||||
await reader.cancel();
|
||||
} catch {}
|
||||
try {
|
||||
reader.releaseLock();
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// WebSocket Parsing
|
||||
// ============================================================================
|
||||
|
||||
const OPENAI_BETA_RESPONSES_WEBSOCKETS = "responses_websockets=2026-02-06";
|
||||
const SESSION_WEBSOCKET_CACHE_TTL_MS = 5 * 60 * 1000;
|
||||
|
||||
type WebSocketEventType = "open" | "message" | "error" | "close";
|
||||
type WebSocketListener = (event: unknown) => void;
|
||||
|
||||
interface WebSocketLike {
|
||||
close(code?: number, reason?: string): void;
|
||||
send(data: string): void;
|
||||
addEventListener(type: WebSocketEventType, listener: WebSocketListener): void;
|
||||
removeEventListener(type: WebSocketEventType, listener: WebSocketListener): void;
|
||||
}
|
||||
|
||||
interface CachedWebSocketConnection {
|
||||
socket: WebSocketLike;
|
||||
busy: boolean;
|
||||
idleTimer?: ReturnType<typeof setTimeout>;
|
||||
}
|
||||
|
||||
const websocketSessionCache = new Map<string, CachedWebSocketConnection>();
|
||||
|
||||
type WebSocketConstructor = new (
|
||||
url: string,
|
||||
protocols?: string | string[] | { headers?: Record<string, string> },
|
||||
) => WebSocketLike;
|
||||
|
||||
function getWebSocketConstructor(): WebSocketConstructor | null {
|
||||
const ctor = (globalThis as { WebSocket?: unknown }).WebSocket;
|
||||
if (typeof ctor !== "function") return null;
|
||||
return ctor as unknown as WebSocketConstructor;
|
||||
}
|
||||
|
||||
function getWebSocketReadyState(socket: WebSocketLike): number | undefined {
|
||||
const readyState = (socket as { readyState?: unknown }).readyState;
|
||||
return typeof readyState === "number" ? readyState : undefined;
|
||||
}
|
||||
|
||||
function isWebSocketReusable(socket: WebSocketLike): boolean {
|
||||
const readyState = getWebSocketReadyState(socket);
|
||||
// If readyState is unavailable, assume the runtime keeps it open/reusable.
|
||||
return readyState === undefined || readyState === 1;
|
||||
}
|
||||
|
||||
function closeWebSocketSilently(socket: WebSocketLike, code = 1000, reason = "done"): void {
|
||||
try {
|
||||
socket.close(code, reason);
|
||||
} catch {}
|
||||
}
|
||||
|
||||
function scheduleSessionWebSocketExpiry(sessionId: string, entry: CachedWebSocketConnection): void {
|
||||
if (entry.idleTimer) {
|
||||
clearTimeout(entry.idleTimer);
|
||||
}
|
||||
entry.idleTimer = setTimeout(() => {
|
||||
if (entry.busy) return;
|
||||
closeWebSocketSilently(entry.socket, 1000, "idle_timeout");
|
||||
websocketSessionCache.delete(sessionId);
|
||||
}, SESSION_WEBSOCKET_CACHE_TTL_MS);
|
||||
}
|
||||
|
||||
async function connectWebSocket(url: string, headers: Headers, signal?: AbortSignal): Promise<WebSocketLike> {
|
||||
const WebSocketCtor = getWebSocketConstructor();
|
||||
if (!WebSocketCtor) {
|
||||
throw new Error("WebSocket transport is not available in this runtime");
|
||||
}
|
||||
|
||||
const wsHeaders = headersToRecord(headers);
|
||||
delete wsHeaders["OpenAI-Beta"];
|
||||
|
||||
return new Promise<WebSocketLike>((resolve, reject) => {
|
||||
let settled = false;
|
||||
let socket: WebSocketLike;
|
||||
|
||||
try {
|
||||
socket = new WebSocketCtor(url, { headers: wsHeaders });
|
||||
} catch (error) {
|
||||
reject(error instanceof Error ? error : new Error(String(error)));
|
||||
return;
|
||||
}
|
||||
|
||||
const onOpen: WebSocketListener = () => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
resolve(socket);
|
||||
};
|
||||
const onError: WebSocketListener = (event) => {
|
||||
const error = extractWebSocketError(event);
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
reject(error);
|
||||
};
|
||||
const onClose: WebSocketListener = (event) => {
|
||||
const error = extractWebSocketCloseError(event);
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
reject(error);
|
||||
};
|
||||
const onAbort = () => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
cleanup();
|
||||
socket.close(1000, "aborted");
|
||||
reject(new Error("Request was aborted"));
|
||||
};
|
||||
|
||||
const cleanup = () => {
|
||||
socket.removeEventListener("open", onOpen);
|
||||
socket.removeEventListener("error", onError);
|
||||
socket.removeEventListener("close", onClose);
|
||||
signal?.removeEventListener("abort", onAbort);
|
||||
};
|
||||
|
||||
socket.addEventListener("open", onOpen);
|
||||
socket.addEventListener("error", onError);
|
||||
socket.addEventListener("close", onClose);
|
||||
signal?.addEventListener("abort", onAbort);
|
||||
});
|
||||
}
|
||||
|
||||
async function acquireWebSocket(
|
||||
url: string,
|
||||
headers: Headers,
|
||||
sessionId: string | undefined,
|
||||
signal?: AbortSignal,
|
||||
): Promise<{ socket: WebSocketLike; release: (options?: { keep?: boolean }) => void }> {
|
||||
if (!sessionId) {
|
||||
const socket = await connectWebSocket(url, headers, signal);
|
||||
return {
|
||||
socket,
|
||||
release: ({ keep } = {}) => {
|
||||
if (keep === false) {
|
||||
closeWebSocketSilently(socket);
|
||||
return;
|
||||
}
|
||||
closeWebSocketSilently(socket);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
const cached = websocketSessionCache.get(sessionId);
|
||||
if (cached) {
|
||||
if (cached.idleTimer) {
|
||||
clearTimeout(cached.idleTimer);
|
||||
cached.idleTimer = undefined;
|
||||
}
|
||||
if (!cached.busy && isWebSocketReusable(cached.socket)) {
|
||||
cached.busy = true;
|
||||
return {
|
||||
socket: cached.socket,
|
||||
release: ({ keep } = {}) => {
|
||||
if (!keep || !isWebSocketReusable(cached.socket)) {
|
||||
closeWebSocketSilently(cached.socket);
|
||||
websocketSessionCache.delete(sessionId);
|
||||
return;
|
||||
}
|
||||
cached.busy = false;
|
||||
scheduleSessionWebSocketExpiry(sessionId, cached);
|
||||
},
|
||||
};
|
||||
}
|
||||
if (cached.busy) {
|
||||
const socket = await connectWebSocket(url, headers, signal);
|
||||
return {
|
||||
socket,
|
||||
release: () => {
|
||||
closeWebSocketSilently(socket);
|
||||
},
|
||||
};
|
||||
}
|
||||
if (!isWebSocketReusable(cached.socket)) {
|
||||
closeWebSocketSilently(cached.socket);
|
||||
websocketSessionCache.delete(sessionId);
|
||||
}
|
||||
}
|
||||
|
||||
const socket = await connectWebSocket(url, headers, signal);
|
||||
const entry: CachedWebSocketConnection = { socket, busy: true };
|
||||
websocketSessionCache.set(sessionId, entry);
|
||||
return {
|
||||
socket,
|
||||
release: ({ keep } = {}) => {
|
||||
if (!keep || !isWebSocketReusable(entry.socket)) {
|
||||
closeWebSocketSilently(entry.socket);
|
||||
if (entry.idleTimer) clearTimeout(entry.idleTimer);
|
||||
if (websocketSessionCache.get(sessionId) === entry) {
|
||||
websocketSessionCache.delete(sessionId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
entry.busy = false;
|
||||
scheduleSessionWebSocketExpiry(sessionId, entry);
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
function extractWebSocketError(event: unknown): Error {
|
||||
if (event && typeof event === "object" && "message" in event) {
|
||||
const message = (event as { message?: unknown }).message;
|
||||
if (typeof message === "string" && message.length > 0) {
|
||||
return new Error(message);
|
||||
}
|
||||
}
|
||||
return new Error("WebSocket error");
|
||||
}
|
||||
|
||||
function extractWebSocketCloseError(event: unknown): Error {
|
||||
if (event && typeof event === "object") {
|
||||
const code = "code" in event ? (event as { code?: unknown }).code : undefined;
|
||||
const reason = "reason" in event ? (event as { reason?: unknown }).reason : undefined;
|
||||
const codeText = typeof code === "number" ? ` ${code}` : "";
|
||||
const reasonText = typeof reason === "string" && reason.length > 0 ? ` ${reason}` : "";
|
||||
return new Error(`WebSocket closed${codeText}${reasonText}`.trim());
|
||||
}
|
||||
return new Error("WebSocket closed");
|
||||
}
|
||||
|
||||
async function decodeWebSocketData(data: unknown): Promise<string | null> {
|
||||
if (typeof data === "string") return data;
|
||||
if (data instanceof ArrayBuffer) {
|
||||
return new TextDecoder().decode(new Uint8Array(data));
|
||||
}
|
||||
if (ArrayBuffer.isView(data)) {
|
||||
const view = data as ArrayBufferView;
|
||||
return new TextDecoder().decode(new Uint8Array(view.buffer, view.byteOffset, view.byteLength));
|
||||
}
|
||||
if (data && typeof data === "object" && "arrayBuffer" in data) {
|
||||
const blobLike = data as { arrayBuffer: () => Promise<ArrayBuffer> };
|
||||
const arrayBuffer = await blobLike.arrayBuffer();
|
||||
return new TextDecoder().decode(new Uint8Array(arrayBuffer));
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): AsyncGenerator<Record<string, unknown>> {
|
||||
const queue: Record<string, unknown>[] = [];
|
||||
let pending: (() => void) | null = null;
|
||||
let done = false;
|
||||
let failed: Error | null = null;
|
||||
let sawCompletion = false;
|
||||
|
||||
const wake = () => {
|
||||
if (!pending) return;
|
||||
const resolve = pending;
|
||||
pending = null;
|
||||
resolve();
|
||||
};
|
||||
|
||||
const onMessage: WebSocketListener = (event) => {
|
||||
void (async () => {
|
||||
if (!event || typeof event !== "object" || !("data" in event)) return;
|
||||
const text = await decodeWebSocketData((event as { data?: unknown }).data);
|
||||
if (!text) return;
|
||||
try {
|
||||
const parsed = JSON.parse(text) as Record<string, unknown>;
|
||||
const type = typeof parsed.type === "string" ? parsed.type : "";
|
||||
if (type === "response.completed" || type === "response.done" || type === "response.incomplete") {
|
||||
sawCompletion = true;
|
||||
done = true;
|
||||
}
|
||||
queue.push(parsed);
|
||||
wake();
|
||||
} catch {}
|
||||
})();
|
||||
};
|
||||
|
||||
const onError: WebSocketListener = (event) => {
|
||||
failed = extractWebSocketError(event);
|
||||
done = true;
|
||||
wake();
|
||||
};
|
||||
|
||||
const onClose: WebSocketListener = (event) => {
|
||||
if (sawCompletion) {
|
||||
done = true;
|
||||
wake();
|
||||
return;
|
||||
}
|
||||
if (!failed) {
|
||||
failed = extractWebSocketCloseError(event);
|
||||
}
|
||||
done = true;
|
||||
wake();
|
||||
};
|
||||
|
||||
const onAbort = () => {
|
||||
failed = new Error("Request was aborted");
|
||||
done = true;
|
||||
wake();
|
||||
};
|
||||
|
||||
socket.addEventListener("message", onMessage);
|
||||
socket.addEventListener("error", onError);
|
||||
socket.addEventListener("close", onClose);
|
||||
signal?.addEventListener("abort", onAbort);
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
if (signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
if (queue.length > 0) {
|
||||
yield queue.shift()!;
|
||||
continue;
|
||||
}
|
||||
if (done) break;
|
||||
await new Promise<void>((resolve) => {
|
||||
pending = resolve;
|
||||
});
|
||||
}
|
||||
|
||||
if (failed) {
|
||||
throw failed;
|
||||
}
|
||||
if (!sawCompletion) {
|
||||
throw new Error("WebSocket stream closed before response.completed");
|
||||
}
|
||||
} finally {
|
||||
socket.removeEventListener("message", onMessage);
|
||||
socket.removeEventListener("error", onError);
|
||||
socket.removeEventListener("close", onClose);
|
||||
signal?.removeEventListener("abort", onAbort);
|
||||
}
|
||||
}
|
||||
|
||||
async function processWebSocketStream(
|
||||
url: string,
|
||||
body: RequestBody,
|
||||
headers: Headers,
|
||||
output: AssistantMessage,
|
||||
stream: AssistantMessageEventStream,
|
||||
model: Model<"openai-codex-responses">,
|
||||
onStart: () => void,
|
||||
options?: OpenAICodexResponsesOptions,
|
||||
): Promise<void> {
|
||||
const { socket, release } = await acquireWebSocket(url, headers, options?.sessionId, options?.signal);
|
||||
let keepConnection = true;
|
||||
try {
|
||||
socket.send(JSON.stringify({ type: "response.create", ...body }));
|
||||
onStart();
|
||||
stream.push({ type: "start", partial: output });
|
||||
await processResponsesStream(mapCodexEvents(parseWebSocket(socket, options?.signal)), output, stream, model, {
|
||||
serviceTier: options?.serviceTier,
|
||||
resolveServiceTier: resolveCodexServiceTier,
|
||||
applyServiceTierPricing,
|
||||
});
|
||||
if (options?.signal?.aborted) {
|
||||
keepConnection = false;
|
||||
}
|
||||
} catch (error) {
|
||||
keepConnection = false;
|
||||
throw error;
|
||||
} finally {
|
||||
release({ keep: keepConnection });
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Error Handling
|
||||
// ============================================================================
|
||||
|
||||
async function parseErrorResponse(response: Response): Promise<{ message: string; friendlyMessage?: string }> {
|
||||
const raw = await response.text();
|
||||
let message = raw || response.statusText || "Request failed";
|
||||
let friendlyMessage: string | undefined;
|
||||
|
||||
try {
|
||||
const parsed = JSON.parse(raw) as {
|
||||
error?: { code?: string; type?: string; message?: string; plan_type?: string; resets_at?: number };
|
||||
};
|
||||
const err = parsed?.error;
|
||||
if (err) {
|
||||
const code = err.code || err.type || "";
|
||||
if (/usage_limit_reached|usage_not_included|rate_limit_exceeded/i.test(code) || response.status === 429) {
|
||||
const plan = err.plan_type ? ` (${err.plan_type.toLowerCase()} plan)` : "";
|
||||
const mins = err.resets_at
|
||||
? Math.max(0, Math.round((err.resets_at * 1000 - Date.now()) / 60000))
|
||||
: undefined;
|
||||
const when = mins !== undefined ? ` Try again in ~${mins} min.` : "";
|
||||
friendlyMessage = `You have hit your ChatGPT usage limit${plan}.${when}`.trim();
|
||||
}
|
||||
message = err.message || friendlyMessage || message;
|
||||
}
|
||||
} catch {}
|
||||
|
||||
return { message, friendlyMessage };
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Auth & Headers
|
||||
// ============================================================================
|
||||
|
||||
function extractAccountId(token: string): string {
|
||||
try {
|
||||
const parts = token.split(".");
|
||||
if (parts.length !== 3) throw new Error("Invalid token");
|
||||
const payload = JSON.parse(atob(parts[1]));
|
||||
const accountId = payload?.[JWT_CLAIM_PATH]?.chatgpt_account_id;
|
||||
if (!accountId) throw new Error("No account ID in token");
|
||||
return accountId;
|
||||
} catch {
|
||||
throw new Error("Failed to extract accountId from token");
|
||||
}
|
||||
}
|
||||
|
||||
function createCodexRequestId(): string {
|
||||
if (typeof globalThis.crypto?.randomUUID === "function") {
|
||||
return globalThis.crypto.randomUUID();
|
||||
}
|
||||
return `codex_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`;
|
||||
}
|
||||
|
||||
function buildBaseCodexHeaders(
|
||||
initHeaders: Record<string, string> | undefined,
|
||||
additionalHeaders: Record<string, string> | undefined,
|
||||
accountId: string,
|
||||
token: string,
|
||||
): Headers {
|
||||
const headers = new Headers(initHeaders);
|
||||
for (const [key, value] of Object.entries(additionalHeaders || {})) {
|
||||
headers.set(key, value);
|
||||
}
|
||||
headers.set("Authorization", `Bearer ${token}`);
|
||||
headers.set("chatgpt-account-id", accountId);
|
||||
headers.set("originator", "pi");
|
||||
const userAgent = _os ? `pi (${_os.platform()} ${_os.release()}; ${_os.arch()})` : "pi (browser)";
|
||||
headers.set("User-Agent", userAgent);
|
||||
return headers;
|
||||
}
|
||||
|
||||
function buildSSEHeaders(
|
||||
initHeaders: Record<string, string> | undefined,
|
||||
additionalHeaders: Record<string, string> | undefined,
|
||||
accountId: string,
|
||||
token: string,
|
||||
sessionId?: string,
|
||||
): Headers {
|
||||
const headers = buildBaseCodexHeaders(initHeaders, additionalHeaders, accountId, token);
|
||||
headers.set("OpenAI-Beta", "responses=experimental");
|
||||
headers.set("accept", "text/event-stream");
|
||||
headers.set("content-type", "application/json");
|
||||
|
||||
if (sessionId) {
|
||||
headers.set("session_id", sessionId);
|
||||
headers.set("x-client-request-id", sessionId);
|
||||
}
|
||||
|
||||
return headers;
|
||||
}
|
||||
|
||||
function buildWebSocketHeaders(
|
||||
initHeaders: Record<string, string> | undefined,
|
||||
additionalHeaders: Record<string, string> | undefined,
|
||||
accountId: string,
|
||||
token: string,
|
||||
requestId: string,
|
||||
): Headers {
|
||||
const headers = buildBaseCodexHeaders(initHeaders, additionalHeaders, accountId, token);
|
||||
headers.delete("accept");
|
||||
headers.delete("content-type");
|
||||
headers.delete("OpenAI-Beta");
|
||||
headers.delete("openai-beta");
|
||||
headers.set("OpenAI-Beta", OPENAI_BETA_RESPONSES_WEBSOCKETS);
|
||||
headers.set("x-client-request-id", requestId);
|
||||
headers.set("session_id", requestId);
|
||||
return headers;
|
||||
}
|
||||
@@ -1,894 +0,0 @@
|
||||
import OpenAI from "openai";
|
||||
import type {
|
||||
ChatCompletionAssistantMessageParam,
|
||||
ChatCompletionChunk,
|
||||
ChatCompletionContentPart,
|
||||
ChatCompletionContentPartImage,
|
||||
ChatCompletionContentPartText,
|
||||
ChatCompletionMessageParam,
|
||||
ChatCompletionToolMessageParam,
|
||||
} from "openai/resources/chat/completions.js";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { calculateCost, supportsXhigh } from "../models.js";
|
||||
import type {
|
||||
AssistantMessage,
|
||||
Context,
|
||||
Message,
|
||||
Model,
|
||||
OpenAICompletionsCompat,
|
||||
SimpleStreamOptions,
|
||||
StopReason,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
TextContent,
|
||||
ThinkingContent,
|
||||
Tool,
|
||||
ToolCall,
|
||||
ToolResultMessage,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { parseStreamingJson } from "../utils/json-parse.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
|
||||
/**
|
||||
* Check if conversation messages contain tool calls or tool results.
|
||||
* This is needed because Anthropic (via proxy) requires the tools param
|
||||
* to be present when messages include tool_calls or tool role messages.
|
||||
*/
|
||||
function hasToolHistory(messages: Message[]): boolean {
|
||||
for (const msg of messages) {
|
||||
if (msg.role === "toolResult") {
|
||||
return true;
|
||||
}
|
||||
if (msg.role === "assistant") {
|
||||
if (msg.content.some((block) => block.type === "toolCall")) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
export interface OpenAICompletionsOptions extends StreamOptions {
|
||||
toolChoice?: "auto" | "none" | "required" | { type: "function"; function: { name: string } };
|
||||
reasoningEffort?: "minimal" | "low" | "medium" | "high" | "xhigh";
|
||||
}
|
||||
|
||||
export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenAICompletionsOptions> = (
|
||||
model: Model<"openai-completions">,
|
||||
context: Context,
|
||||
options?: OpenAICompletionsOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
(async () => {
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
try {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider) || "";
|
||||
const client = createClient(model, context, apiKey, options?.headers);
|
||||
let params = buildParams(model, context, options);
|
||||
const nextParams = await options?.onPayload?.(params, model);
|
||||
if (nextParams !== undefined) {
|
||||
params = nextParams as OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming;
|
||||
}
|
||||
const { data: openaiStream, response } = await client.chat.completions
|
||||
.create(params, { signal: options?.signal })
|
||||
.withResponse();
|
||||
await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model);
|
||||
stream.push({ type: "start", partial: output });
|
||||
|
||||
let currentBlock: TextContent | ThinkingContent | (ToolCall & { partialArgs?: string }) | null = null;
|
||||
const blocks = output.content;
|
||||
const blockIndex = () => blocks.length - 1;
|
||||
const finishCurrentBlock = (block?: typeof currentBlock) => {
|
||||
if (block) {
|
||||
if (block.type === "text") {
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: block.text,
|
||||
partial: output,
|
||||
});
|
||||
} else if (block.type === "thinking") {
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: block.thinking,
|
||||
partial: output,
|
||||
});
|
||||
} else if (block.type === "toolCall") {
|
||||
block.arguments = parseStreamingJson(block.partialArgs);
|
||||
// Finalize in-place and strip the scratch buffer so replay only
|
||||
// carries parsed arguments.
|
||||
delete block.partialArgs;
|
||||
stream.push({
|
||||
type: "toolcall_end",
|
||||
contentIndex: blockIndex(),
|
||||
toolCall: block,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
for await (const chunk of openaiStream) {
|
||||
if (!chunk || typeof chunk !== "object") continue;
|
||||
|
||||
// OpenAI documents ChatCompletionChunk.id as the unique chat completion identifier,
|
||||
// and each chunk in a streamed completion carries the same id.
|
||||
output.responseId ||= chunk.id;
|
||||
if (chunk.usage) {
|
||||
output.usage = parseChunkUsage(chunk.usage, model);
|
||||
}
|
||||
|
||||
const choice = Array.isArray(chunk.choices) ? chunk.choices[0] : undefined;
|
||||
if (!choice) continue;
|
||||
|
||||
// Fallback: some providers (e.g., Moonshot) return usage
|
||||
// in choice.usage instead of the standard chunk.usage
|
||||
if (!chunk.usage && (choice as any).usage) {
|
||||
output.usage = parseChunkUsage((choice as any).usage, model);
|
||||
}
|
||||
|
||||
if (choice.finish_reason) {
|
||||
const finishReasonResult = mapStopReason(choice.finish_reason);
|
||||
output.stopReason = finishReasonResult.stopReason;
|
||||
if (finishReasonResult.errorMessage) {
|
||||
output.errorMessage = finishReasonResult.errorMessage;
|
||||
}
|
||||
}
|
||||
|
||||
if (choice.delta) {
|
||||
if (
|
||||
choice.delta.content !== null &&
|
||||
choice.delta.content !== undefined &&
|
||||
choice.delta.content.length > 0
|
||||
) {
|
||||
if (!currentBlock || currentBlock.type !== "text") {
|
||||
finishCurrentBlock(currentBlock);
|
||||
currentBlock = { type: "text", text: "" };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
|
||||
if (currentBlock.type === "text") {
|
||||
currentBlock.text += choice.delta.content;
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: choice.delta.content,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Some endpoints return reasoning in reasoning_content (llama.cpp),
|
||||
// or reasoning (other openai compatible endpoints)
|
||||
// Use the first non-empty reasoning field to avoid duplication
|
||||
// (e.g., chutes.ai returns both reasoning_content and reasoning with same content)
|
||||
const reasoningFields = ["reasoning_content", "reasoning", "reasoning_text"];
|
||||
let foundReasoningField: string | null = null;
|
||||
for (const field of reasoningFields) {
|
||||
if (
|
||||
(choice.delta as any)[field] !== null &&
|
||||
(choice.delta as any)[field] !== undefined &&
|
||||
(choice.delta as any)[field].length > 0
|
||||
) {
|
||||
if (!foundReasoningField) {
|
||||
foundReasoningField = field;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (foundReasoningField) {
|
||||
if (!currentBlock || currentBlock.type !== "thinking") {
|
||||
finishCurrentBlock(currentBlock);
|
||||
currentBlock = {
|
||||
type: "thinking",
|
||||
thinking: "",
|
||||
thinkingSignature: foundReasoningField,
|
||||
};
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
|
||||
if (currentBlock.type === "thinking") {
|
||||
const delta = (choice.delta as any)[foundReasoningField];
|
||||
currentBlock.thinking += delta;
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
if (choice?.delta?.tool_calls) {
|
||||
for (const toolCall of choice.delta.tool_calls) {
|
||||
if (
|
||||
!currentBlock ||
|
||||
currentBlock.type !== "toolCall" ||
|
||||
(toolCall.id && currentBlock.id !== toolCall.id)
|
||||
) {
|
||||
finishCurrentBlock(currentBlock);
|
||||
currentBlock = {
|
||||
type: "toolCall",
|
||||
id: toolCall.id || "",
|
||||
name: toolCall.function?.name || "",
|
||||
arguments: {},
|
||||
partialArgs: "",
|
||||
};
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
|
||||
if (currentBlock.type === "toolCall") {
|
||||
if (toolCall.id) currentBlock.id = toolCall.id;
|
||||
if (toolCall.function?.name) currentBlock.name = toolCall.function.name;
|
||||
let delta = "";
|
||||
if (toolCall.function?.arguments) {
|
||||
delta = toolCall.function.arguments;
|
||||
currentBlock.partialArgs += toolCall.function.arguments;
|
||||
currentBlock.arguments = parseStreamingJson(currentBlock.partialArgs);
|
||||
}
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const reasoningDetails = (choice.delta as any).reasoning_details;
|
||||
if (reasoningDetails && Array.isArray(reasoningDetails)) {
|
||||
for (const detail of reasoningDetails) {
|
||||
if (detail.type === "reasoning.encrypted" && detail.id && detail.data) {
|
||||
const matchingToolCall = output.content.find(
|
||||
(b) => b.type === "toolCall" && b.id === detail.id,
|
||||
) as ToolCall | undefined;
|
||||
if (matchingToolCall) {
|
||||
matchingToolCall.thoughtSignature = JSON.stringify(detail);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
finishCurrentBlock(currentBlock);
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "aborted") {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
if (output.stopReason === "error") {
|
||||
throw new Error(output.errorMessage || "Provider returned an error stop reason");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
for (const block of output.content) {
|
||||
delete (block as { index?: number }).index;
|
||||
// partialArgs is only a streaming scratch buffer; never persist it.
|
||||
delete (block as { partialArgs?: string }).partialArgs;
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
|
||||
// Some providers via OpenRouter give additional information in this field.
|
||||
const rawMetadata = (error as any)?.error?.metadata?.raw;
|
||||
if (rawMetadata) output.errorMessage += `\n${rawMetadata}`;
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
export const streamSimpleOpenAICompletions: StreamFunction<"openai-completions", SimpleStreamOptions> = (
|
||||
model: Model<"openai-completions">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
|
||||
if (!apiKey) {
|
||||
throw new Error(`No API key for provider: ${model.provider}`);
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoningEffort = supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning);
|
||||
const toolChoice = (options as OpenAICompletionsOptions | undefined)?.toolChoice;
|
||||
|
||||
return streamOpenAICompletions(model, context, {
|
||||
...base,
|
||||
reasoningEffort,
|
||||
toolChoice,
|
||||
} satisfies OpenAICompletionsOptions);
|
||||
};
|
||||
|
||||
function createClient(
|
||||
model: Model<"openai-completions">,
|
||||
context: Context,
|
||||
apiKey?: string,
|
||||
optionsHeaders?: Record<string, string>,
|
||||
) {
|
||||
if (!apiKey) {
|
||||
if (!process.env.OPENAI_API_KEY) {
|
||||
throw new Error(
|
||||
"OpenAI API key is required. Set OPENAI_API_KEY environment variable or pass it as an argument.",
|
||||
);
|
||||
}
|
||||
apiKey = process.env.OPENAI_API_KEY;
|
||||
}
|
||||
|
||||
const headers = { ...model.headers };
|
||||
if (model.provider === "github-copilot") {
|
||||
const hasImages = hasCopilotVisionInput(context.messages);
|
||||
const copilotHeaders = buildCopilotDynamicHeaders({
|
||||
messages: context.messages,
|
||||
hasImages,
|
||||
});
|
||||
Object.assign(headers, copilotHeaders);
|
||||
}
|
||||
|
||||
// Merge options headers last so they can override defaults
|
||||
if (optionsHeaders) {
|
||||
Object.assign(headers, optionsHeaders);
|
||||
}
|
||||
|
||||
return new OpenAI({
|
||||
apiKey,
|
||||
baseURL: model.baseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: headers,
|
||||
});
|
||||
}
|
||||
|
||||
function buildParams(model: Model<"openai-completions">, context: Context, options?: OpenAICompletionsOptions) {
|
||||
const compat = getCompat(model);
|
||||
const messages = convertMessages(model, context, compat);
|
||||
maybeAddOpenRouterAnthropicCacheControl(model, messages);
|
||||
|
||||
const params: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = {
|
||||
model: model.id,
|
||||
messages,
|
||||
stream: true,
|
||||
};
|
||||
|
||||
if (compat.supportsUsageInStreaming !== false) {
|
||||
(params as any).stream_options = { include_usage: true };
|
||||
}
|
||||
|
||||
if (compat.supportsStore) {
|
||||
params.store = false;
|
||||
}
|
||||
|
||||
if (options?.maxTokens) {
|
||||
if (compat.maxTokensField === "max_tokens") {
|
||||
(params as any).max_tokens = options.maxTokens;
|
||||
} else {
|
||||
params.max_completion_tokens = options.maxTokens;
|
||||
}
|
||||
}
|
||||
|
||||
if (options?.temperature !== undefined) {
|
||||
params.temperature = options.temperature;
|
||||
}
|
||||
|
||||
if (context.tools) {
|
||||
params.tools = convertTools(context.tools, compat);
|
||||
if (compat.zaiToolStream) {
|
||||
(params as any).tool_stream = true;
|
||||
}
|
||||
} else if (hasToolHistory(context.messages)) {
|
||||
// Anthropic (via LiteLLM/proxy) requires tools param when conversation has tool_calls/tool_results
|
||||
params.tools = [];
|
||||
}
|
||||
|
||||
if (options?.toolChoice) {
|
||||
params.tool_choice = options.toolChoice;
|
||||
}
|
||||
|
||||
if (compat.thinkingFormat === "zai" && model.reasoning) {
|
||||
(params as any).enable_thinking = !!options?.reasoningEffort;
|
||||
} else if (compat.thinkingFormat === "qwen" && model.reasoning) {
|
||||
(params as any).enable_thinking = !!options?.reasoningEffort;
|
||||
} else if (compat.thinkingFormat === "qwen-chat-template" && model.reasoning) {
|
||||
(params as any).chat_template_kwargs = {
|
||||
enable_thinking: !!options?.reasoningEffort,
|
||||
preserve_thinking: true,
|
||||
};
|
||||
} else if (compat.thinkingFormat === "openrouter" && model.reasoning) {
|
||||
// OpenRouter normalizes reasoning across providers via a nested reasoning object.
|
||||
const openRouterParams = params as typeof params & { reasoning?: { effort?: string } };
|
||||
if (options?.reasoningEffort) {
|
||||
openRouterParams.reasoning = {
|
||||
effort: mapReasoningEffort(options.reasoningEffort, compat.reasoningEffortMap),
|
||||
};
|
||||
} else {
|
||||
openRouterParams.reasoning = { effort: "none" };
|
||||
}
|
||||
} else if (options?.reasoningEffort && model.reasoning && compat.supportsReasoningEffort) {
|
||||
// OpenAI-style reasoning_effort
|
||||
(params as any).reasoning_effort = mapReasoningEffort(options.reasoningEffort, compat.reasoningEffortMap);
|
||||
}
|
||||
|
||||
// OpenRouter provider routing preferences
|
||||
if (model.baseUrl.includes("openrouter.ai") && model.compat?.openRouterRouting) {
|
||||
(params as any).provider = model.compat.openRouterRouting;
|
||||
}
|
||||
|
||||
// Vercel AI Gateway provider routing preferences
|
||||
if (model.baseUrl.includes("ai-gateway.vercel.sh") && model.compat?.vercelGatewayRouting) {
|
||||
const routing = model.compat.vercelGatewayRouting;
|
||||
if (routing.only || routing.order) {
|
||||
const gatewayOptions: Record<string, string[]> = {};
|
||||
if (routing.only) gatewayOptions.only = routing.only;
|
||||
if (routing.order) gatewayOptions.order = routing.order;
|
||||
(params as any).providerOptions = { gateway: gatewayOptions };
|
||||
}
|
||||
}
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
function mapReasoningEffort(
|
||||
effort: NonNullable<OpenAICompletionsOptions["reasoningEffort"]>,
|
||||
reasoningEffortMap: Partial<Record<NonNullable<OpenAICompletionsOptions["reasoningEffort"]>, string>>,
|
||||
): string {
|
||||
return reasoningEffortMap[effort] ?? effort;
|
||||
}
|
||||
|
||||
function maybeAddOpenRouterAnthropicCacheControl(
|
||||
model: Model<"openai-completions">,
|
||||
messages: ChatCompletionMessageParam[],
|
||||
): void {
|
||||
if (model.provider !== "openrouter" || !model.id.startsWith("anthropic/")) return;
|
||||
|
||||
// Anthropic-style caching requires cache_control on a text part. Add a breakpoint
|
||||
// on the last user/assistant message (walking backwards until we find text content).
|
||||
for (let i = messages.length - 1; i >= 0; i--) {
|
||||
const msg = messages[i];
|
||||
if (msg.role !== "user" && msg.role !== "assistant") continue;
|
||||
|
||||
const content = msg.content;
|
||||
if (typeof content === "string") {
|
||||
msg.content = [
|
||||
Object.assign({ type: "text" as const, text: content }, { cache_control: { type: "ephemeral" } }),
|
||||
];
|
||||
return;
|
||||
}
|
||||
|
||||
if (!Array.isArray(content)) continue;
|
||||
|
||||
// Find last text part and add cache_control
|
||||
for (let j = content.length - 1; j >= 0; j--) {
|
||||
const part = content[j];
|
||||
if (part?.type === "text") {
|
||||
Object.assign(part, { cache_control: { type: "ephemeral" } });
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export function convertMessages(
|
||||
model: Model<"openai-completions">,
|
||||
context: Context,
|
||||
compat: Required<OpenAICompletionsCompat>,
|
||||
): ChatCompletionMessageParam[] {
|
||||
const params: ChatCompletionMessageParam[] = [];
|
||||
|
||||
const normalizeToolCallId = (id: string): string => {
|
||||
// Handle pipe-separated IDs from OpenAI Responses API
|
||||
// Format: {call_id}|{id} where {id} can be 400+ chars with special chars (+, /, =)
|
||||
// These come from providers like github-copilot, openai-codex, opencode
|
||||
// Extract just the call_id part and normalize it
|
||||
if (id.includes("|")) {
|
||||
const [callId] = id.split("|");
|
||||
// Sanitize to allowed chars and truncate to 40 chars (OpenAI limit)
|
||||
return callId.replace(/[^a-zA-Z0-9_-]/g, "_").slice(0, 40);
|
||||
}
|
||||
|
||||
if (model.provider === "openai") return id.length > 40 ? id.slice(0, 40) : id;
|
||||
return id;
|
||||
};
|
||||
|
||||
const transformedMessages = transformMessages(context.messages, model, (id) => normalizeToolCallId(id));
|
||||
|
||||
if (context.systemPrompt) {
|
||||
const useDeveloperRole = model.reasoning && compat.supportsDeveloperRole;
|
||||
const role = useDeveloperRole ? "developer" : "system";
|
||||
params.push({ role: role, content: sanitizeSurrogates(context.systemPrompt) });
|
||||
}
|
||||
|
||||
let lastRole: string | null = null;
|
||||
|
||||
for (let i = 0; i < transformedMessages.length; i++) {
|
||||
const msg = transformedMessages[i];
|
||||
// Some providers don't allow user messages directly after tool results
|
||||
// Insert a synthetic assistant message to bridge the gap
|
||||
if (compat.requiresAssistantAfterToolResult && lastRole === "toolResult" && msg.role === "user") {
|
||||
params.push({
|
||||
role: "assistant",
|
||||
content: "I have processed the tool results.",
|
||||
});
|
||||
}
|
||||
|
||||
if (msg.role === "user") {
|
||||
if (typeof msg.content === "string") {
|
||||
params.push({
|
||||
role: "user",
|
||||
content: sanitizeSurrogates(msg.content),
|
||||
});
|
||||
} else {
|
||||
const content: ChatCompletionContentPart[] = msg.content.map((item): ChatCompletionContentPart => {
|
||||
if (item.type === "text") {
|
||||
return {
|
||||
type: "text",
|
||||
text: sanitizeSurrogates(item.text),
|
||||
} satisfies ChatCompletionContentPartText;
|
||||
} else {
|
||||
return {
|
||||
type: "image_url",
|
||||
image_url: {
|
||||
url: `data:${item.mimeType};base64,${item.data}`,
|
||||
},
|
||||
} satisfies ChatCompletionContentPartImage;
|
||||
}
|
||||
});
|
||||
const filteredContent = !model.input.includes("image")
|
||||
? content.filter((c) => c.type !== "image_url")
|
||||
: content;
|
||||
if (filteredContent.length === 0) continue;
|
||||
params.push({
|
||||
role: "user",
|
||||
content: filteredContent,
|
||||
});
|
||||
}
|
||||
} else if (msg.role === "assistant") {
|
||||
// Some providers don't accept null content, use empty string instead
|
||||
const assistantMsg: ChatCompletionAssistantMessageParam = {
|
||||
role: "assistant",
|
||||
content: compat.requiresAssistantAfterToolResult ? "" : null,
|
||||
};
|
||||
|
||||
const textBlocks = msg.content.filter((b) => b.type === "text") as TextContent[];
|
||||
// Filter out empty text blocks to avoid API validation errors
|
||||
const nonEmptyTextBlocks = textBlocks.filter((b) => b.text && b.text.trim().length > 0);
|
||||
if (nonEmptyTextBlocks.length > 0) {
|
||||
// Always send assistant content as a plain string (OpenAI Chat Completions
|
||||
// API standard format). Sending as an array of {type:"text", text:"..."}
|
||||
// objects is non-standard and causes some models (e.g. DeepSeek V3.2 via
|
||||
// NVIDIA NIM) to mirror the content-block structure literally in their
|
||||
// output, producing recursive nesting like [{'type':'text','text':'[{...}]'}].
|
||||
assistantMsg.content = nonEmptyTextBlocks.map((b) => sanitizeSurrogates(b.text)).join("");
|
||||
}
|
||||
|
||||
// Handle thinking blocks
|
||||
const thinkingBlocks = msg.content.filter((b) => b.type === "thinking") as ThinkingContent[];
|
||||
// Filter out empty thinking blocks to avoid API validation errors
|
||||
const nonEmptyThinkingBlocks = thinkingBlocks.filter((b) => b.thinking && b.thinking.trim().length > 0);
|
||||
if (nonEmptyThinkingBlocks.length > 0) {
|
||||
if (compat.requiresThinkingAsText) {
|
||||
// Convert thinking blocks to plain text (no tags to avoid model mimicking them)
|
||||
const thinkingText = nonEmptyThinkingBlocks.map((b) => b.thinking).join("\n\n");
|
||||
const textContent = assistantMsg.content as Array<{ type: "text"; text: string }> | null;
|
||||
if (textContent) {
|
||||
textContent.unshift({ type: "text", text: thinkingText });
|
||||
} else {
|
||||
assistantMsg.content = [{ type: "text", text: thinkingText }];
|
||||
}
|
||||
} else {
|
||||
// Use the signature from the first thinking block if available (for llama.cpp server + gpt-oss)
|
||||
const signature = nonEmptyThinkingBlocks[0].thinkingSignature;
|
||||
if (signature && signature.length > 0) {
|
||||
(assistantMsg as any)[signature] = nonEmptyThinkingBlocks.map((b) => b.thinking).join("\n");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const toolCalls = msg.content.filter((b) => b.type === "toolCall") as ToolCall[];
|
||||
if (toolCalls.length > 0) {
|
||||
assistantMsg.tool_calls = toolCalls.map((tc) => ({
|
||||
id: tc.id,
|
||||
type: "function" as const,
|
||||
function: {
|
||||
name: tc.name,
|
||||
arguments: JSON.stringify(tc.arguments),
|
||||
},
|
||||
}));
|
||||
const reasoningDetails = toolCalls
|
||||
.filter((tc) => tc.thoughtSignature)
|
||||
.map((tc) => {
|
||||
try {
|
||||
return JSON.parse(tc.thoughtSignature!);
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
})
|
||||
.filter(Boolean);
|
||||
if (reasoningDetails.length > 0) {
|
||||
(assistantMsg as any).reasoning_details = reasoningDetails;
|
||||
}
|
||||
}
|
||||
// Skip assistant messages that have no content and no tool calls.
|
||||
// Some providers require "either content or tool_calls, but not none".
|
||||
// Other providers also don't accept empty assistant messages.
|
||||
// This handles aborted assistant responses that got no content.
|
||||
const content = assistantMsg.content;
|
||||
const hasContent =
|
||||
content !== null &&
|
||||
content !== undefined &&
|
||||
(typeof content === "string" ? content.length > 0 : content.length > 0);
|
||||
if (!hasContent && !assistantMsg.tool_calls) {
|
||||
continue;
|
||||
}
|
||||
params.push(assistantMsg);
|
||||
} else if (msg.role === "toolResult") {
|
||||
const imageBlocks: Array<{ type: "image_url"; image_url: { url: string } }> = [];
|
||||
let j = i;
|
||||
|
||||
for (; j < transformedMessages.length && transformedMessages[j].role === "toolResult"; j++) {
|
||||
const toolMsg = transformedMessages[j] as ToolResultMessage;
|
||||
|
||||
// Extract text and image content
|
||||
const textResult = toolMsg.content
|
||||
.filter((c) => c.type === "text")
|
||||
.map((c) => (c as any).text)
|
||||
.join("\n");
|
||||
const hasImages = toolMsg.content.some((c) => c.type === "image");
|
||||
|
||||
// Always send tool result with text (or placeholder if only images)
|
||||
const hasText = textResult.length > 0;
|
||||
// Some providers require the 'name' field in tool results
|
||||
const toolResultMsg: ChatCompletionToolMessageParam = {
|
||||
role: "tool",
|
||||
content: sanitizeSurrogates(hasText ? textResult : "(see attached image)"),
|
||||
tool_call_id: toolMsg.toolCallId,
|
||||
};
|
||||
if (compat.requiresToolResultName && toolMsg.toolName) {
|
||||
(toolResultMsg as any).name = toolMsg.toolName;
|
||||
}
|
||||
params.push(toolResultMsg);
|
||||
|
||||
if (hasImages && model.input.includes("image")) {
|
||||
for (const block of toolMsg.content) {
|
||||
if (block.type === "image") {
|
||||
imageBlocks.push({
|
||||
type: "image_url",
|
||||
image_url: {
|
||||
url: `data:${(block as any).mimeType};base64,${(block as any).data}`,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
i = j - 1;
|
||||
|
||||
if (imageBlocks.length > 0) {
|
||||
if (compat.requiresAssistantAfterToolResult) {
|
||||
params.push({
|
||||
role: "assistant",
|
||||
content: "I have processed the tool results.",
|
||||
});
|
||||
}
|
||||
|
||||
params.push({
|
||||
role: "user",
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: "Attached image(s) from tool result:",
|
||||
},
|
||||
...imageBlocks,
|
||||
],
|
||||
});
|
||||
lastRole = "user";
|
||||
} else {
|
||||
lastRole = "toolResult";
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
lastRole = msg.role;
|
||||
}
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
function convertTools(
|
||||
tools: Tool[],
|
||||
compat: Required<OpenAICompletionsCompat>,
|
||||
): OpenAI.Chat.Completions.ChatCompletionTool[] {
|
||||
return tools.map((tool) => ({
|
||||
type: "function",
|
||||
function: {
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
parameters: tool.parameters as any, // TypeBox already generates JSON Schema
|
||||
// Only include strict if provider supports it. Some reject unknown fields.
|
||||
...(compat.supportsStrictMode !== false && { strict: false }),
|
||||
},
|
||||
}));
|
||||
}
|
||||
|
||||
function parseChunkUsage(
|
||||
rawUsage: {
|
||||
prompt_tokens?: number;
|
||||
completion_tokens?: number;
|
||||
prompt_tokens_details?: { cached_tokens?: number; cache_write_tokens?: number };
|
||||
completion_tokens_details?: { reasoning_tokens?: number };
|
||||
},
|
||||
model: Model<"openai-completions">,
|
||||
): AssistantMessage["usage"] {
|
||||
const promptTokens = rawUsage.prompt_tokens || 0;
|
||||
const reportedCachedTokens = rawUsage.prompt_tokens_details?.cached_tokens || 0;
|
||||
const cacheWriteTokens = rawUsage.prompt_tokens_details?.cache_write_tokens || 0;
|
||||
const reasoningTokens = rawUsage.completion_tokens_details?.reasoning_tokens || 0;
|
||||
|
||||
// Normalize to pi-ai semantics:
|
||||
// - cacheRead: hits from cache created by previous requests only
|
||||
// - cacheWrite: tokens written to cache in this request
|
||||
// Some OpenAI-compatible providers (observed on OpenRouter) report cached_tokens
|
||||
// as (previous hits + current writes). In that case, remove cacheWrite from cacheRead.
|
||||
const cacheReadTokens =
|
||||
cacheWriteTokens > 0 ? Math.max(0, reportedCachedTokens - cacheWriteTokens) : reportedCachedTokens;
|
||||
|
||||
const input = Math.max(0, promptTokens - cacheReadTokens - cacheWriteTokens);
|
||||
// Compute totalTokens ourselves since we add reasoning_tokens to output
|
||||
// and some providers (e.g., Groq) don't include them in total_tokens
|
||||
const outputTokens = (rawUsage.completion_tokens || 0) + reasoningTokens;
|
||||
const usage: AssistantMessage["usage"] = {
|
||||
input,
|
||||
output: outputTokens,
|
||||
cacheRead: cacheReadTokens,
|
||||
cacheWrite: cacheWriteTokens,
|
||||
totalTokens: input + outputTokens + cacheReadTokens + cacheWriteTokens,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
calculateCost(model, usage);
|
||||
return usage;
|
||||
}
|
||||
|
||||
function mapStopReason(reason: ChatCompletionChunk.Choice["finish_reason"] | string): {
|
||||
stopReason: StopReason;
|
||||
errorMessage?: string;
|
||||
} {
|
||||
if (reason === null) return { stopReason: "stop" };
|
||||
switch (reason) {
|
||||
case "stop":
|
||||
case "end":
|
||||
return { stopReason: "stop" };
|
||||
case "length":
|
||||
return { stopReason: "length" };
|
||||
case "function_call":
|
||||
case "tool_calls":
|
||||
return { stopReason: "toolUse" };
|
||||
case "content_filter":
|
||||
return { stopReason: "error", errorMessage: "Provider finish_reason: content_filter" };
|
||||
case "network_error":
|
||||
return { stopReason: "error", errorMessage: "Provider finish_reason: network_error" };
|
||||
default:
|
||||
return {
|
||||
stopReason: "error",
|
||||
errorMessage: `Provider finish_reason: ${reason}`,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Detect compatibility settings from provider and baseUrl for known providers.
|
||||
* Provider takes precedence over URL-based detection since it's explicitly configured.
|
||||
* Returns a fully resolved OpenAICompletionsCompat object with all fields set.
|
||||
*/
|
||||
function detectCompat(model: Model<"openai-completions">): Required<OpenAICompletionsCompat> {
|
||||
const provider = model.provider;
|
||||
const baseUrl = model.baseUrl;
|
||||
|
||||
const isZai = provider === "zai" || baseUrl.includes("api.z.ai");
|
||||
|
||||
const isNonStandard =
|
||||
provider === "cerebras" ||
|
||||
baseUrl.includes("cerebras.ai") ||
|
||||
provider === "xai" ||
|
||||
baseUrl.includes("api.x.ai") ||
|
||||
baseUrl.includes("chutes.ai") ||
|
||||
baseUrl.includes("deepseek.com") ||
|
||||
isZai ||
|
||||
provider === "opencode" ||
|
||||
baseUrl.includes("opencode.ai");
|
||||
|
||||
const useMaxTokens = baseUrl.includes("chutes.ai");
|
||||
|
||||
const isGrok = provider === "xai" || baseUrl.includes("api.x.ai");
|
||||
const isGroq = provider === "groq" || baseUrl.includes("groq.com");
|
||||
|
||||
const reasoningEffortMap =
|
||||
isGroq && model.id === "qwen/qwen3-32b"
|
||||
? {
|
||||
minimal: "default",
|
||||
low: "default",
|
||||
medium: "default",
|
||||
high: "default",
|
||||
xhigh: "default",
|
||||
}
|
||||
: {};
|
||||
return {
|
||||
supportsStore: !isNonStandard,
|
||||
supportsDeveloperRole: !isNonStandard,
|
||||
supportsReasoningEffort: !isGrok && !isZai,
|
||||
reasoningEffortMap,
|
||||
supportsUsageInStreaming: true,
|
||||
maxTokensField: useMaxTokens ? "max_tokens" : "max_completion_tokens",
|
||||
requiresToolResultName: false,
|
||||
requiresAssistantAfterToolResult: false,
|
||||
requiresThinkingAsText: false,
|
||||
thinkingFormat: isZai
|
||||
? "zai"
|
||||
: provider === "openrouter" || baseUrl.includes("openrouter.ai")
|
||||
? "openrouter"
|
||||
: "openai",
|
||||
openRouterRouting: {},
|
||||
vercelGatewayRouting: {},
|
||||
zaiToolStream: false,
|
||||
supportsStrictMode: true,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Get resolved compatibility settings for a model.
|
||||
* Uses explicit model.compat if provided, otherwise auto-detects from provider/URL.
|
||||
*/
|
||||
function getCompat(model: Model<"openai-completions">): Required<OpenAICompletionsCompat> {
|
||||
const detected = detectCompat(model);
|
||||
if (!model.compat) return detected;
|
||||
|
||||
return {
|
||||
supportsStore: model.compat.supportsStore ?? detected.supportsStore,
|
||||
supportsDeveloperRole: model.compat.supportsDeveloperRole ?? detected.supportsDeveloperRole,
|
||||
supportsReasoningEffort: model.compat.supportsReasoningEffort ?? detected.supportsReasoningEffort,
|
||||
reasoningEffortMap: model.compat.reasoningEffortMap ?? detected.reasoningEffortMap,
|
||||
supportsUsageInStreaming: model.compat.supportsUsageInStreaming ?? detected.supportsUsageInStreaming,
|
||||
maxTokensField: model.compat.maxTokensField ?? detected.maxTokensField,
|
||||
requiresToolResultName: model.compat.requiresToolResultName ?? detected.requiresToolResultName,
|
||||
requiresAssistantAfterToolResult:
|
||||
model.compat.requiresAssistantAfterToolResult ?? detected.requiresAssistantAfterToolResult,
|
||||
requiresThinkingAsText: model.compat.requiresThinkingAsText ?? detected.requiresThinkingAsText,
|
||||
thinkingFormat: model.compat.thinkingFormat ?? detected.thinkingFormat,
|
||||
openRouterRouting: model.compat.openRouterRouting ?? {},
|
||||
vercelGatewayRouting: model.compat.vercelGatewayRouting ?? detected.vercelGatewayRouting,
|
||||
zaiToolStream: model.compat.zaiToolStream ?? detected.zaiToolStream,
|
||||
supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode,
|
||||
};
|
||||
}
|
||||
@@ -1,542 +0,0 @@
|
||||
import type OpenAI from "openai";
|
||||
import type {
|
||||
Tool as OpenAITool,
|
||||
ResponseCreateParamsStreaming,
|
||||
ResponseFunctionCallOutputItemList,
|
||||
ResponseFunctionToolCall,
|
||||
ResponseInput,
|
||||
ResponseInputContent,
|
||||
ResponseInputImage,
|
||||
ResponseInputText,
|
||||
ResponseOutputMessage,
|
||||
ResponseReasoningItem,
|
||||
ResponseStreamEvent,
|
||||
} from "openai/resources/responses/responses.js";
|
||||
import { calculateCost } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
Context,
|
||||
ImageContent,
|
||||
Model,
|
||||
StopReason,
|
||||
TextContent,
|
||||
TextSignatureV1,
|
||||
ThinkingContent,
|
||||
Tool,
|
||||
ToolCall,
|
||||
Usage,
|
||||
} from "../types.js";
|
||||
import type { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { shortHash } from "../utils/hash.js";
|
||||
import { parseStreamingJson } from "../utils/json-parse.js";
|
||||
import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
|
||||
import { transformMessages } from "./transform-messages.js";
|
||||
|
||||
// =============================================================================
|
||||
// Utilities
|
||||
// =============================================================================
|
||||
|
||||
function encodeTextSignatureV1(id: string, phase?: TextSignatureV1["phase"]): string {
|
||||
const payload: TextSignatureV1 = { v: 1, id };
|
||||
if (phase) payload.phase = phase;
|
||||
return JSON.stringify(payload);
|
||||
}
|
||||
|
||||
function parseTextSignature(
|
||||
signature: string | undefined,
|
||||
): { id: string; phase?: TextSignatureV1["phase"] } | undefined {
|
||||
if (!signature) return undefined;
|
||||
if (signature.startsWith("{")) {
|
||||
try {
|
||||
const parsed = JSON.parse(signature) as Partial<TextSignatureV1>;
|
||||
if (parsed.v === 1 && typeof parsed.id === "string") {
|
||||
if (parsed.phase === "commentary" || parsed.phase === "final_answer") {
|
||||
return { id: parsed.id, phase: parsed.phase };
|
||||
}
|
||||
return { id: parsed.id };
|
||||
}
|
||||
} catch {
|
||||
// Fall through to legacy plain-string handling.
|
||||
}
|
||||
}
|
||||
return { id: signature };
|
||||
}
|
||||
|
||||
export interface OpenAIResponsesStreamOptions {
|
||||
serviceTier?: ResponseCreateParamsStreaming["service_tier"];
|
||||
resolveServiceTier?: (
|
||||
responseServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
|
||||
requestServiceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
|
||||
) => ResponseCreateParamsStreaming["service_tier"] | undefined;
|
||||
applyServiceTierPricing?: (
|
||||
usage: Usage,
|
||||
serviceTier: ResponseCreateParamsStreaming["service_tier"] | undefined,
|
||||
) => void;
|
||||
}
|
||||
|
||||
export interface ConvertResponsesMessagesOptions {
|
||||
includeSystemPrompt?: boolean;
|
||||
}
|
||||
|
||||
export interface ConvertResponsesToolsOptions {
|
||||
strict?: boolean | null;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Message conversion
|
||||
// =============================================================================
|
||||
|
||||
export function convertResponsesMessages<TApi extends Api>(
|
||||
model: Model<TApi>,
|
||||
context: Context,
|
||||
allowedToolCallProviders: ReadonlySet<string>,
|
||||
options?: ConvertResponsesMessagesOptions,
|
||||
): ResponseInput {
|
||||
const messages: ResponseInput = [];
|
||||
|
||||
const normalizeIdPart = (part: string): string => {
|
||||
const sanitized = part.replace(/[^a-zA-Z0-9_-]/g, "_");
|
||||
const normalized = sanitized.length > 64 ? sanitized.slice(0, 64) : sanitized;
|
||||
return normalized.replace(/_+$/, "");
|
||||
};
|
||||
|
||||
const buildForeignResponsesItemId = (itemId: string): string => {
|
||||
const normalized = `fc_${shortHash(itemId)}`;
|
||||
return normalized.length > 64 ? normalized.slice(0, 64) : normalized;
|
||||
};
|
||||
|
||||
const normalizeToolCallId = (id: string, _targetModel: Model<TApi>, source: AssistantMessage): string => {
|
||||
if (!allowedToolCallProviders.has(model.provider)) return normalizeIdPart(id);
|
||||
if (!id.includes("|")) return normalizeIdPart(id);
|
||||
const [callId, itemId] = id.split("|");
|
||||
const normalizedCallId = normalizeIdPart(callId);
|
||||
const isForeignToolCall = source.provider !== model.provider || source.api !== model.api;
|
||||
let normalizedItemId = isForeignToolCall ? buildForeignResponsesItemId(itemId) : normalizeIdPart(itemId);
|
||||
// OpenAI Responses API requires item id to start with "fc"
|
||||
if (!normalizedItemId.startsWith("fc_")) {
|
||||
normalizedItemId = normalizeIdPart(`fc_${normalizedItemId}`);
|
||||
}
|
||||
return `${normalizedCallId}|${normalizedItemId}`;
|
||||
};
|
||||
|
||||
const transformedMessages = transformMessages(context.messages, model, normalizeToolCallId);
|
||||
|
||||
const includeSystemPrompt = options?.includeSystemPrompt ?? true;
|
||||
if (includeSystemPrompt && context.systemPrompt) {
|
||||
const role = model.reasoning ? "developer" : "system";
|
||||
messages.push({
|
||||
role,
|
||||
content: sanitizeSurrogates(context.systemPrompt),
|
||||
});
|
||||
}
|
||||
|
||||
let msgIndex = 0;
|
||||
for (const msg of transformedMessages) {
|
||||
if (msg.role === "user") {
|
||||
if (typeof msg.content === "string") {
|
||||
messages.push({
|
||||
role: "user",
|
||||
content: [{ type: "input_text", text: sanitizeSurrogates(msg.content) }],
|
||||
});
|
||||
} else {
|
||||
const content: ResponseInputContent[] = msg.content.map((item): ResponseInputContent => {
|
||||
if (item.type === "text") {
|
||||
return {
|
||||
type: "input_text",
|
||||
text: sanitizeSurrogates(item.text),
|
||||
} satisfies ResponseInputText;
|
||||
}
|
||||
return {
|
||||
type: "input_image",
|
||||
detail: "auto",
|
||||
image_url: `data:${item.mimeType};base64,${item.data}`,
|
||||
} satisfies ResponseInputImage;
|
||||
});
|
||||
const filteredContent = !model.input.includes("image")
|
||||
? content.filter((c) => c.type !== "input_image")
|
||||
: content;
|
||||
if (filteredContent.length === 0) continue;
|
||||
messages.push({
|
||||
role: "user",
|
||||
content: filteredContent,
|
||||
});
|
||||
}
|
||||
} else if (msg.role === "assistant") {
|
||||
const output: ResponseInput = [];
|
||||
const assistantMsg = msg as AssistantMessage;
|
||||
const isDifferentModel =
|
||||
assistantMsg.model !== model.id &&
|
||||
assistantMsg.provider === model.provider &&
|
||||
assistantMsg.api === model.api;
|
||||
|
||||
for (const block of msg.content) {
|
||||
if (block.type === "thinking") {
|
||||
if (block.thinkingSignature) {
|
||||
const reasoningItem = JSON.parse(block.thinkingSignature) as ResponseReasoningItem;
|
||||
output.push(reasoningItem);
|
||||
}
|
||||
} else if (block.type === "text") {
|
||||
const textBlock = block as TextContent;
|
||||
const parsedSignature = parseTextSignature(textBlock.textSignature);
|
||||
// OpenAI requires id to be max 64 characters
|
||||
let msgId = parsedSignature?.id;
|
||||
if (!msgId) {
|
||||
msgId = `msg_${msgIndex}`;
|
||||
} else if (msgId.length > 64) {
|
||||
msgId = `msg_${shortHash(msgId)}`;
|
||||
}
|
||||
output.push({
|
||||
type: "message",
|
||||
role: "assistant",
|
||||
content: [{ type: "output_text", text: sanitizeSurrogates(textBlock.text), annotations: [] }],
|
||||
status: "completed",
|
||||
id: msgId,
|
||||
phase: parsedSignature?.phase,
|
||||
} satisfies ResponseOutputMessage);
|
||||
} else if (block.type === "toolCall") {
|
||||
const toolCall = block as ToolCall;
|
||||
const [callId, itemIdRaw] = toolCall.id.split("|");
|
||||
let itemId: string | undefined = itemIdRaw;
|
||||
|
||||
// For different-model messages, set id to undefined to avoid pairing validation.
|
||||
// OpenAI tracks which fc_xxx IDs were paired with rs_xxx reasoning items.
|
||||
// By omitting the id, we avoid triggering that validation (like cross-provider does).
|
||||
if (isDifferentModel && itemId?.startsWith("fc_")) {
|
||||
itemId = undefined;
|
||||
}
|
||||
|
||||
output.push({
|
||||
type: "function_call",
|
||||
id: itemId,
|
||||
call_id: callId,
|
||||
name: toolCall.name,
|
||||
arguments: JSON.stringify(toolCall.arguments),
|
||||
});
|
||||
}
|
||||
}
|
||||
if (output.length === 0) continue;
|
||||
messages.push(...output);
|
||||
} else if (msg.role === "toolResult") {
|
||||
const textResult = msg.content
|
||||
.filter((c): c is TextContent => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join("\n");
|
||||
const hasImages = msg.content.some((c): c is ImageContent => c.type === "image");
|
||||
const hasText = textResult.length > 0;
|
||||
const [callId] = msg.toolCallId.split("|");
|
||||
|
||||
let output: string | ResponseFunctionCallOutputItemList;
|
||||
if (hasImages && model.input.includes("image")) {
|
||||
const contentParts: ResponseFunctionCallOutputItemList = [];
|
||||
|
||||
if (hasText) {
|
||||
contentParts.push({
|
||||
type: "input_text",
|
||||
text: sanitizeSurrogates(textResult),
|
||||
});
|
||||
}
|
||||
|
||||
for (const block of msg.content) {
|
||||
if (block.type === "image") {
|
||||
contentParts.push({
|
||||
type: "input_image",
|
||||
detail: "auto",
|
||||
image_url: `data:${block.mimeType};base64,${block.data}`,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
output = contentParts;
|
||||
} else {
|
||||
output = sanitizeSurrogates(hasText ? textResult : "(see attached image)");
|
||||
}
|
||||
|
||||
messages.push({
|
||||
type: "function_call_output",
|
||||
call_id: callId,
|
||||
output,
|
||||
});
|
||||
}
|
||||
msgIndex++;
|
||||
}
|
||||
|
||||
return messages;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tool conversion
|
||||
// =============================================================================
|
||||
|
||||
export function convertResponsesTools(tools: Tool[], options?: ConvertResponsesToolsOptions): OpenAITool[] {
|
||||
const strict = options?.strict === undefined ? false : options.strict;
|
||||
return tools.map((tool) => ({
|
||||
type: "function",
|
||||
name: tool.name,
|
||||
description: tool.description,
|
||||
parameters: tool.parameters as any, // TypeBox already generates JSON Schema
|
||||
strict,
|
||||
}));
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Stream processing
|
||||
// =============================================================================
|
||||
|
||||
export async function processResponsesStream<TApi extends Api>(
|
||||
openaiStream: AsyncIterable<ResponseStreamEvent>,
|
||||
output: AssistantMessage,
|
||||
stream: AssistantMessageEventStream,
|
||||
model: Model<TApi>,
|
||||
options?: OpenAIResponsesStreamOptions,
|
||||
): Promise<void> {
|
||||
let currentItem: ResponseReasoningItem | ResponseOutputMessage | ResponseFunctionToolCall | null = null;
|
||||
let currentBlock: ThinkingContent | TextContent | (ToolCall & { partialJson: string }) | null = null;
|
||||
const blocks = output.content;
|
||||
const blockIndex = () => blocks.length - 1;
|
||||
|
||||
for await (const event of openaiStream) {
|
||||
if (event.type === "response.created") {
|
||||
output.responseId = event.response.id;
|
||||
} else if (event.type === "response.output_item.added") {
|
||||
const item = event.item;
|
||||
if (item.type === "reasoning") {
|
||||
currentItem = item;
|
||||
currentBlock = { type: "thinking", thinking: "" };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "thinking_start", contentIndex: blockIndex(), partial: output });
|
||||
} else if (item.type === "message") {
|
||||
currentItem = item;
|
||||
currentBlock = { type: "text", text: "" };
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "text_start", contentIndex: blockIndex(), partial: output });
|
||||
} else if (item.type === "function_call") {
|
||||
currentItem = item;
|
||||
currentBlock = {
|
||||
type: "toolCall",
|
||||
id: `${item.call_id}|${item.id}`,
|
||||
name: item.name,
|
||||
arguments: {},
|
||||
partialJson: item.arguments || "",
|
||||
};
|
||||
output.content.push(currentBlock);
|
||||
stream.push({ type: "toolcall_start", contentIndex: blockIndex(), partial: output });
|
||||
}
|
||||
} else if (event.type === "response.reasoning_summary_part.added") {
|
||||
if (currentItem && currentItem.type === "reasoning") {
|
||||
currentItem.summary = currentItem.summary || [];
|
||||
currentItem.summary.push(event.part);
|
||||
}
|
||||
} else if (event.type === "response.reasoning_summary_text.delta") {
|
||||
if (currentItem?.type === "reasoning" && currentBlock?.type === "thinking") {
|
||||
currentItem.summary = currentItem.summary || [];
|
||||
const lastPart = currentItem.summary[currentItem.summary.length - 1];
|
||||
if (lastPart) {
|
||||
currentBlock.thinking += event.delta;
|
||||
lastPart.text += event.delta;
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: event.delta,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
} else if (event.type === "response.reasoning_summary_part.done") {
|
||||
if (currentItem?.type === "reasoning" && currentBlock?.type === "thinking") {
|
||||
currentItem.summary = currentItem.summary || [];
|
||||
const lastPart = currentItem.summary[currentItem.summary.length - 1];
|
||||
if (lastPart) {
|
||||
currentBlock.thinking += "\n\n";
|
||||
lastPart.text += "\n\n";
|
||||
stream.push({
|
||||
type: "thinking_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: "\n\n",
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
} else if (event.type === "response.content_part.added") {
|
||||
if (currentItem?.type === "message") {
|
||||
currentItem.content = currentItem.content || [];
|
||||
// Filter out ReasoningText, only accept output_text and refusal
|
||||
if (event.part.type === "output_text" || event.part.type === "refusal") {
|
||||
currentItem.content.push(event.part);
|
||||
}
|
||||
}
|
||||
} else if (event.type === "response.output_text.delta") {
|
||||
if (currentItem?.type === "message" && currentBlock?.type === "text") {
|
||||
if (!currentItem.content || currentItem.content.length === 0) {
|
||||
continue;
|
||||
}
|
||||
const lastPart = currentItem.content[currentItem.content.length - 1];
|
||||
if (lastPart?.type === "output_text") {
|
||||
currentBlock.text += event.delta;
|
||||
lastPart.text += event.delta;
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: event.delta,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
} else if (event.type === "response.refusal.delta") {
|
||||
if (currentItem?.type === "message" && currentBlock?.type === "text") {
|
||||
if (!currentItem.content || currentItem.content.length === 0) {
|
||||
continue;
|
||||
}
|
||||
const lastPart = currentItem.content[currentItem.content.length - 1];
|
||||
if (lastPart?.type === "refusal") {
|
||||
currentBlock.text += event.delta;
|
||||
lastPart.refusal += event.delta;
|
||||
stream.push({
|
||||
type: "text_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: event.delta,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
} else if (event.type === "response.function_call_arguments.delta") {
|
||||
if (currentItem?.type === "function_call" && currentBlock?.type === "toolCall") {
|
||||
currentBlock.partialJson += event.delta;
|
||||
currentBlock.arguments = parseStreamingJson(currentBlock.partialJson);
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta: event.delta,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
} else if (event.type === "response.function_call_arguments.done") {
|
||||
if (currentItem?.type === "function_call" && currentBlock?.type === "toolCall") {
|
||||
const previousPartialJson = currentBlock.partialJson;
|
||||
currentBlock.partialJson = event.arguments;
|
||||
currentBlock.arguments = parseStreamingJson(currentBlock.partialJson);
|
||||
|
||||
if (event.arguments.startsWith(previousPartialJson)) {
|
||||
const delta = event.arguments.slice(previousPartialJson.length);
|
||||
if (delta.length > 0) {
|
||||
stream.push({
|
||||
type: "toolcall_delta",
|
||||
contentIndex: blockIndex(),
|
||||
delta,
|
||||
partial: output,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if (event.type === "response.output_item.done") {
|
||||
const item = event.item;
|
||||
|
||||
if (item.type === "reasoning" && currentBlock?.type === "thinking") {
|
||||
currentBlock.thinking = item.summary?.map((s) => s.text).join("\n\n") || "";
|
||||
currentBlock.thinkingSignature = JSON.stringify(item);
|
||||
stream.push({
|
||||
type: "thinking_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.thinking,
|
||||
partial: output,
|
||||
});
|
||||
currentBlock = null;
|
||||
} else if (item.type === "message" && currentBlock?.type === "text") {
|
||||
currentBlock.text = item.content.map((c) => (c.type === "output_text" ? c.text : c.refusal)).join("");
|
||||
currentBlock.textSignature = encodeTextSignatureV1(item.id, item.phase ?? undefined);
|
||||
stream.push({
|
||||
type: "text_end",
|
||||
contentIndex: blockIndex(),
|
||||
content: currentBlock.text,
|
||||
partial: output,
|
||||
});
|
||||
currentBlock = null;
|
||||
} else if (item.type === "function_call") {
|
||||
const args =
|
||||
currentBlock?.type === "toolCall" && currentBlock.partialJson
|
||||
? parseStreamingJson(currentBlock.partialJson)
|
||||
: parseStreamingJson(item.arguments || "{}");
|
||||
|
||||
let toolCall: ToolCall;
|
||||
if (currentBlock?.type === "toolCall") {
|
||||
// Finalize in-place and strip the scratch buffer so replay only
|
||||
// carries parsed arguments.
|
||||
currentBlock.arguments = args;
|
||||
delete (currentBlock as { partialJson?: string }).partialJson;
|
||||
toolCall = currentBlock;
|
||||
} else {
|
||||
toolCall = {
|
||||
type: "toolCall",
|
||||
id: `${item.call_id}|${item.id}`,
|
||||
name: item.name,
|
||||
arguments: args,
|
||||
};
|
||||
}
|
||||
|
||||
currentBlock = null;
|
||||
stream.push({ type: "toolcall_end", contentIndex: blockIndex(), toolCall, partial: output });
|
||||
}
|
||||
} else if (event.type === "response.completed") {
|
||||
const response = event.response;
|
||||
if (response?.id) {
|
||||
output.responseId = response.id;
|
||||
}
|
||||
if (response?.usage) {
|
||||
const cachedTokens = response.usage.input_tokens_details?.cached_tokens || 0;
|
||||
output.usage = {
|
||||
// OpenAI includes cached tokens in input_tokens, so subtract to get non-cached input
|
||||
input: (response.usage.input_tokens || 0) - cachedTokens,
|
||||
output: response.usage.output_tokens || 0,
|
||||
cacheRead: cachedTokens,
|
||||
cacheWrite: 0,
|
||||
totalTokens: response.usage.total_tokens || 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
}
|
||||
calculateCost(model, output.usage);
|
||||
if (options?.applyServiceTierPricing) {
|
||||
const serviceTier = options.resolveServiceTier
|
||||
? options.resolveServiceTier(response?.service_tier, options.serviceTier)
|
||||
: (response?.service_tier ?? options.serviceTier);
|
||||
options.applyServiceTierPricing(output.usage, serviceTier);
|
||||
}
|
||||
// Map status to stop reason
|
||||
output.stopReason = mapStopReason(response?.status);
|
||||
if (output.content.some((b) => b.type === "toolCall") && output.stopReason === "stop") {
|
||||
output.stopReason = "toolUse";
|
||||
}
|
||||
} else if (event.type === "error") {
|
||||
throw new Error(`Error Code ${event.code}: ${event.message}` || "Unknown error");
|
||||
} else if (event.type === "response.failed") {
|
||||
const error = event.response?.error;
|
||||
const details = event.response?.incomplete_details;
|
||||
const msg = error
|
||||
? `${error.code || "unknown"}: ${error.message || "no message"}`
|
||||
: details?.reason
|
||||
? `incomplete: ${details.reason}`
|
||||
: "Unknown error (no error details in response)";
|
||||
throw new Error(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function mapStopReason(status: OpenAI.Responses.ResponseStatus | undefined): StopReason {
|
||||
if (!status) return "stop";
|
||||
switch (status) {
|
||||
case "completed":
|
||||
return "stop";
|
||||
case "incomplete":
|
||||
return "length";
|
||||
case "failed":
|
||||
case "cancelled":
|
||||
return "error";
|
||||
// These two are wonky ...
|
||||
case "in_progress":
|
||||
case "queued":
|
||||
return "stop";
|
||||
default: {
|
||||
const _exhaustive: never = status;
|
||||
throw new Error(`Unhandled stop reason: ${_exhaustive}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,264 +0,0 @@
|
||||
import OpenAI from "openai";
|
||||
import type { ResponseCreateParamsStreaming } from "openai/resources/responses/responses.js";
|
||||
import { getEnvApiKey } from "../env-api-keys.js";
|
||||
import { supportsXhigh } from "../models.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
CacheRetention,
|
||||
Context,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
Usage,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import { headersToRecord } from "../utils/headers.js";
|
||||
import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js";
|
||||
import { convertResponsesMessages, convertResponsesTools, processResponsesStream } from "./openai-responses-shared.js";
|
||||
import { buildBaseOptions, clampReasoning } from "./simple-options.js";
|
||||
|
||||
const OPENAI_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode"]);
|
||||
|
||||
/**
|
||||
* Resolve cache retention preference.
|
||||
* Defaults to "short" and uses PI_CACHE_RETENTION for backward compatibility.
|
||||
*/
|
||||
function resolveCacheRetention(cacheRetention?: CacheRetention): CacheRetention {
|
||||
if (cacheRetention) {
|
||||
return cacheRetention;
|
||||
}
|
||||
if (typeof process !== "undefined" && process.env.PI_CACHE_RETENTION === "long") {
|
||||
return "long";
|
||||
}
|
||||
return "short";
|
||||
}
|
||||
|
||||
/**
|
||||
* Get prompt cache retention based on cacheRetention and base URL.
|
||||
* Only applies to direct OpenAI API calls (api.openai.com).
|
||||
*/
|
||||
function getPromptCacheRetention(baseUrl: string, cacheRetention: CacheRetention): "24h" | undefined {
|
||||
if (cacheRetention !== "long") {
|
||||
return undefined;
|
||||
}
|
||||
if (baseUrl.includes("api.openai.com")) {
|
||||
return "24h";
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
// OpenAI Responses-specific options
|
||||
export interface OpenAIResponsesOptions extends StreamOptions {
|
||||
reasoningEffort?: "minimal" | "low" | "medium" | "high" | "xhigh";
|
||||
reasoningSummary?: "auto" | "detailed" | "concise" | null;
|
||||
serviceTier?: ResponseCreateParamsStreaming["service_tier"];
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate function for OpenAI Responses API
|
||||
*/
|
||||
export const streamOpenAIResponses: StreamFunction<"openai-responses", OpenAIResponsesOptions> = (
|
||||
model: Model<"openai-responses">,
|
||||
context: Context,
|
||||
options?: OpenAIResponsesOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
// Start async processing
|
||||
(async () => {
|
||||
const output: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: model.api as Api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
try {
|
||||
// Create OpenAI client
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider) || "";
|
||||
const cacheRetention = resolveCacheRetention(options?.cacheRetention);
|
||||
const cacheSessionId = cacheRetention === "none" ? undefined : options?.sessionId;
|
||||
const client = createClient(model, context, apiKey, options?.headers, cacheSessionId);
|
||||
let params = buildParams(model, context, options);
|
||||
const nextParams = await options?.onPayload?.(params, model);
|
||||
if (nextParams !== undefined) {
|
||||
params = nextParams as ResponseCreateParamsStreaming;
|
||||
}
|
||||
const { data: openaiStream, response } = await client.responses
|
||||
.create(params, options?.signal ? { signal: options.signal } : undefined)
|
||||
.withResponse();
|
||||
await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model);
|
||||
stream.push({ type: "start", partial: output });
|
||||
|
||||
await processResponsesStream(openaiStream, output, stream, model, {
|
||||
serviceTier: options?.serviceTier,
|
||||
applyServiceTierPricing,
|
||||
});
|
||||
|
||||
if (options?.signal?.aborted) {
|
||||
throw new Error("Request was aborted");
|
||||
}
|
||||
|
||||
if (output.stopReason === "aborted" || output.stopReason === "error") {
|
||||
throw new Error("An unknown error occurred");
|
||||
}
|
||||
|
||||
stream.push({ type: "done", reason: output.stopReason, message: output });
|
||||
stream.end();
|
||||
} catch (error) {
|
||||
for (const block of output.content) {
|
||||
delete (block as { index?: number }).index;
|
||||
// partialJson is only a streaming scratch buffer; never persist it.
|
||||
delete (block as { partialJson?: string }).partialJson;
|
||||
}
|
||||
output.stopReason = options?.signal?.aborted ? "aborted" : "error";
|
||||
output.errorMessage = error instanceof Error ? error.message : JSON.stringify(error);
|
||||
stream.push({ type: "error", reason: output.stopReason, error: output });
|
||||
stream.end();
|
||||
}
|
||||
})();
|
||||
|
||||
return stream;
|
||||
};
|
||||
|
||||
export const streamSimpleOpenAIResponses: StreamFunction<"openai-responses", SimpleStreamOptions> = (
|
||||
model: Model<"openai-responses">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream => {
|
||||
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
|
||||
if (!apiKey) {
|
||||
throw new Error(`No API key for provider: ${model.provider}`);
|
||||
}
|
||||
|
||||
const base = buildBaseOptions(model, options, apiKey);
|
||||
const reasoningEffort = supportsXhigh(model) ? options?.reasoning : clampReasoning(options?.reasoning);
|
||||
|
||||
return streamOpenAIResponses(model, context, {
|
||||
...base,
|
||||
reasoningEffort,
|
||||
} satisfies OpenAIResponsesOptions);
|
||||
};
|
||||
|
||||
function createClient(
|
||||
model: Model<"openai-responses">,
|
||||
context: Context,
|
||||
apiKey?: string,
|
||||
optionsHeaders?: Record<string, string>,
|
||||
sessionId?: string,
|
||||
) {
|
||||
if (!apiKey) {
|
||||
if (!process.env.OPENAI_API_KEY) {
|
||||
throw new Error(
|
||||
"OpenAI API key is required. Set OPENAI_API_KEY environment variable or pass it as an argument.",
|
||||
);
|
||||
}
|
||||
apiKey = process.env.OPENAI_API_KEY;
|
||||
}
|
||||
|
||||
const headers = { ...model.headers };
|
||||
if (model.provider === "github-copilot") {
|
||||
const hasImages = hasCopilotVisionInput(context.messages);
|
||||
const copilotHeaders = buildCopilotDynamicHeaders({
|
||||
messages: context.messages,
|
||||
hasImages,
|
||||
});
|
||||
Object.assign(headers, copilotHeaders);
|
||||
}
|
||||
|
||||
if (sessionId) {
|
||||
headers.session_id = sessionId;
|
||||
headers["x-client-request-id"] = sessionId;
|
||||
}
|
||||
|
||||
// Merge options headers last so they can override defaults
|
||||
if (optionsHeaders) {
|
||||
Object.assign(headers, optionsHeaders);
|
||||
}
|
||||
|
||||
return new OpenAI({
|
||||
apiKey,
|
||||
baseURL: model.baseUrl,
|
||||
dangerouslyAllowBrowser: true,
|
||||
defaultHeaders: headers,
|
||||
});
|
||||
}
|
||||
|
||||
function buildParams(model: Model<"openai-responses">, context: Context, options?: OpenAIResponsesOptions) {
|
||||
const messages = convertResponsesMessages(model, context, OPENAI_TOOL_CALL_PROVIDERS);
|
||||
|
||||
const cacheRetention = resolveCacheRetention(options?.cacheRetention);
|
||||
const params: ResponseCreateParamsStreaming = {
|
||||
model: model.id,
|
||||
input: messages,
|
||||
stream: true,
|
||||
prompt_cache_key: cacheRetention === "none" ? undefined : options?.sessionId,
|
||||
prompt_cache_retention: getPromptCacheRetention(model.baseUrl, cacheRetention),
|
||||
store: false,
|
||||
};
|
||||
|
||||
if (options?.maxTokens) {
|
||||
params.max_output_tokens = options?.maxTokens;
|
||||
}
|
||||
|
||||
if (options?.temperature !== undefined) {
|
||||
params.temperature = options?.temperature;
|
||||
}
|
||||
|
||||
if (options?.serviceTier !== undefined) {
|
||||
params.service_tier = options.serviceTier;
|
||||
}
|
||||
|
||||
if (context.tools) {
|
||||
params.tools = convertResponsesTools(context.tools);
|
||||
}
|
||||
|
||||
if (model.reasoning) {
|
||||
if (options?.reasoningEffort || options?.reasoningSummary) {
|
||||
params.reasoning = {
|
||||
effort: options?.reasoningEffort || "medium",
|
||||
summary: options?.reasoningSummary || "auto",
|
||||
};
|
||||
params.include = ["reasoning.encrypted_content"];
|
||||
} else if (model.provider !== "github-copilot") {
|
||||
params.reasoning = { effort: "none" };
|
||||
}
|
||||
}
|
||||
|
||||
return params;
|
||||
}
|
||||
|
||||
function getServiceTierCostMultiplier(serviceTier: ResponseCreateParamsStreaming["service_tier"] | undefined): number {
|
||||
switch (serviceTier) {
|
||||
case "flex":
|
||||
return 0.5;
|
||||
case "priority":
|
||||
return 2;
|
||||
default:
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
function applyServiceTierPricing(usage: Usage, serviceTier: ResponseCreateParamsStreaming["service_tier"] | undefined) {
|
||||
const multiplier = getServiceTierCostMultiplier(serviceTier);
|
||||
if (multiplier === 1) return;
|
||||
|
||||
usage.cost.input *= multiplier;
|
||||
usage.cost.output *= multiplier;
|
||||
usage.cost.cacheRead *= multiplier;
|
||||
usage.cost.cacheWrite *= multiplier;
|
||||
usage.cost.total = usage.cost.input + usage.cost.output + usage.cost.cacheRead + usage.cost.cacheWrite;
|
||||
}
|
||||
@@ -1,433 +0,0 @@
|
||||
import { clearApiProviders, registerApiProvider } from "../api-registry.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
AssistantMessageEvent,
|
||||
Context,
|
||||
Model,
|
||||
SimpleStreamOptions,
|
||||
StreamFunction,
|
||||
StreamOptions,
|
||||
} from "../types.js";
|
||||
import { AssistantMessageEventStream } from "../utils/event-stream.js";
|
||||
import type { BedrockOptions } from "./amazon-bedrock.js";
|
||||
import type { AnthropicOptions } from "./anthropic.js";
|
||||
import type { AzureOpenAIResponsesOptions } from "./azure-openai-responses.js";
|
||||
import type { GoogleOptions } from "./google.js";
|
||||
import type { GoogleGeminiCliOptions } from "./google-gemini-cli.js";
|
||||
import type { GoogleVertexOptions } from "./google-vertex.js";
|
||||
import type { MistralOptions } from "./mistral.js";
|
||||
import type { OpenAICodexResponsesOptions } from "./openai-codex-responses.js";
|
||||
import type { OpenAICompletionsOptions } from "./openai-completions.js";
|
||||
import type { OpenAIResponsesOptions } from "./openai-responses.js";
|
||||
|
||||
interface LazyProviderModule<
|
||||
TApi extends Api,
|
||||
TOptions extends StreamOptions,
|
||||
TSimpleOptions extends SimpleStreamOptions,
|
||||
> {
|
||||
stream: (model: Model<TApi>, context: Context, options?: TOptions) => AsyncIterable<AssistantMessageEvent>;
|
||||
streamSimple: (
|
||||
model: Model<TApi>,
|
||||
context: Context,
|
||||
options?: TSimpleOptions,
|
||||
) => AsyncIterable<AssistantMessageEvent>;
|
||||
}
|
||||
|
||||
interface AnthropicProviderModule {
|
||||
streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOptions>;
|
||||
streamSimpleAnthropic: StreamFunction<"anthropic-messages", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface AzureOpenAIResponsesProviderModule {
|
||||
streamAzureOpenAIResponses: StreamFunction<"azure-openai-responses", AzureOpenAIResponsesOptions>;
|
||||
streamSimpleAzureOpenAIResponses: StreamFunction<"azure-openai-responses", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface GoogleProviderModule {
|
||||
streamGoogle: StreamFunction<"google-generative-ai", GoogleOptions>;
|
||||
streamSimpleGoogle: StreamFunction<"google-generative-ai", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface GoogleGeminiCliProviderModule {
|
||||
streamGoogleGeminiCli: StreamFunction<"google-gemini-cli", GoogleGeminiCliOptions>;
|
||||
streamSimpleGoogleGeminiCli: StreamFunction<"google-gemini-cli", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface GoogleVertexProviderModule {
|
||||
streamGoogleVertex: StreamFunction<"google-vertex", GoogleVertexOptions>;
|
||||
streamSimpleGoogleVertex: StreamFunction<"google-vertex", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface MistralProviderModule {
|
||||
streamMistral: StreamFunction<"mistral-conversations", MistralOptions>;
|
||||
streamSimpleMistral: StreamFunction<"mistral-conversations", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface OpenAICodexResponsesProviderModule {
|
||||
streamOpenAICodexResponses: StreamFunction<"openai-codex-responses", OpenAICodexResponsesOptions>;
|
||||
streamSimpleOpenAICodexResponses: StreamFunction<"openai-codex-responses", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface OpenAICompletionsProviderModule {
|
||||
streamOpenAICompletions: StreamFunction<"openai-completions", OpenAICompletionsOptions>;
|
||||
streamSimpleOpenAICompletions: StreamFunction<"openai-completions", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface OpenAIResponsesProviderModule {
|
||||
streamOpenAIResponses: StreamFunction<"openai-responses", OpenAIResponsesOptions>;
|
||||
streamSimpleOpenAIResponses: StreamFunction<"openai-responses", SimpleStreamOptions>;
|
||||
}
|
||||
|
||||
interface BedrockProviderModule {
|
||||
streamBedrock: (
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
context: Context,
|
||||
options?: BedrockOptions,
|
||||
) => AsyncIterable<AssistantMessageEvent>;
|
||||
streamSimpleBedrock: (
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
) => AsyncIterable<AssistantMessageEvent>;
|
||||
}
|
||||
|
||||
const importNodeOnlyProvider = (specifier: string): Promise<unknown> => import(specifier);
|
||||
|
||||
let anthropicProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"anthropic-messages", AnthropicOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let azureOpenAIResponsesProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"azure-openai-responses", AzureOpenAIResponsesOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let googleProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"google-generative-ai", GoogleOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let googleGeminiCliProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"google-gemini-cli", GoogleGeminiCliOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let googleVertexProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"google-vertex", GoogleVertexOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let mistralProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"mistral-conversations", MistralOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let openAICodexResponsesProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"openai-codex-responses", OpenAICodexResponsesOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let openAICompletionsProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"openai-completions", OpenAICompletionsOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let openAIResponsesProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"openai-responses", OpenAIResponsesOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
let bedrockProviderModuleOverride:
|
||||
| LazyProviderModule<"bedrock-converse-stream", BedrockOptions, SimpleStreamOptions>
|
||||
| undefined;
|
||||
let bedrockProviderModulePromise:
|
||||
| Promise<LazyProviderModule<"bedrock-converse-stream", BedrockOptions, SimpleStreamOptions>>
|
||||
| undefined;
|
||||
|
||||
export function setBedrockProviderModule(module: BedrockProviderModule): void {
|
||||
bedrockProviderModuleOverride = {
|
||||
stream: module.streamBedrock,
|
||||
streamSimple: module.streamSimpleBedrock,
|
||||
};
|
||||
}
|
||||
|
||||
function forwardStream(target: AssistantMessageEventStream, source: AsyncIterable<AssistantMessageEvent>): void {
|
||||
(async () => {
|
||||
for await (const event of source) {
|
||||
target.push(event);
|
||||
}
|
||||
target.end();
|
||||
})();
|
||||
}
|
||||
|
||||
function createLazyLoadErrorMessage<TApi extends Api>(model: Model<TApi>, error: unknown): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "error",
|
||||
errorMessage: error instanceof Error ? error.message : String(error),
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createLazyStream<TApi extends Api, TOptions extends StreamOptions, TSimpleOptions extends SimpleStreamOptions>(
|
||||
loadModule: () => Promise<LazyProviderModule<TApi, TOptions, TSimpleOptions>>,
|
||||
): StreamFunction<TApi, TOptions> {
|
||||
return (model, context, options) => {
|
||||
const outer = new AssistantMessageEventStream();
|
||||
|
||||
loadModule()
|
||||
.then((module) => {
|
||||
const inner = module.stream(model, context, options);
|
||||
forwardStream(outer, inner);
|
||||
})
|
||||
.catch((error) => {
|
||||
const message = createLazyLoadErrorMessage(model, error);
|
||||
outer.push({ type: "error", reason: "error", error: message });
|
||||
outer.end(message);
|
||||
});
|
||||
|
||||
return outer;
|
||||
};
|
||||
}
|
||||
|
||||
function createLazySimpleStream<
|
||||
TApi extends Api,
|
||||
TOptions extends StreamOptions,
|
||||
TSimpleOptions extends SimpleStreamOptions,
|
||||
>(loadModule: () => Promise<LazyProviderModule<TApi, TOptions, TSimpleOptions>>): StreamFunction<TApi, TSimpleOptions> {
|
||||
return (model, context, options) => {
|
||||
const outer = new AssistantMessageEventStream();
|
||||
|
||||
loadModule()
|
||||
.then((module) => {
|
||||
const inner = module.streamSimple(model, context, options);
|
||||
forwardStream(outer, inner);
|
||||
})
|
||||
.catch((error) => {
|
||||
const message = createLazyLoadErrorMessage(model, error);
|
||||
outer.push({ type: "error", reason: "error", error: message });
|
||||
outer.end(message);
|
||||
});
|
||||
|
||||
return outer;
|
||||
};
|
||||
}
|
||||
|
||||
function loadAnthropicProviderModule(): Promise<
|
||||
LazyProviderModule<"anthropic-messages", AnthropicOptions, SimpleStreamOptions>
|
||||
> {
|
||||
anthropicProviderModulePromise ||= import("./anthropic.js").then((module) => {
|
||||
const provider = module as AnthropicProviderModule;
|
||||
return {
|
||||
stream: provider.streamAnthropic,
|
||||
streamSimple: provider.streamSimpleAnthropic,
|
||||
};
|
||||
});
|
||||
return anthropicProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadAzureOpenAIResponsesProviderModule(): Promise<
|
||||
LazyProviderModule<"azure-openai-responses", AzureOpenAIResponsesOptions, SimpleStreamOptions>
|
||||
> {
|
||||
azureOpenAIResponsesProviderModulePromise ||= import("./azure-openai-responses.js").then((module) => {
|
||||
const provider = module as AzureOpenAIResponsesProviderModule;
|
||||
return {
|
||||
stream: provider.streamAzureOpenAIResponses,
|
||||
streamSimple: provider.streamSimpleAzureOpenAIResponses,
|
||||
};
|
||||
});
|
||||
return azureOpenAIResponsesProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadGoogleProviderModule(): Promise<
|
||||
LazyProviderModule<"google-generative-ai", GoogleOptions, SimpleStreamOptions>
|
||||
> {
|
||||
googleProviderModulePromise ||= import("./google.js").then((module) => {
|
||||
const provider = module as GoogleProviderModule;
|
||||
return {
|
||||
stream: provider.streamGoogle,
|
||||
streamSimple: provider.streamSimpleGoogle,
|
||||
};
|
||||
});
|
||||
return googleProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadGoogleGeminiCliProviderModule(): Promise<
|
||||
LazyProviderModule<"google-gemini-cli", GoogleGeminiCliOptions, SimpleStreamOptions>
|
||||
> {
|
||||
googleGeminiCliProviderModulePromise ||= import("./google-gemini-cli.js").then((module) => {
|
||||
const provider = module as GoogleGeminiCliProviderModule;
|
||||
return {
|
||||
stream: provider.streamGoogleGeminiCli,
|
||||
streamSimple: provider.streamSimpleGoogleGeminiCli,
|
||||
};
|
||||
});
|
||||
return googleGeminiCliProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadGoogleVertexProviderModule(): Promise<
|
||||
LazyProviderModule<"google-vertex", GoogleVertexOptions, SimpleStreamOptions>
|
||||
> {
|
||||
googleVertexProviderModulePromise ||= import("./google-vertex.js").then((module) => {
|
||||
const provider = module as GoogleVertexProviderModule;
|
||||
return {
|
||||
stream: provider.streamGoogleVertex,
|
||||
streamSimple: provider.streamSimpleGoogleVertex,
|
||||
};
|
||||
});
|
||||
return googleVertexProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadMistralProviderModule(): Promise<
|
||||
LazyProviderModule<"mistral-conversations", MistralOptions, SimpleStreamOptions>
|
||||
> {
|
||||
mistralProviderModulePromise ||= import("./mistral.js").then((module) => {
|
||||
const provider = module as MistralProviderModule;
|
||||
return {
|
||||
stream: provider.streamMistral,
|
||||
streamSimple: provider.streamSimpleMistral,
|
||||
};
|
||||
});
|
||||
return mistralProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadOpenAICodexResponsesProviderModule(): Promise<
|
||||
LazyProviderModule<"openai-codex-responses", OpenAICodexResponsesOptions, SimpleStreamOptions>
|
||||
> {
|
||||
openAICodexResponsesProviderModulePromise ||= import("./openai-codex-responses.js").then((module) => {
|
||||
const provider = module as OpenAICodexResponsesProviderModule;
|
||||
return {
|
||||
stream: provider.streamOpenAICodexResponses,
|
||||
streamSimple: provider.streamSimpleOpenAICodexResponses,
|
||||
};
|
||||
});
|
||||
return openAICodexResponsesProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadOpenAICompletionsProviderModule(): Promise<
|
||||
LazyProviderModule<"openai-completions", OpenAICompletionsOptions, SimpleStreamOptions>
|
||||
> {
|
||||
openAICompletionsProviderModulePromise ||= import("./openai-completions.js").then((module) => {
|
||||
const provider = module as OpenAICompletionsProviderModule;
|
||||
return {
|
||||
stream: provider.streamOpenAICompletions,
|
||||
streamSimple: provider.streamSimpleOpenAICompletions,
|
||||
};
|
||||
});
|
||||
return openAICompletionsProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadOpenAIResponsesProviderModule(): Promise<
|
||||
LazyProviderModule<"openai-responses", OpenAIResponsesOptions, SimpleStreamOptions>
|
||||
> {
|
||||
openAIResponsesProviderModulePromise ||= import("./openai-responses.js").then((module) => {
|
||||
const provider = module as OpenAIResponsesProviderModule;
|
||||
return {
|
||||
stream: provider.streamOpenAIResponses,
|
||||
streamSimple: provider.streamSimpleOpenAIResponses,
|
||||
};
|
||||
});
|
||||
return openAIResponsesProviderModulePromise;
|
||||
}
|
||||
|
||||
function loadBedrockProviderModule(): Promise<
|
||||
LazyProviderModule<"bedrock-converse-stream", BedrockOptions, SimpleStreamOptions>
|
||||
> {
|
||||
if (bedrockProviderModuleOverride) {
|
||||
return Promise.resolve(bedrockProviderModuleOverride);
|
||||
}
|
||||
bedrockProviderModulePromise ||= importNodeOnlyProvider("./amazon-bedrock.js").then((module) => {
|
||||
const provider = module as BedrockProviderModule;
|
||||
return {
|
||||
stream: provider.streamBedrock,
|
||||
streamSimple: provider.streamSimpleBedrock,
|
||||
};
|
||||
});
|
||||
return bedrockProviderModulePromise;
|
||||
}
|
||||
|
||||
export const streamAnthropic = createLazyStream(loadAnthropicProviderModule);
|
||||
export const streamSimpleAnthropic = createLazySimpleStream(loadAnthropicProviderModule);
|
||||
export const streamAzureOpenAIResponses = createLazyStream(loadAzureOpenAIResponsesProviderModule);
|
||||
export const streamSimpleAzureOpenAIResponses = createLazySimpleStream(loadAzureOpenAIResponsesProviderModule);
|
||||
export const streamGoogle = createLazyStream(loadGoogleProviderModule);
|
||||
export const streamSimpleGoogle = createLazySimpleStream(loadGoogleProviderModule);
|
||||
export const streamGoogleGeminiCli = createLazyStream(loadGoogleGeminiCliProviderModule);
|
||||
export const streamSimpleGoogleGeminiCli = createLazySimpleStream(loadGoogleGeminiCliProviderModule);
|
||||
export const streamGoogleVertex = createLazyStream(loadGoogleVertexProviderModule);
|
||||
export const streamSimpleGoogleVertex = createLazySimpleStream(loadGoogleVertexProviderModule);
|
||||
export const streamMistral = createLazyStream(loadMistralProviderModule);
|
||||
export const streamSimpleMistral = createLazySimpleStream(loadMistralProviderModule);
|
||||
export const streamOpenAICodexResponses = createLazyStream(loadOpenAICodexResponsesProviderModule);
|
||||
export const streamSimpleOpenAICodexResponses = createLazySimpleStream(loadOpenAICodexResponsesProviderModule);
|
||||
export const streamOpenAICompletions = createLazyStream(loadOpenAICompletionsProviderModule);
|
||||
export const streamSimpleOpenAICompletions = createLazySimpleStream(loadOpenAICompletionsProviderModule);
|
||||
export const streamOpenAIResponses = createLazyStream(loadOpenAIResponsesProviderModule);
|
||||
export const streamSimpleOpenAIResponses = createLazySimpleStream(loadOpenAIResponsesProviderModule);
|
||||
const streamBedrockLazy = createLazyStream(loadBedrockProviderModule);
|
||||
const streamSimpleBedrockLazy = createLazySimpleStream(loadBedrockProviderModule);
|
||||
|
||||
export function registerBuiltInApiProviders(): void {
|
||||
registerApiProvider({
|
||||
api: "anthropic-messages",
|
||||
stream: streamAnthropic,
|
||||
streamSimple: streamSimpleAnthropic,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "openai-completions",
|
||||
stream: streamOpenAICompletions,
|
||||
streamSimple: streamSimpleOpenAICompletions,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "mistral-conversations",
|
||||
stream: streamMistral,
|
||||
streamSimple: streamSimpleMistral,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "openai-responses",
|
||||
stream: streamOpenAIResponses,
|
||||
streamSimple: streamSimpleOpenAIResponses,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "azure-openai-responses",
|
||||
stream: streamAzureOpenAIResponses,
|
||||
streamSimple: streamSimpleAzureOpenAIResponses,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "openai-codex-responses",
|
||||
stream: streamOpenAICodexResponses,
|
||||
streamSimple: streamSimpleOpenAICodexResponses,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "google-generative-ai",
|
||||
stream: streamGoogle,
|
||||
streamSimple: streamSimpleGoogle,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "google-gemini-cli",
|
||||
stream: streamGoogleGeminiCli,
|
||||
streamSimple: streamSimpleGoogleGeminiCli,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "google-vertex",
|
||||
stream: streamGoogleVertex,
|
||||
streamSimple: streamSimpleGoogleVertex,
|
||||
});
|
||||
|
||||
registerApiProvider({
|
||||
api: "bedrock-converse-stream",
|
||||
stream: streamBedrockLazy,
|
||||
streamSimple: streamSimpleBedrockLazy,
|
||||
});
|
||||
}
|
||||
|
||||
export function resetApiProviders(): void {
|
||||
clearApiProviders();
|
||||
registerBuiltInApiProviders();
|
||||
}
|
||||
|
||||
registerBuiltInApiProviders();
|
||||
@@ -1,47 +0,0 @@
|
||||
import type { Api, Model, SimpleStreamOptions, StreamOptions, ThinkingBudgets, ThinkingLevel } from "../types.js";
|
||||
|
||||
export function buildBaseOptions(model: Model<Api>, options?: SimpleStreamOptions, apiKey?: string): StreamOptions {
|
||||
return {
|
||||
temperature: options?.temperature,
|
||||
maxTokens: options?.maxTokens || Math.min(model.maxTokens, 32000),
|
||||
signal: options?.signal,
|
||||
apiKey: apiKey || options?.apiKey,
|
||||
cacheRetention: options?.cacheRetention,
|
||||
sessionId: options?.sessionId,
|
||||
headers: options?.headers,
|
||||
onPayload: options?.onPayload,
|
||||
onResponse: options?.onResponse,
|
||||
maxRetryDelayMs: options?.maxRetryDelayMs,
|
||||
metadata: options?.metadata,
|
||||
};
|
||||
}
|
||||
|
||||
export function clampReasoning(effort: ThinkingLevel | undefined): Exclude<ThinkingLevel, "xhigh"> | undefined {
|
||||
return effort === "xhigh" ? "high" : effort;
|
||||
}
|
||||
|
||||
export function adjustMaxTokensForThinking(
|
||||
baseMaxTokens: number,
|
||||
modelMaxTokens: number,
|
||||
reasoningLevel: ThinkingLevel,
|
||||
customBudgets?: ThinkingBudgets,
|
||||
): { maxTokens: number; thinkingBudget: number } {
|
||||
const defaultBudgets: ThinkingBudgets = {
|
||||
minimal: 1024,
|
||||
low: 2048,
|
||||
medium: 8192,
|
||||
high: 16384,
|
||||
};
|
||||
const budgets = { ...defaultBudgets, ...customBudgets };
|
||||
|
||||
const minOutputTokens = 1024;
|
||||
const level = clampReasoning(reasoningLevel)!;
|
||||
let thinkingBudget = budgets[level]!;
|
||||
const maxTokens = Math.min(baseMaxTokens + thinkingBudget, modelMaxTokens);
|
||||
|
||||
if (maxTokens <= thinkingBudget) {
|
||||
thinkingBudget = Math.max(0, maxTokens - minOutputTokens);
|
||||
}
|
||||
|
||||
return { maxTokens, thinkingBudget };
|
||||
}
|
||||
@@ -1,160 +0,0 @@
|
||||
import type { Api, AssistantMessage, Message, Model, ToolCall, ToolResultMessage } from "../types.js";
|
||||
|
||||
/**
|
||||
* Normalize tool call ID for cross-provider compatibility.
|
||||
* OpenAI Responses API generates IDs that are 450+ chars with special characters like `|`.
|
||||
* Anthropic APIs require IDs matching ^[a-zA-Z0-9_-]+$ (max 64 chars).
|
||||
*/
|
||||
export function transformMessages<TApi extends Api>(
|
||||
messages: Message[],
|
||||
model: Model<TApi>,
|
||||
normalizeToolCallId?: (id: string, model: Model<TApi>, source: AssistantMessage) => string,
|
||||
): Message[] {
|
||||
// Build a map of original tool call IDs to normalized IDs
|
||||
const toolCallIdMap = new Map<string, string>();
|
||||
|
||||
// First pass: transform messages (thinking blocks, tool call ID normalization)
|
||||
const transformed = messages.map((msg) => {
|
||||
// User messages pass through unchanged
|
||||
if (msg.role === "user") {
|
||||
return msg;
|
||||
}
|
||||
|
||||
// Handle toolResult messages - normalize toolCallId if we have a mapping
|
||||
if (msg.role === "toolResult") {
|
||||
const normalizedId = toolCallIdMap.get(msg.toolCallId);
|
||||
if (normalizedId && normalizedId !== msg.toolCallId) {
|
||||
return { ...msg, toolCallId: normalizedId };
|
||||
}
|
||||
return msg;
|
||||
}
|
||||
|
||||
// Assistant messages need transformation check
|
||||
if (msg.role === "assistant") {
|
||||
const assistantMsg = msg as AssistantMessage;
|
||||
const isSameModel =
|
||||
assistantMsg.provider === model.provider &&
|
||||
assistantMsg.api === model.api &&
|
||||
assistantMsg.model === model.id;
|
||||
|
||||
const transformedContent = assistantMsg.content.flatMap((block) => {
|
||||
if (block.type === "thinking") {
|
||||
// Redacted thinking is opaque encrypted content, only valid for the same model.
|
||||
// Drop it for cross-model to avoid API errors.
|
||||
if (block.redacted) {
|
||||
return isSameModel ? block : [];
|
||||
}
|
||||
// For same model: keep thinking blocks with signatures (needed for replay)
|
||||
// even if the thinking text is empty (OpenAI encrypted reasoning)
|
||||
if (isSameModel && block.thinkingSignature) return block;
|
||||
// Skip empty thinking blocks, convert others to plain text
|
||||
if (!block.thinking || block.thinking.trim() === "") return [];
|
||||
if (isSameModel) return block;
|
||||
return {
|
||||
type: "text" as const,
|
||||
text: block.thinking,
|
||||
};
|
||||
}
|
||||
|
||||
if (block.type === "text") {
|
||||
if (isSameModel) return block;
|
||||
return {
|
||||
type: "text" as const,
|
||||
text: block.text,
|
||||
};
|
||||
}
|
||||
|
||||
if (block.type === "toolCall") {
|
||||
const toolCall = block as ToolCall;
|
||||
let normalizedToolCall: ToolCall = toolCall;
|
||||
|
||||
if (!isSameModel && toolCall.thoughtSignature) {
|
||||
normalizedToolCall = { ...toolCall };
|
||||
delete (normalizedToolCall as { thoughtSignature?: string }).thoughtSignature;
|
||||
}
|
||||
|
||||
if (!isSameModel && normalizeToolCallId) {
|
||||
const normalizedId = normalizeToolCallId(toolCall.id, model, assistantMsg);
|
||||
if (normalizedId !== toolCall.id) {
|
||||
toolCallIdMap.set(toolCall.id, normalizedId);
|
||||
normalizedToolCall = { ...normalizedToolCall, id: normalizedId };
|
||||
}
|
||||
}
|
||||
|
||||
return normalizedToolCall;
|
||||
}
|
||||
|
||||
return block;
|
||||
});
|
||||
|
||||
return {
|
||||
...assistantMsg,
|
||||
content: transformedContent,
|
||||
};
|
||||
}
|
||||
return msg;
|
||||
});
|
||||
|
||||
// Second pass: insert synthetic empty tool results for orphaned tool calls
|
||||
// This preserves thinking signatures and satisfies API requirements
|
||||
const result: Message[] = [];
|
||||
let pendingToolCalls: ToolCall[] = [];
|
||||
let existingToolResultIds = new Set<string>();
|
||||
const insertSyntheticToolResults = () => {
|
||||
if (pendingToolCalls.length > 0) {
|
||||
for (const tc of pendingToolCalls) {
|
||||
if (!existingToolResultIds.has(tc.id)) {
|
||||
result.push({
|
||||
role: "toolResult",
|
||||
toolCallId: tc.id,
|
||||
toolName: tc.name,
|
||||
content: [{ type: "text", text: "No result provided" }],
|
||||
isError: true,
|
||||
timestamp: Date.now(),
|
||||
} as ToolResultMessage);
|
||||
}
|
||||
}
|
||||
pendingToolCalls = [];
|
||||
existingToolResultIds = new Set();
|
||||
}
|
||||
};
|
||||
|
||||
for (let i = 0; i < transformed.length; i++) {
|
||||
const msg = transformed[i];
|
||||
|
||||
if (msg.role === "assistant") {
|
||||
// If we have pending orphaned tool calls from a previous assistant, insert synthetic results now
|
||||
insertSyntheticToolResults();
|
||||
|
||||
// Skip errored/aborted assistant messages entirely.
|
||||
// These are incomplete turns that shouldn't be replayed:
|
||||
// - May have partial content (reasoning without message, incomplete tool calls)
|
||||
// - Replaying them can cause API errors (e.g., OpenAI "reasoning without following item")
|
||||
// - The model should retry from the last valid state
|
||||
const assistantMsg = msg as AssistantMessage;
|
||||
if (assistantMsg.stopReason === "error" || assistantMsg.stopReason === "aborted") {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Track tool calls from this assistant message
|
||||
const toolCalls = assistantMsg.content.filter((b) => b.type === "toolCall") as ToolCall[];
|
||||
if (toolCalls.length > 0) {
|
||||
pendingToolCalls = toolCalls;
|
||||
existingToolResultIds = new Set();
|
||||
}
|
||||
|
||||
result.push(msg);
|
||||
} else if (msg.role === "toolResult") {
|
||||
existingToolResultIds.add(msg.toolCallId);
|
||||
result.push(msg);
|
||||
} else if (msg.role === "user") {
|
||||
// User message interrupts tool flow - insert synthetic results for orphaned calls
|
||||
insertSyntheticToolResults();
|
||||
result.push(msg);
|
||||
} else {
|
||||
result.push(msg);
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
@@ -1,59 +0,0 @@
|
||||
import "./providers/register-builtins.js";
|
||||
|
||||
import { getApiProvider } from "./api-registry.js";
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
AssistantMessageEventStream,
|
||||
Context,
|
||||
Model,
|
||||
ProviderStreamOptions,
|
||||
SimpleStreamOptions,
|
||||
StreamOptions,
|
||||
} from "./types.js";
|
||||
|
||||
export { getEnvApiKey } from "./env-api-keys.js";
|
||||
|
||||
function resolveApiProvider(api: Api) {
|
||||
const provider = getApiProvider(api);
|
||||
if (!provider) {
|
||||
throw new Error(`No API provider registered for api: ${api}`);
|
||||
}
|
||||
return provider;
|
||||
}
|
||||
|
||||
export function stream<TApi extends Api>(
|
||||
model: Model<TApi>,
|
||||
context: Context,
|
||||
options?: ProviderStreamOptions,
|
||||
): AssistantMessageEventStream {
|
||||
const provider = resolveApiProvider(model.api);
|
||||
return provider.stream(model, context, options as StreamOptions);
|
||||
}
|
||||
|
||||
export async function complete<TApi extends Api>(
|
||||
model: Model<TApi>,
|
||||
context: Context,
|
||||
options?: ProviderStreamOptions,
|
||||
): Promise<AssistantMessage> {
|
||||
const s = stream(model, context, options);
|
||||
return s.result();
|
||||
}
|
||||
|
||||
export function streamSimple<TApi extends Api>(
|
||||
model: Model<TApi>,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): AssistantMessageEventStream {
|
||||
const provider = resolveApiProvider(model.api);
|
||||
return provider.streamSimple(model, context, options);
|
||||
}
|
||||
|
||||
export async function completeSimple<TApi extends Api>(
|
||||
model: Model<TApi>,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
): Promise<AssistantMessage> {
|
||||
const s = streamSimple(model, context, options);
|
||||
return s.result();
|
||||
}
|
||||
@@ -1,412 +0,0 @@
|
||||
import type { AssistantMessageEventStream } from "./utils/event-stream.js";
|
||||
|
||||
export type { AssistantMessageEventStream } from "./utils/event-stream.js";
|
||||
|
||||
export type KnownApi =
|
||||
| "openai-completions"
|
||||
| "mistral-conversations"
|
||||
| "openai-responses"
|
||||
| "azure-openai-responses"
|
||||
| "openai-codex-responses"
|
||||
| "anthropic-messages"
|
||||
| "bedrock-converse-stream"
|
||||
| "google-generative-ai"
|
||||
| "google-gemini-cli"
|
||||
| "google-vertex";
|
||||
|
||||
export type Api = KnownApi | (string & {});
|
||||
|
||||
export type KnownProvider =
|
||||
| "amazon-bedrock"
|
||||
| "anthropic"
|
||||
| "google"
|
||||
| "google-gemini-cli"
|
||||
| "google-antigravity"
|
||||
| "google-vertex"
|
||||
| "openai"
|
||||
| "azure-openai-responses"
|
||||
| "openai-codex"
|
||||
| "github-copilot"
|
||||
| "xai"
|
||||
| "groq"
|
||||
| "cerebras"
|
||||
| "openrouter"
|
||||
| "vercel-ai-gateway"
|
||||
| "zai"
|
||||
| "mistral"
|
||||
| "minimax"
|
||||
| "minimax-cn"
|
||||
| "huggingface"
|
||||
| "opencode"
|
||||
| "opencode-go"
|
||||
| "kimi-coding";
|
||||
export type Provider = KnownProvider | string;
|
||||
|
||||
export type ThinkingLevel = "minimal" | "low" | "medium" | "high" | "xhigh";
|
||||
|
||||
/** Token budgets for each thinking level (token-based providers only) */
|
||||
export interface ThinkingBudgets {
|
||||
minimal?: number;
|
||||
low?: number;
|
||||
medium?: number;
|
||||
high?: number;
|
||||
}
|
||||
|
||||
// Base options all providers share
|
||||
export type CacheRetention = "none" | "short" | "long";
|
||||
|
||||
export type Transport = "sse" | "websocket" | "auto";
|
||||
|
||||
export interface ProviderResponse {
|
||||
status: number;
|
||||
headers: Record<string, string>;
|
||||
}
|
||||
|
||||
export interface StreamOptions {
|
||||
temperature?: number;
|
||||
maxTokens?: number;
|
||||
signal?: AbortSignal;
|
||||
apiKey?: string;
|
||||
/**
|
||||
* Preferred transport for providers that support multiple transports.
|
||||
* Providers that do not support this option ignore it.
|
||||
*/
|
||||
transport?: Transport;
|
||||
/**
|
||||
* Prompt cache retention preference. Providers map this to their supported values.
|
||||
* Default: "short".
|
||||
*/
|
||||
cacheRetention?: CacheRetention;
|
||||
/**
|
||||
* Optional session identifier for providers that support session-based caching.
|
||||
* Providers can use this to enable prompt caching, request routing, or other
|
||||
* session-aware features. Ignored by providers that don't support it.
|
||||
*/
|
||||
sessionId?: string;
|
||||
/**
|
||||
* Optional callback for inspecting or replacing provider payloads before sending.
|
||||
* Return undefined to keep the payload unchanged.
|
||||
*/
|
||||
onPayload?: (payload: unknown, model: Model<Api>) => unknown | undefined | Promise<unknown | undefined>;
|
||||
/**
|
||||
* Optional callback invoked after an HTTP response is received and before
|
||||
* its body stream is consumed.
|
||||
*/
|
||||
onResponse?: (response: ProviderResponse, model: Model<Api>) => void | Promise<void>;
|
||||
/**
|
||||
* Optional custom HTTP headers to include in API requests.
|
||||
* Merged with provider defaults; can override default headers.
|
||||
* Not supported by all providers (e.g., AWS Bedrock uses SDK auth).
|
||||
*/
|
||||
headers?: Record<string, string>;
|
||||
/**
|
||||
* Maximum delay in milliseconds to wait for a retry when the server requests a long wait.
|
||||
* If the server's requested delay exceeds this value, the request fails immediately
|
||||
* with an error containing the requested delay, allowing higher-level retry logic
|
||||
* to handle it with user visibility.
|
||||
* Default: 60000 (60 seconds). Set to 0 to disable the cap.
|
||||
*/
|
||||
maxRetryDelayMs?: number;
|
||||
/**
|
||||
* Optional metadata to include in API requests.
|
||||
* Providers extract the fields they understand and ignore the rest.
|
||||
* For example, Anthropic uses `user_id` for abuse tracking and rate limiting.
|
||||
*/
|
||||
metadata?: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export type ProviderStreamOptions = StreamOptions & Record<string, unknown>;
|
||||
|
||||
// Unified options with reasoning passed to streamSimple() and completeSimple()
|
||||
export interface SimpleStreamOptions extends StreamOptions {
|
||||
reasoning?: ThinkingLevel;
|
||||
/** Custom token budgets for thinking levels (token-based providers only) */
|
||||
thinkingBudgets?: ThinkingBudgets;
|
||||
}
|
||||
|
||||
// Generic StreamFunction with typed options.
|
||||
//
|
||||
// Contract:
|
||||
// - Must return an AssistantMessageEventStream.
|
||||
// - Once invoked, request/model/runtime failures should be encoded in the
|
||||
// returned stream, not thrown.
|
||||
// - Error termination must produce an AssistantMessage with stopReason
|
||||
// "error" or "aborted" and errorMessage, emitted via the stream protocol.
|
||||
export type StreamFunction<TApi extends Api = Api, TOptions extends StreamOptions = StreamOptions> = (
|
||||
model: Model<TApi>,
|
||||
context: Context,
|
||||
options?: TOptions,
|
||||
) => AssistantMessageEventStream;
|
||||
|
||||
export interface TextSignatureV1 {
|
||||
v: 1;
|
||||
id: string;
|
||||
phase?: "commentary" | "final_answer";
|
||||
}
|
||||
|
||||
export interface TextContent {
|
||||
type: "text";
|
||||
text: string;
|
||||
textSignature?: string; // e.g., for OpenAI responses, message metadata (legacy id string or TextSignatureV1 JSON)
|
||||
}
|
||||
|
||||
export interface ThinkingContent {
|
||||
type: "thinking";
|
||||
thinking: string;
|
||||
thinkingSignature?: string; // e.g., for OpenAI responses, the reasoning item ID
|
||||
/** When true, the thinking content was redacted by safety filters. The opaque
|
||||
* encrypted payload is stored in `thinkingSignature` so it can be passed back
|
||||
* to the API for multi-turn continuity. */
|
||||
redacted?: boolean;
|
||||
}
|
||||
|
||||
export interface ImageContent {
|
||||
type: "image";
|
||||
data: string; // base64 encoded image data
|
||||
mimeType: string; // e.g., "image/jpeg", "image/png"
|
||||
}
|
||||
|
||||
export interface ToolCall {
|
||||
type: "toolCall";
|
||||
id: string;
|
||||
name: string;
|
||||
arguments: Record<string, any>;
|
||||
thoughtSignature?: string; // Google-specific: opaque signature for reusing thought context
|
||||
}
|
||||
|
||||
export interface Usage {
|
||||
input: number;
|
||||
output: number;
|
||||
cacheRead: number;
|
||||
cacheWrite: number;
|
||||
totalTokens: number;
|
||||
cost: {
|
||||
input: number;
|
||||
output: number;
|
||||
cacheRead: number;
|
||||
cacheWrite: number;
|
||||
total: number;
|
||||
};
|
||||
}
|
||||
|
||||
export type StopReason = "stop" | "length" | "toolUse" | "error" | "aborted";
|
||||
|
||||
export interface UserMessage {
|
||||
role: "user";
|
||||
content: string | (TextContent | ImageContent)[];
|
||||
timestamp: number; // Unix timestamp in milliseconds
|
||||
}
|
||||
|
||||
export interface AssistantMessage {
|
||||
role: "assistant";
|
||||
content: (TextContent | ThinkingContent | ToolCall)[];
|
||||
api: Api;
|
||||
provider: Provider;
|
||||
model: string;
|
||||
responseId?: string; // Provider-specific response/message identifier when the upstream API exposes one
|
||||
usage: Usage;
|
||||
stopReason: StopReason;
|
||||
errorMessage?: string;
|
||||
timestamp: number; // Unix timestamp in milliseconds
|
||||
}
|
||||
|
||||
export interface ToolResultMessage<TDetails = any> {
|
||||
role: "toolResult";
|
||||
toolCallId: string;
|
||||
toolName: string;
|
||||
content: (TextContent | ImageContent)[]; // Supports text and images
|
||||
details?: TDetails;
|
||||
isError: boolean;
|
||||
timestamp: number; // Unix timestamp in milliseconds
|
||||
}
|
||||
|
||||
export type Message = UserMessage | AssistantMessage | ToolResultMessage;
|
||||
|
||||
import type { TSchema } from "@sinclair/typebox";
|
||||
|
||||
export interface Tool<TParameters extends TSchema = TSchema> {
|
||||
name: string;
|
||||
description: string;
|
||||
parameters: TParameters;
|
||||
}
|
||||
|
||||
export interface Context {
|
||||
systemPrompt?: string;
|
||||
messages: Message[];
|
||||
tools?: Tool[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Event protocol for AssistantMessageEventStream.
|
||||
*
|
||||
* Streams should emit `start` before partial updates, then terminate with either:
|
||||
* - `done` carrying the final successful AssistantMessage, or
|
||||
* - `error` carrying the final AssistantMessage with stopReason "error" or "aborted"
|
||||
* and errorMessage.
|
||||
*/
|
||||
export type AssistantMessageEvent =
|
||||
| { type: "start"; partial: AssistantMessage }
|
||||
| { type: "text_start"; contentIndex: number; partial: AssistantMessage }
|
||||
| { type: "text_delta"; contentIndex: number; delta: string; partial: AssistantMessage }
|
||||
| { type: "text_end"; contentIndex: number; content: string; partial: AssistantMessage }
|
||||
| { type: "thinking_start"; contentIndex: number; partial: AssistantMessage }
|
||||
| { type: "thinking_delta"; contentIndex: number; delta: string; partial: AssistantMessage }
|
||||
| { type: "thinking_end"; contentIndex: number; content: string; partial: AssistantMessage }
|
||||
| { type: "toolcall_start"; contentIndex: number; partial: AssistantMessage }
|
||||
| { type: "toolcall_delta"; contentIndex: number; delta: string; partial: AssistantMessage }
|
||||
| { type: "toolcall_end"; contentIndex: number; toolCall: ToolCall; partial: AssistantMessage }
|
||||
| { type: "done"; reason: Extract<StopReason, "stop" | "length" | "toolUse">; message: AssistantMessage }
|
||||
| { type: "error"; reason: Extract<StopReason, "aborted" | "error">; error: AssistantMessage };
|
||||
|
||||
/**
|
||||
* Compatibility settings for OpenAI-compatible completions APIs.
|
||||
* Use this to override URL-based auto-detection for custom providers.
|
||||
*/
|
||||
export interface OpenAICompletionsCompat {
|
||||
/** Whether the provider supports the `store` field. Default: auto-detected from URL. */
|
||||
supportsStore?: boolean;
|
||||
/** Whether the provider supports the `developer` role (vs `system`). Default: auto-detected from URL. */
|
||||
supportsDeveloperRole?: boolean;
|
||||
/** Whether the provider supports `reasoning_effort`. Default: auto-detected from URL. */
|
||||
supportsReasoningEffort?: boolean;
|
||||
/** Optional mapping from pi-ai reasoning levels to provider/model-specific `reasoning_effort` values. */
|
||||
reasoningEffortMap?: Partial<Record<ThinkingLevel, string>>;
|
||||
/** Whether the provider supports `stream_options: { include_usage: true }` for token usage in streaming responses. Default: true. */
|
||||
supportsUsageInStreaming?: boolean;
|
||||
/** Which field to use for max tokens. Default: auto-detected from URL. */
|
||||
maxTokensField?: "max_completion_tokens" | "max_tokens";
|
||||
/** Whether tool results require the `name` field. Default: auto-detected from URL. */
|
||||
requiresToolResultName?: boolean;
|
||||
/** Whether a user message after tool results requires an assistant message in between. Default: auto-detected from URL. */
|
||||
requiresAssistantAfterToolResult?: boolean;
|
||||
/** Whether thinking blocks must be converted to text blocks with <thinking> delimiters. Default: auto-detected from URL. */
|
||||
requiresThinkingAsText?: boolean;
|
||||
/** Format for reasoning/thinking parameter. "openai" uses reasoning_effort, "openrouter" uses reasoning: { effort }, "zai" uses top-level enable_thinking: boolean, "qwen" uses top-level enable_thinking: boolean, and "qwen-chat-template" uses chat_template_kwargs.enable_thinking. Default: "openai". */
|
||||
thinkingFormat?: "openai" | "openrouter" | "zai" | "qwen" | "qwen-chat-template";
|
||||
/** OpenRouter-specific routing preferences. Only used when baseUrl points to OpenRouter. */
|
||||
openRouterRouting?: OpenRouterRouting;
|
||||
/** Vercel AI Gateway routing preferences. Only used when baseUrl points to Vercel AI Gateway. */
|
||||
vercelGatewayRouting?: VercelGatewayRouting;
|
||||
/** Whether z.ai supports top-level `tool_stream: true` for streaming tool call deltas. Default: false. */
|
||||
zaiToolStream?: boolean;
|
||||
/** Whether the provider supports the `strict` field in tool definitions. Default: true. */
|
||||
supportsStrictMode?: boolean;
|
||||
}
|
||||
|
||||
/** Compatibility settings for OpenAI Responses APIs. */
|
||||
export interface OpenAIResponsesCompat {
|
||||
// Reserved for future use
|
||||
}
|
||||
|
||||
/**
|
||||
* OpenRouter provider routing preferences.
|
||||
* Controls which upstream providers OpenRouter routes requests to.
|
||||
* Sent as the `provider` field in the OpenRouter API request body.
|
||||
* @see https://openrouter.ai/docs/guides/routing/provider-selection
|
||||
*/
|
||||
export interface OpenRouterRouting {
|
||||
/** Whether to allow backup providers to serve requests. Default: true. */
|
||||
allow_fallbacks?: boolean;
|
||||
/** Whether to filter providers to only those that support all parameters in the request. Default: false. */
|
||||
require_parameters?: boolean;
|
||||
/** Data collection setting. "allow" (default): allow providers that may store/train on data. "deny": only use providers that don't collect user data. */
|
||||
data_collection?: "deny" | "allow";
|
||||
/** Whether to restrict routing to only ZDR (Zero Data Retention) endpoints. */
|
||||
zdr?: boolean;
|
||||
/** Whether to restrict routing to only models that allow text distillation. */
|
||||
enforce_distillable_text?: boolean;
|
||||
/** An ordered list of provider names/slugs to try in sequence, falling back to the next if unavailable. */
|
||||
order?: string[];
|
||||
/** List of provider names/slugs to exclusively allow for this request. */
|
||||
only?: string[];
|
||||
/** List of provider names/slugs to skip for this request. */
|
||||
ignore?: string[];
|
||||
/** A list of quantization levels to filter providers by (e.g., ["fp16", "bf16", "fp8", "fp6", "int8", "int4", "fp4", "fp32"]). */
|
||||
quantizations?: string[];
|
||||
/** Sorting strategy. Can be a string (e.g., "price", "throughput", "latency") or an object with `by` and `partition`. */
|
||||
sort?:
|
||||
| string
|
||||
| {
|
||||
/** The sorting metric: "price", "throughput", "latency". */
|
||||
by?: string;
|
||||
/** Partitioning strategy: "model" (default) or "none". */
|
||||
partition?: string | null;
|
||||
};
|
||||
/** Maximum price per million tokens (USD). */
|
||||
max_price?: {
|
||||
/** Price per million prompt tokens. */
|
||||
prompt?: number | string;
|
||||
/** Price per million completion tokens. */
|
||||
completion?: number | string;
|
||||
/** Price per image. */
|
||||
image?: number | string;
|
||||
/** Price per audio unit. */
|
||||
audio?: number | string;
|
||||
/** Price per request. */
|
||||
request?: number | string;
|
||||
};
|
||||
/** Preferred minimum throughput (tokens/second). Can be a number (applies to p50) or an object with percentile-specific cutoffs. */
|
||||
preferred_min_throughput?:
|
||||
| number
|
||||
| {
|
||||
/** Minimum tokens/second at the 50th percentile. */
|
||||
p50?: number;
|
||||
/** Minimum tokens/second at the 75th percentile. */
|
||||
p75?: number;
|
||||
/** Minimum tokens/second at the 90th percentile. */
|
||||
p90?: number;
|
||||
/** Minimum tokens/second at the 99th percentile. */
|
||||
p99?: number;
|
||||
};
|
||||
/** Preferred maximum latency (seconds). Can be a number (applies to p50) or an object with percentile-specific cutoffs. */
|
||||
preferred_max_latency?:
|
||||
| number
|
||||
| {
|
||||
/** Maximum latency in seconds at the 50th percentile. */
|
||||
p50?: number;
|
||||
/** Maximum latency in seconds at the 75th percentile. */
|
||||
p75?: number;
|
||||
/** Maximum latency in seconds at the 90th percentile. */
|
||||
p90?: number;
|
||||
/** Maximum latency in seconds at the 99th percentile. */
|
||||
p99?: number;
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Vercel AI Gateway routing preferences.
|
||||
* Controls which upstream providers the gateway routes requests to.
|
||||
* @see https://vercel.com/docs/ai-gateway/models-and-providers/provider-options
|
||||
*/
|
||||
export interface VercelGatewayRouting {
|
||||
/** List of provider slugs to exclusively use for this request (e.g., ["bedrock", "anthropic"]). */
|
||||
only?: string[];
|
||||
/** List of provider slugs to try in order (e.g., ["anthropic", "openai"]). */
|
||||
order?: string[];
|
||||
}
|
||||
|
||||
// Model interface for the unified model system
|
||||
export interface Model<TApi extends Api> {
|
||||
id: string;
|
||||
name: string;
|
||||
api: TApi;
|
||||
provider: Provider;
|
||||
baseUrl: string;
|
||||
reasoning: boolean;
|
||||
input: ("text" | "image")[];
|
||||
cost: {
|
||||
input: number; // $/million tokens
|
||||
output: number; // $/million tokens
|
||||
cacheRead: number; // $/million tokens
|
||||
cacheWrite: number; // $/million tokens
|
||||
};
|
||||
contextWindow: number;
|
||||
maxTokens: number;
|
||||
headers?: Record<string, string>;
|
||||
/** Compatibility overrides for OpenAI-compatible APIs. If not set, auto-detected from baseUrl. */
|
||||
compat?: TApi extends "openai-completions"
|
||||
? OpenAICompletionsCompat
|
||||
: TApi extends "openai-responses"
|
||||
? OpenAIResponsesCompat
|
||||
: never;
|
||||
}
|
||||
@@ -1,87 +0,0 @@
|
||||
import type { AssistantMessage, AssistantMessageEvent } from "../types.js";
|
||||
|
||||
// Generic event stream class for async iteration
|
||||
export class EventStream<T, R = T> implements AsyncIterable<T> {
|
||||
private queue: T[] = [];
|
||||
private waiting: ((value: IteratorResult<T>) => void)[] = [];
|
||||
private done = false;
|
||||
private finalResultPromise: Promise<R>;
|
||||
private resolveFinalResult!: (result: R) => void;
|
||||
|
||||
constructor(
|
||||
private isComplete: (event: T) => boolean,
|
||||
private extractResult: (event: T) => R,
|
||||
) {
|
||||
this.finalResultPromise = new Promise((resolve) => {
|
||||
this.resolveFinalResult = resolve;
|
||||
});
|
||||
}
|
||||
|
||||
push(event: T): void {
|
||||
if (this.done) return;
|
||||
|
||||
if (this.isComplete(event)) {
|
||||
this.done = true;
|
||||
this.resolveFinalResult(this.extractResult(event));
|
||||
}
|
||||
|
||||
// Deliver to waiting consumer or queue it
|
||||
const waiter = this.waiting.shift();
|
||||
if (waiter) {
|
||||
waiter({ value: event, done: false });
|
||||
} else {
|
||||
this.queue.push(event);
|
||||
}
|
||||
}
|
||||
|
||||
end(result?: R): void {
|
||||
this.done = true;
|
||||
if (result !== undefined) {
|
||||
this.resolveFinalResult(result);
|
||||
}
|
||||
// Notify all waiting consumers that we're done
|
||||
while (this.waiting.length > 0) {
|
||||
const waiter = this.waiting.shift()!;
|
||||
waiter({ value: undefined as any, done: true });
|
||||
}
|
||||
}
|
||||
|
||||
async *[Symbol.asyncIterator](): AsyncIterator<T> {
|
||||
while (true) {
|
||||
if (this.queue.length > 0) {
|
||||
yield this.queue.shift()!;
|
||||
} else if (this.done) {
|
||||
return;
|
||||
} else {
|
||||
const result = await new Promise<IteratorResult<T>>((resolve) => this.waiting.push(resolve));
|
||||
if (result.done) return;
|
||||
yield result.value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result(): Promise<R> {
|
||||
return this.finalResultPromise;
|
||||
}
|
||||
}
|
||||
|
||||
export class AssistantMessageEventStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
|
||||
constructor() {
|
||||
super(
|
||||
(event) => event.type === "done" || event.type === "error",
|
||||
(event) => {
|
||||
if (event.type === "done") {
|
||||
return event.message;
|
||||
} else if (event.type === "error") {
|
||||
return event.error;
|
||||
}
|
||||
throw new Error("Unexpected event type for final result");
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/** Factory function for AssistantMessageEventStream (for use in extensions) */
|
||||
export function createAssistantMessageEventStream(): AssistantMessageEventStream {
|
||||
return new AssistantMessageEventStream();
|
||||
}
|
||||
@@ -1,13 +0,0 @@
|
||||
/** Fast deterministic hash to shorten long strings */
|
||||
export function shortHash(str: string): string {
|
||||
let h1 = 0xdeadbeef;
|
||||
let h2 = 0x41c6ce57;
|
||||
for (let i = 0; i < str.length; i++) {
|
||||
const ch = str.charCodeAt(i);
|
||||
h1 = Math.imul(h1 ^ ch, 2654435761);
|
||||
h2 = Math.imul(h2 ^ ch, 1597334677);
|
||||
}
|
||||
h1 = Math.imul(h1 ^ (h1 >>> 16), 2246822507) ^ Math.imul(h2 ^ (h2 >>> 13), 3266489909);
|
||||
h2 = Math.imul(h2 ^ (h2 >>> 16), 2246822507) ^ Math.imul(h1 ^ (h1 >>> 13), 3266489909);
|
||||
return (h2 >>> 0).toString(36) + (h1 >>> 0).toString(36);
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
export function headersToRecord(headers: Headers): Record<string, string> {
|
||||
const result: Record<string, string> = {};
|
||||
for (const [key, value] of headers.entries()) {
|
||||
result[key] = value;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
@@ -1,28 +0,0 @@
|
||||
import { parse as partialParse } from "partial-json";
|
||||
|
||||
/**
|
||||
* Attempts to parse potentially incomplete JSON during streaming.
|
||||
* Always returns a valid object, even if the JSON is incomplete.
|
||||
*
|
||||
* @param partialJson The partial JSON string from streaming
|
||||
* @returns Parsed object or empty object if parsing fails
|
||||
*/
|
||||
export function parseStreamingJson<T = any>(partialJson: string | undefined): T {
|
||||
if (!partialJson || partialJson.trim() === "") {
|
||||
return {} as T;
|
||||
}
|
||||
|
||||
// Try standard parsing first (fastest for complete JSON)
|
||||
try {
|
||||
return JSON.parse(partialJson) as T;
|
||||
} catch {
|
||||
// Try partial-json for incomplete JSON
|
||||
try {
|
||||
const result = partialParse(partialJson);
|
||||
return (result ?? {}) as T;
|
||||
} catch {
|
||||
// If all parsing fails, return empty object
|
||||
return {} as T;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,402 +0,0 @@
|
||||
/**
|
||||
* Anthropic OAuth flow (Claude Pro/Max)
|
||||
*
|
||||
* NOTE: This module uses Node.js http.createServer for the OAuth callback server.
|
||||
* It is only intended for CLI use, not browser environments.
|
||||
*/
|
||||
|
||||
import type { Server } from "node:http";
|
||||
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.js";
|
||||
import { generatePKCE } from "./pkce.js";
|
||||
import type { OAuthCredentials, OAuthLoginCallbacks, OAuthPrompt, OAuthProviderInterface } from "./types.js";
|
||||
|
||||
type CallbackServerInfo = {
|
||||
server: Server;
|
||||
redirectUri: string;
|
||||
cancelWait: () => void;
|
||||
waitForCode: () => Promise<{ code: string; state: string } | null>;
|
||||
};
|
||||
|
||||
type NodeApis = {
|
||||
createServer: typeof import("node:http").createServer;
|
||||
};
|
||||
|
||||
let nodeApis: NodeApis | null = null;
|
||||
let nodeApisPromise: Promise<NodeApis> | null = null;
|
||||
|
||||
const decode = (s: string) => atob(s);
|
||||
const CLIENT_ID = decode("OWQxYzI1MGEtZTYxYi00NGQ5LTg4ZWQtNTk0NGQxOTYyZjVl");
|
||||
const AUTHORIZE_URL = "https://claude.ai/oauth/authorize";
|
||||
const TOKEN_URL = "https://platform.claude.com/v1/oauth/token";
|
||||
const CALLBACK_HOST = "127.0.0.1";
|
||||
const CALLBACK_PORT = 53692;
|
||||
const CALLBACK_PATH = "/callback";
|
||||
const REDIRECT_URI = `http://localhost:${CALLBACK_PORT}${CALLBACK_PATH}`;
|
||||
const SCOPES =
|
||||
"org:create_api_key user:profile user:inference user:sessions:claude_code user:mcp_servers user:file_upload";
|
||||
async function getNodeApis(): Promise<NodeApis> {
|
||||
if (nodeApis) return nodeApis;
|
||||
if (!nodeApisPromise) {
|
||||
if (typeof process === "undefined" || (!process.versions?.node && !process.versions?.bun)) {
|
||||
throw new Error("Anthropic OAuth is only available in Node.js environments");
|
||||
}
|
||||
nodeApisPromise = import("node:http").then((httpModule) => ({
|
||||
createServer: httpModule.createServer,
|
||||
}));
|
||||
}
|
||||
nodeApis = await nodeApisPromise;
|
||||
return nodeApis;
|
||||
}
|
||||
|
||||
function parseAuthorizationInput(input: string): { code?: string; state?: string } {
|
||||
const value = input.trim();
|
||||
if (!value) return {};
|
||||
|
||||
try {
|
||||
const url = new URL(value);
|
||||
return {
|
||||
code: url.searchParams.get("code") ?? undefined,
|
||||
state: url.searchParams.get("state") ?? undefined,
|
||||
};
|
||||
} catch {
|
||||
// not a URL
|
||||
}
|
||||
|
||||
if (value.includes("#")) {
|
||||
const [code, state] = value.split("#", 2);
|
||||
return { code, state };
|
||||
}
|
||||
|
||||
if (value.includes("code=")) {
|
||||
const params = new URLSearchParams(value);
|
||||
return {
|
||||
code: params.get("code") ?? undefined,
|
||||
state: params.get("state") ?? undefined,
|
||||
};
|
||||
}
|
||||
|
||||
return { code: value };
|
||||
}
|
||||
|
||||
function formatErrorDetails(error: unknown): string {
|
||||
if (error instanceof Error) {
|
||||
const details: string[] = [`${error.name}: ${error.message}`];
|
||||
const errorWithCode = error as Error & { code?: string; errno?: number | string; cause?: unknown };
|
||||
if (errorWithCode.code) details.push(`code=${errorWithCode.code}`);
|
||||
if (typeof errorWithCode.errno !== "undefined") details.push(`errno=${String(errorWithCode.errno)}`);
|
||||
if (typeof error.cause !== "undefined") {
|
||||
details.push(`cause=${formatErrorDetails(error.cause)}`);
|
||||
}
|
||||
if (error.stack) {
|
||||
details.push(`stack=${error.stack}`);
|
||||
}
|
||||
return details.join("; ");
|
||||
}
|
||||
return String(error);
|
||||
}
|
||||
|
||||
async function startCallbackServer(expectedState: string): Promise<CallbackServerInfo> {
|
||||
const { createServer } = await getNodeApis();
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
let settleWait: ((value: { code: string; state: string } | null) => void) | undefined;
|
||||
const waitForCodePromise = new Promise<{ code: string; state: string } | null>((resolveWait) => {
|
||||
let settled = false;
|
||||
settleWait = (value) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
resolveWait(value);
|
||||
};
|
||||
});
|
||||
|
||||
const server = createServer((req, res) => {
|
||||
try {
|
||||
const url = new URL(req.url || "", "http://localhost");
|
||||
if (url.pathname !== CALLBACK_PATH) {
|
||||
res.writeHead(404, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Callback route not found."));
|
||||
return;
|
||||
}
|
||||
|
||||
const code = url.searchParams.get("code");
|
||||
const state = url.searchParams.get("state");
|
||||
const error = url.searchParams.get("error");
|
||||
|
||||
if (error) {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Anthropic authentication did not complete.", `Error: ${error}`));
|
||||
return;
|
||||
}
|
||||
|
||||
if (!code || !state) {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Missing code or state parameter."));
|
||||
return;
|
||||
}
|
||||
|
||||
if (state !== expectedState) {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("State mismatch."));
|
||||
return;
|
||||
}
|
||||
|
||||
res.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthSuccessHtml("Anthropic authentication completed. You can close this window."));
|
||||
settleWait?.({ code, state });
|
||||
} catch {
|
||||
res.writeHead(500, { "Content-Type": "text/plain; charset=utf-8" });
|
||||
res.end("Internal error");
|
||||
}
|
||||
});
|
||||
|
||||
server.on("error", (err) => {
|
||||
reject(err);
|
||||
});
|
||||
|
||||
server.listen(CALLBACK_PORT, CALLBACK_HOST, () => {
|
||||
resolve({
|
||||
server,
|
||||
redirectUri: REDIRECT_URI,
|
||||
cancelWait: () => {
|
||||
settleWait?.(null);
|
||||
},
|
||||
waitForCode: () => waitForCodePromise,
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async function postJson(url: string, body: Record<string, string | number>): Promise<string> {
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
Accept: "application/json",
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
signal: AbortSignal.timeout(30_000),
|
||||
});
|
||||
|
||||
const responseBody = await response.text();
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`HTTP request failed. status=${response.status}; url=${url}; body=${responseBody}`);
|
||||
}
|
||||
|
||||
return responseBody;
|
||||
}
|
||||
|
||||
async function exchangeAuthorizationCode(
|
||||
code: string,
|
||||
state: string,
|
||||
verifier: string,
|
||||
redirectUri: string,
|
||||
): Promise<OAuthCredentials> {
|
||||
let responseBody: string;
|
||||
try {
|
||||
responseBody = await postJson(TOKEN_URL, {
|
||||
grant_type: "authorization_code",
|
||||
client_id: CLIENT_ID,
|
||||
code,
|
||||
state,
|
||||
redirect_uri: redirectUri,
|
||||
code_verifier: verifier,
|
||||
});
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Token exchange request failed. url=${TOKEN_URL}; redirect_uri=${redirectUri}; response_type=authorization_code; details=${formatErrorDetails(error)}`,
|
||||
);
|
||||
}
|
||||
|
||||
let tokenData: { access_token: string; refresh_token: string; expires_in: number };
|
||||
try {
|
||||
tokenData = JSON.parse(responseBody) as { access_token: string; refresh_token: string; expires_in: number };
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Token exchange returned invalid JSON. url=${TOKEN_URL}; body=${responseBody}; details=${formatErrorDetails(error)}`,
|
||||
);
|
||||
}
|
||||
|
||||
return {
|
||||
refresh: tokenData.refresh_token,
|
||||
access: tokenData.access_token,
|
||||
expires: Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Login with Anthropic OAuth (authorization code + PKCE)
|
||||
*/
|
||||
export async function loginAnthropic(options: {
|
||||
onAuth: (info: { url: string; instructions?: string }) => void;
|
||||
onPrompt: (prompt: OAuthPrompt) => Promise<string>;
|
||||
onProgress?: (message: string) => void;
|
||||
onManualCodeInput?: () => Promise<string>;
|
||||
}): Promise<OAuthCredentials> {
|
||||
const { verifier, challenge } = await generatePKCE();
|
||||
const server = await startCallbackServer(verifier);
|
||||
|
||||
let code: string | undefined;
|
||||
let state: string | undefined;
|
||||
let redirectUriForExchange = REDIRECT_URI;
|
||||
|
||||
try {
|
||||
const authParams = new URLSearchParams({
|
||||
code: "true",
|
||||
client_id: CLIENT_ID,
|
||||
response_type: "code",
|
||||
redirect_uri: REDIRECT_URI,
|
||||
scope: SCOPES,
|
||||
code_challenge: challenge,
|
||||
code_challenge_method: "S256",
|
||||
state: verifier,
|
||||
});
|
||||
|
||||
options.onAuth({
|
||||
url: `${AUTHORIZE_URL}?${authParams.toString()}`,
|
||||
instructions:
|
||||
"Complete login in your browser. If the browser is on another machine, paste the final redirect URL here.",
|
||||
});
|
||||
|
||||
if (options.onManualCodeInput) {
|
||||
let manualInput: string | undefined;
|
||||
let manualError: Error | undefined;
|
||||
const manualPromise = options
|
||||
.onManualCodeInput()
|
||||
.then((input) => {
|
||||
manualInput = input;
|
||||
server.cancelWait();
|
||||
})
|
||||
.catch((err) => {
|
||||
manualError = err instanceof Error ? err : new Error(String(err));
|
||||
server.cancelWait();
|
||||
});
|
||||
|
||||
const result = await server.waitForCode();
|
||||
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
|
||||
if (result?.code) {
|
||||
code = result.code;
|
||||
state = result.state;
|
||||
redirectUriForExchange = REDIRECT_URI;
|
||||
} else if (manualInput) {
|
||||
const parsed = parseAuthorizationInput(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch");
|
||||
}
|
||||
code = parsed.code;
|
||||
state = parsed.state ?? verifier;
|
||||
}
|
||||
|
||||
if (!code) {
|
||||
await manualPromise;
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
if (manualInput) {
|
||||
const parsed = parseAuthorizationInput(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch");
|
||||
}
|
||||
code = parsed.code;
|
||||
state = parsed.state ?? verifier;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
const result = await server.waitForCode();
|
||||
if (result?.code) {
|
||||
code = result.code;
|
||||
state = result.state;
|
||||
redirectUriForExchange = REDIRECT_URI;
|
||||
}
|
||||
}
|
||||
|
||||
if (!code) {
|
||||
const input = await options.onPrompt({
|
||||
message: "Paste the authorization code or full redirect URL:",
|
||||
placeholder: REDIRECT_URI,
|
||||
});
|
||||
const parsed = parseAuthorizationInput(input);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch");
|
||||
}
|
||||
code = parsed.code;
|
||||
state = parsed.state ?? verifier;
|
||||
}
|
||||
|
||||
if (!code) {
|
||||
throw new Error("Missing authorization code");
|
||||
}
|
||||
|
||||
if (!state) {
|
||||
throw new Error("Missing OAuth state");
|
||||
}
|
||||
|
||||
options.onProgress?.("Exchanging authorization code for tokens...");
|
||||
return exchangeAuthorizationCode(code, state, verifier, redirectUriForExchange);
|
||||
} finally {
|
||||
server.server.close();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh Anthropic OAuth token
|
||||
*/
|
||||
export async function refreshAnthropicToken(refreshToken: string): Promise<OAuthCredentials> {
|
||||
let responseBody: string;
|
||||
try {
|
||||
responseBody = await postJson(TOKEN_URL, {
|
||||
grant_type: "refresh_token",
|
||||
client_id: CLIENT_ID,
|
||||
refresh_token: refreshToken,
|
||||
});
|
||||
} catch (error) {
|
||||
throw new Error(`Anthropic token refresh request failed. url=${TOKEN_URL}; details=${formatErrorDetails(error)}`);
|
||||
}
|
||||
|
||||
let data: { access_token: string; refresh_token: string; expires_in: number; scope?: string };
|
||||
try {
|
||||
data = JSON.parse(responseBody) as {
|
||||
access_token: string;
|
||||
refresh_token: string;
|
||||
expires_in: number;
|
||||
scope?: string;
|
||||
};
|
||||
} catch (error) {
|
||||
throw new Error(
|
||||
`Anthropic token refresh returned invalid JSON. url=${TOKEN_URL}; body=${responseBody}; details=${formatErrorDetails(error)}`,
|
||||
);
|
||||
}
|
||||
|
||||
return {
|
||||
refresh: data.refresh_token,
|
||||
access: data.access_token,
|
||||
expires: Date.now() + data.expires_in * 1000 - 5 * 60 * 1000,
|
||||
};
|
||||
}
|
||||
|
||||
export const anthropicOAuthProvider: OAuthProviderInterface = {
|
||||
id: "anthropic",
|
||||
name: "Anthropic (Claude Pro/Max)",
|
||||
usesCallbackServer: true,
|
||||
|
||||
async login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
|
||||
return loginAnthropic({
|
||||
onAuth: callbacks.onAuth,
|
||||
onPrompt: callbacks.onPrompt,
|
||||
onProgress: callbacks.onProgress,
|
||||
onManualCodeInput: callbacks.onManualCodeInput,
|
||||
});
|
||||
},
|
||||
|
||||
async refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
|
||||
return refreshAnthropicToken(credentials.refresh);
|
||||
},
|
||||
|
||||
getApiKey(credentials: OAuthCredentials): string {
|
||||
return credentials.access;
|
||||
},
|
||||
};
|
||||
@@ -1,396 +0,0 @@
|
||||
/**
|
||||
* GitHub Copilot OAuth flow
|
||||
*/
|
||||
|
||||
import { getModels } from "../../models.js";
|
||||
import type { Api, Model } from "../../types.js";
|
||||
import type { OAuthCredentials, OAuthLoginCallbacks, OAuthProviderInterface } from "./types.js";
|
||||
|
||||
type CopilotCredentials = OAuthCredentials & {
|
||||
enterpriseUrl?: string;
|
||||
};
|
||||
|
||||
const decode = (s: string) => atob(s);
|
||||
const CLIENT_ID = decode("SXYxLmI1MDdhMDhjODdlY2ZlOTg=");
|
||||
|
||||
const COPILOT_HEADERS = {
|
||||
"User-Agent": "GitHubCopilotChat/0.35.0",
|
||||
"Editor-Version": "vscode/1.107.0",
|
||||
"Editor-Plugin-Version": "copilot-chat/0.35.0",
|
||||
"Copilot-Integration-Id": "vscode-chat",
|
||||
} as const;
|
||||
|
||||
const INITIAL_POLL_INTERVAL_MULTIPLIER = 1.2;
|
||||
const SLOW_DOWN_POLL_INTERVAL_MULTIPLIER = 1.4;
|
||||
|
||||
type DeviceCodeResponse = {
|
||||
device_code: string;
|
||||
user_code: string;
|
||||
verification_uri: string;
|
||||
interval: number;
|
||||
expires_in: number;
|
||||
};
|
||||
|
||||
type DeviceTokenSuccessResponse = {
|
||||
access_token: string;
|
||||
token_type?: string;
|
||||
scope?: string;
|
||||
};
|
||||
|
||||
type DeviceTokenErrorResponse = {
|
||||
error: string;
|
||||
error_description?: string;
|
||||
interval?: number;
|
||||
};
|
||||
|
||||
export function normalizeDomain(input: string): string | null {
|
||||
const trimmed = input.trim();
|
||||
if (!trimmed) return null;
|
||||
try {
|
||||
const url = trimmed.includes("://") ? new URL(trimmed) : new URL(`https://${trimmed}`);
|
||||
return url.hostname;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function getUrls(domain: string): {
|
||||
deviceCodeUrl: string;
|
||||
accessTokenUrl: string;
|
||||
copilotTokenUrl: string;
|
||||
} {
|
||||
return {
|
||||
deviceCodeUrl: `https://${domain}/login/device/code`,
|
||||
accessTokenUrl: `https://${domain}/login/oauth/access_token`,
|
||||
copilotTokenUrl: `https://api.${domain}/copilot_internal/v2/token`,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse the proxy-ep from a Copilot token and convert to API base URL.
|
||||
* Token format: tid=...;exp=...;proxy-ep=proxy.individual.githubcopilot.com;...
|
||||
* Returns API URL like https://api.individual.githubcopilot.com
|
||||
*/
|
||||
function getBaseUrlFromToken(token: string): string | null {
|
||||
const match = token.match(/proxy-ep=([^;]+)/);
|
||||
if (!match) return null;
|
||||
const proxyHost = match[1];
|
||||
// Convert proxy.xxx to api.xxx
|
||||
const apiHost = proxyHost.replace(/^proxy\./, "api.");
|
||||
return `https://${apiHost}`;
|
||||
}
|
||||
|
||||
export function getGitHubCopilotBaseUrl(token?: string, enterpriseDomain?: string): string {
|
||||
// If we have a token, extract the base URL from proxy-ep
|
||||
if (token) {
|
||||
const urlFromToken = getBaseUrlFromToken(token);
|
||||
if (urlFromToken) return urlFromToken;
|
||||
}
|
||||
// Fallback for enterprise or if token parsing fails
|
||||
if (enterpriseDomain) return `https://copilot-api.${enterpriseDomain}`;
|
||||
return "https://api.individual.githubcopilot.com";
|
||||
}
|
||||
|
||||
async function fetchJson(url: string, init: RequestInit): Promise<unknown> {
|
||||
const response = await fetch(url, init);
|
||||
if (!response.ok) {
|
||||
const text = await response.text();
|
||||
throw new Error(`${response.status} ${response.statusText}: ${text}`);
|
||||
}
|
||||
return response.json();
|
||||
}
|
||||
|
||||
async function startDeviceFlow(domain: string): Promise<DeviceCodeResponse> {
|
||||
const urls = getUrls(domain);
|
||||
const data = await fetchJson(urls.deviceCodeUrl, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
Accept: "application/json",
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
"User-Agent": "GitHubCopilotChat/0.35.0",
|
||||
},
|
||||
body: new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
scope: "read:user",
|
||||
}),
|
||||
});
|
||||
|
||||
if (!data || typeof data !== "object") {
|
||||
throw new Error("Invalid device code response");
|
||||
}
|
||||
|
||||
const deviceCode = (data as Record<string, unknown>).device_code;
|
||||
const userCode = (data as Record<string, unknown>).user_code;
|
||||
const verificationUri = (data as Record<string, unknown>).verification_uri;
|
||||
const interval = (data as Record<string, unknown>).interval;
|
||||
const expiresIn = (data as Record<string, unknown>).expires_in;
|
||||
|
||||
if (
|
||||
typeof deviceCode !== "string" ||
|
||||
typeof userCode !== "string" ||
|
||||
typeof verificationUri !== "string" ||
|
||||
typeof interval !== "number" ||
|
||||
typeof expiresIn !== "number"
|
||||
) {
|
||||
throw new Error("Invalid device code response fields");
|
||||
}
|
||||
|
||||
return {
|
||||
device_code: deviceCode,
|
||||
user_code: userCode,
|
||||
verification_uri: verificationUri,
|
||||
interval,
|
||||
expires_in: expiresIn,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Sleep that can be interrupted by an AbortSignal
|
||||
*/
|
||||
function abortableSleep(ms: number, signal?: AbortSignal): Promise<void> {
|
||||
return new Promise((resolve, reject) => {
|
||||
if (signal?.aborted) {
|
||||
reject(new Error("Login cancelled"));
|
||||
return;
|
||||
}
|
||||
|
||||
const timeout = setTimeout(resolve, ms);
|
||||
|
||||
signal?.addEventListener(
|
||||
"abort",
|
||||
() => {
|
||||
clearTimeout(timeout);
|
||||
reject(new Error("Login cancelled"));
|
||||
},
|
||||
{ once: true },
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
async function pollForGitHubAccessToken(
|
||||
domain: string,
|
||||
deviceCode: string,
|
||||
intervalSeconds: number,
|
||||
expiresIn: number,
|
||||
signal?: AbortSignal,
|
||||
) {
|
||||
const urls = getUrls(domain);
|
||||
const deadline = Date.now() + expiresIn * 1000;
|
||||
let intervalMs = Math.max(1000, Math.floor(intervalSeconds * 1000));
|
||||
let intervalMultiplier = INITIAL_POLL_INTERVAL_MULTIPLIER;
|
||||
let slowDownResponses = 0;
|
||||
|
||||
while (Date.now() < deadline) {
|
||||
if (signal?.aborted) {
|
||||
throw new Error("Login cancelled");
|
||||
}
|
||||
|
||||
const remainingMs = deadline - Date.now();
|
||||
const waitMs = Math.min(Math.ceil(intervalMs * intervalMultiplier), remainingMs);
|
||||
await abortableSleep(waitMs, signal);
|
||||
|
||||
const raw = await fetchJson(urls.accessTokenUrl, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
Accept: "application/json",
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
"User-Agent": "GitHubCopilotChat/0.35.0",
|
||||
},
|
||||
body: new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
device_code: deviceCode,
|
||||
grant_type: "urn:ietf:params:oauth:grant-type:device_code",
|
||||
}),
|
||||
});
|
||||
|
||||
if (raw && typeof raw === "object" && typeof (raw as DeviceTokenSuccessResponse).access_token === "string") {
|
||||
return (raw as DeviceTokenSuccessResponse).access_token;
|
||||
}
|
||||
|
||||
if (raw && typeof raw === "object" && typeof (raw as DeviceTokenErrorResponse).error === "string") {
|
||||
const { error, error_description: description, interval } = raw as DeviceTokenErrorResponse;
|
||||
if (error === "authorization_pending") {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (error === "slow_down") {
|
||||
slowDownResponses += 1;
|
||||
intervalMs =
|
||||
typeof interval === "number" && interval > 0 ? interval * 1000 : Math.max(1000, intervalMs + 5000);
|
||||
intervalMultiplier = SLOW_DOWN_POLL_INTERVAL_MULTIPLIER;
|
||||
continue;
|
||||
}
|
||||
|
||||
const descriptionSuffix = description ? `: ${description}` : "";
|
||||
throw new Error(`Device flow failed: ${error}${descriptionSuffix}`);
|
||||
}
|
||||
}
|
||||
|
||||
if (slowDownResponses > 0) {
|
||||
throw new Error(
|
||||
"Device flow timed out after one or more slow_down responses. This is often caused by clock drift in WSL or VM environments. Please sync or restart the VM clock and try again.",
|
||||
);
|
||||
}
|
||||
|
||||
throw new Error("Device flow timed out");
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh GitHub Copilot token
|
||||
*/
|
||||
export async function refreshGitHubCopilotToken(
|
||||
refreshToken: string,
|
||||
enterpriseDomain?: string,
|
||||
): Promise<OAuthCredentials> {
|
||||
const domain = enterpriseDomain || "github.com";
|
||||
const urls = getUrls(domain);
|
||||
|
||||
const raw = await fetchJson(urls.copilotTokenUrl, {
|
||||
headers: {
|
||||
Accept: "application/json",
|
||||
Authorization: `Bearer ${refreshToken}`,
|
||||
...COPILOT_HEADERS,
|
||||
},
|
||||
});
|
||||
|
||||
if (!raw || typeof raw !== "object") {
|
||||
throw new Error("Invalid Copilot token response");
|
||||
}
|
||||
|
||||
const token = (raw as Record<string, unknown>).token;
|
||||
const expiresAt = (raw as Record<string, unknown>).expires_at;
|
||||
|
||||
if (typeof token !== "string" || typeof expiresAt !== "number") {
|
||||
throw new Error("Invalid Copilot token response fields");
|
||||
}
|
||||
|
||||
return {
|
||||
refresh: refreshToken,
|
||||
access: token,
|
||||
expires: expiresAt * 1000 - 5 * 60 * 1000,
|
||||
enterpriseUrl: enterpriseDomain,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Enable a model for the user's GitHub Copilot account.
|
||||
* This is required for some models (like Claude, Grok) before they can be used.
|
||||
*/
|
||||
async function enableGitHubCopilotModel(token: string, modelId: string, enterpriseDomain?: string): Promise<boolean> {
|
||||
const baseUrl = getGitHubCopilotBaseUrl(token, enterpriseDomain);
|
||||
const url = `${baseUrl}/models/${modelId}/policy`;
|
||||
|
||||
try {
|
||||
const response = await fetch(url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
Authorization: `Bearer ${token}`,
|
||||
...COPILOT_HEADERS,
|
||||
"openai-intent": "chat-policy",
|
||||
"x-interaction-type": "chat-policy",
|
||||
},
|
||||
body: JSON.stringify({ state: "enabled" }),
|
||||
});
|
||||
return response.ok;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Enable all known GitHub Copilot models that may require policy acceptance.
|
||||
* Called after successful login to ensure all models are available.
|
||||
*/
|
||||
async function enableAllGitHubCopilotModels(
|
||||
token: string,
|
||||
enterpriseDomain?: string,
|
||||
onProgress?: (model: string, success: boolean) => void,
|
||||
): Promise<void> {
|
||||
const models = getModels("github-copilot");
|
||||
await Promise.all(
|
||||
models.map(async (model) => {
|
||||
const success = await enableGitHubCopilotModel(token, model.id, enterpriseDomain);
|
||||
onProgress?.(model.id, success);
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Login with GitHub Copilot OAuth (device code flow)
|
||||
*
|
||||
* @param options.onAuth - Callback with URL and optional instructions (user code)
|
||||
* @param options.onPrompt - Callback to prompt user for input
|
||||
* @param options.onProgress - Optional progress callback
|
||||
* @param options.signal - Optional AbortSignal for cancellation
|
||||
*/
|
||||
export async function loginGitHubCopilot(options: {
|
||||
onAuth: (url: string, instructions?: string) => void;
|
||||
onPrompt: (prompt: { message: string; placeholder?: string; allowEmpty?: boolean }) => Promise<string>;
|
||||
onProgress?: (message: string) => void;
|
||||
signal?: AbortSignal;
|
||||
}): Promise<OAuthCredentials> {
|
||||
const input = await options.onPrompt({
|
||||
message: "GitHub Enterprise URL/domain (blank for github.com)",
|
||||
placeholder: "company.ghe.com",
|
||||
allowEmpty: true,
|
||||
});
|
||||
|
||||
if (options.signal?.aborted) {
|
||||
throw new Error("Login cancelled");
|
||||
}
|
||||
|
||||
const trimmed = input.trim();
|
||||
const enterpriseDomain = normalizeDomain(input);
|
||||
if (trimmed && !enterpriseDomain) {
|
||||
throw new Error("Invalid GitHub Enterprise URL/domain");
|
||||
}
|
||||
const domain = enterpriseDomain || "github.com";
|
||||
|
||||
const device = await startDeviceFlow(domain);
|
||||
options.onAuth(device.verification_uri, `Enter code: ${device.user_code}`);
|
||||
|
||||
const githubAccessToken = await pollForGitHubAccessToken(
|
||||
domain,
|
||||
device.device_code,
|
||||
device.interval,
|
||||
device.expires_in,
|
||||
options.signal,
|
||||
);
|
||||
const credentials = await refreshGitHubCopilotToken(githubAccessToken, enterpriseDomain ?? undefined);
|
||||
|
||||
// Enable all models after successful login
|
||||
options.onProgress?.("Enabling models...");
|
||||
await enableAllGitHubCopilotModels(credentials.access, enterpriseDomain ?? undefined);
|
||||
return credentials;
|
||||
}
|
||||
|
||||
export const githubCopilotOAuthProvider: OAuthProviderInterface = {
|
||||
id: "github-copilot",
|
||||
name: "GitHub Copilot",
|
||||
|
||||
async login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
|
||||
return loginGitHubCopilot({
|
||||
onAuth: (url, instructions) => callbacks.onAuth({ url, instructions }),
|
||||
onPrompt: callbacks.onPrompt,
|
||||
onProgress: callbacks.onProgress,
|
||||
signal: callbacks.signal,
|
||||
});
|
||||
},
|
||||
|
||||
async refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
|
||||
const creds = credentials as CopilotCredentials;
|
||||
return refreshGitHubCopilotToken(creds.refresh, creds.enterpriseUrl);
|
||||
},
|
||||
|
||||
getApiKey(credentials: OAuthCredentials): string {
|
||||
return credentials.access;
|
||||
},
|
||||
|
||||
modifyModels(models: Model<Api>[], credentials: OAuthCredentials): Model<Api>[] {
|
||||
const creds = credentials as CopilotCredentials;
|
||||
const domain = creds.enterpriseUrl ? (normalizeDomain(creds.enterpriseUrl) ?? undefined) : undefined;
|
||||
const baseUrl = getGitHubCopilotBaseUrl(creds.access, domain);
|
||||
return models.map((m) => (m.provider === "github-copilot" ? { ...m, baseUrl } : m));
|
||||
},
|
||||
};
|
||||
@@ -1,453 +0,0 @@
|
||||
/**
|
||||
* Antigravity OAuth flow (Gemini 3, Claude, GPT-OSS via Google Cloud)
|
||||
* Uses different OAuth credentials than google-gemini-cli for access to additional models.
|
||||
*
|
||||
* NOTE: This module uses Node.js http.createServer for the OAuth callback.
|
||||
* It is only intended for CLI use, not browser environments.
|
||||
*/
|
||||
|
||||
import type { Server } from "node:http";
|
||||
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.js";
|
||||
import { generatePKCE } from "./pkce.js";
|
||||
import type { OAuthCredentials, OAuthLoginCallbacks, OAuthProviderInterface } from "./types.js";
|
||||
|
||||
type AntigravityCredentials = OAuthCredentials & {
|
||||
projectId: string;
|
||||
};
|
||||
|
||||
let _createServer: typeof import("node:http").createServer | null = null;
|
||||
let _httpImportPromise: Promise<void> | null = null;
|
||||
if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) {
|
||||
_httpImportPromise = import("node:http").then((m) => {
|
||||
_createServer = m.createServer;
|
||||
});
|
||||
}
|
||||
|
||||
// Antigravity OAuth credentials (different from Gemini CLI)
|
||||
const decode = (s: string) => atob(s);
|
||||
const CLIENT_ID = decode(
|
||||
"MTA3MTAwNjA2MDU5MS10bWhzc2luMmgyMWxjcmUyMzV2dG9sb2poNGc0MDNlcC5hcHBzLmdvb2dsZXVzZXJjb250ZW50LmNvbQ==",
|
||||
);
|
||||
const CLIENT_SECRET = decode("R09DU1BYLUs1OEZXUjQ4NkxkTEoxbUxCOHNYQzR6NnFEQWY=");
|
||||
const REDIRECT_URI = "http://localhost:51121/oauth-callback";
|
||||
|
||||
// Antigravity requires additional scopes
|
||||
const SCOPES = [
|
||||
"https://www.googleapis.com/auth/cloud-platform",
|
||||
"https://www.googleapis.com/auth/userinfo.email",
|
||||
"https://www.googleapis.com/auth/userinfo.profile",
|
||||
"https://www.googleapis.com/auth/cclog",
|
||||
"https://www.googleapis.com/auth/experimentsandconfigs",
|
||||
];
|
||||
|
||||
const AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth";
|
||||
const TOKEN_URL = "https://oauth2.googleapis.com/token";
|
||||
|
||||
// Fallback project ID when discovery fails
|
||||
const DEFAULT_PROJECT_ID = "rising-fact-p41fc";
|
||||
|
||||
type CallbackServerInfo = {
|
||||
server: Server;
|
||||
cancelWait: () => void;
|
||||
waitForCode: () => Promise<{ code: string; state: string } | null>;
|
||||
};
|
||||
|
||||
/**
|
||||
* Start a local HTTP server to receive the OAuth callback
|
||||
*/
|
||||
async function getNodeCreateServer(): Promise<typeof import("node:http").createServer> {
|
||||
if (_createServer) return _createServer;
|
||||
if (_httpImportPromise) {
|
||||
await _httpImportPromise;
|
||||
}
|
||||
if (_createServer) return _createServer;
|
||||
throw new Error("Antigravity OAuth is only available in Node.js environments");
|
||||
}
|
||||
|
||||
async function startCallbackServer(): Promise<CallbackServerInfo> {
|
||||
const createServer = await getNodeCreateServer();
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
let settleWait: ((value: { code: string; state: string } | null) => void) | undefined;
|
||||
const waitForCodePromise = new Promise<{ code: string; state: string } | null>((resolveWait) => {
|
||||
let settled = false;
|
||||
settleWait = (value) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
resolveWait(value);
|
||||
};
|
||||
});
|
||||
|
||||
const server = createServer((req, res) => {
|
||||
const url = new URL(req.url || "", `http://localhost:51121`);
|
||||
|
||||
if (url.pathname === "/oauth-callback") {
|
||||
const code = url.searchParams.get("code");
|
||||
const state = url.searchParams.get("state");
|
||||
const error = url.searchParams.get("error");
|
||||
|
||||
if (error) {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Google authentication did not complete.", `Error: ${error}`));
|
||||
return;
|
||||
}
|
||||
|
||||
if (code && state) {
|
||||
res.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthSuccessHtml("Google authentication completed. You can close this window."));
|
||||
settleWait?.({ code, state });
|
||||
} else {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Missing code or state parameter."));
|
||||
}
|
||||
} else {
|
||||
res.writeHead(404, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Callback route not found."));
|
||||
}
|
||||
});
|
||||
|
||||
server.on("error", (err) => {
|
||||
reject(err);
|
||||
});
|
||||
|
||||
server.listen(51121, "127.0.0.1", () => {
|
||||
resolve({
|
||||
server,
|
||||
cancelWait: () => {
|
||||
settleWait?.(null);
|
||||
},
|
||||
waitForCode: () => waitForCodePromise,
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse redirect URL to extract code and state
|
||||
*/
|
||||
function parseRedirectUrl(input: string): { code?: string; state?: string } {
|
||||
const value = input.trim();
|
||||
if (!value) return {};
|
||||
|
||||
try {
|
||||
const url = new URL(value);
|
||||
return {
|
||||
code: url.searchParams.get("code") ?? undefined,
|
||||
state: url.searchParams.get("state") ?? undefined,
|
||||
};
|
||||
} catch {
|
||||
// Not a URL, return empty
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
interface LoadCodeAssistPayload {
|
||||
cloudaicompanionProject?: string | { id?: string };
|
||||
currentTier?: { id?: string };
|
||||
allowedTiers?: Array<{ id?: string; isDefault?: boolean }>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Discover or provision a project for the user
|
||||
*/
|
||||
async function discoverProject(accessToken: string, onProgress?: (message: string) => void): Promise<string> {
|
||||
const headers = {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": "google-api-nodejs-client/9.15.1",
|
||||
"X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
|
||||
"Client-Metadata": JSON.stringify({
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
}),
|
||||
};
|
||||
|
||||
// Try endpoints in order: prod first, then sandbox
|
||||
const endpoints = ["https://cloudcode-pa.googleapis.com", "https://daily-cloudcode-pa.sandbox.googleapis.com"];
|
||||
|
||||
onProgress?.("Checking for existing project...");
|
||||
|
||||
for (const endpoint of endpoints) {
|
||||
try {
|
||||
const loadResponse = await fetch(`${endpoint}/v1internal:loadCodeAssist`, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify({
|
||||
metadata: {
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
},
|
||||
}),
|
||||
});
|
||||
|
||||
if (loadResponse.ok) {
|
||||
const data = (await loadResponse.json()) as LoadCodeAssistPayload;
|
||||
|
||||
// Handle both string and object formats
|
||||
if (typeof data.cloudaicompanionProject === "string" && data.cloudaicompanionProject) {
|
||||
return data.cloudaicompanionProject;
|
||||
}
|
||||
if (
|
||||
data.cloudaicompanionProject &&
|
||||
typeof data.cloudaicompanionProject === "object" &&
|
||||
data.cloudaicompanionProject.id
|
||||
) {
|
||||
return data.cloudaicompanionProject.id;
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
// Try next endpoint
|
||||
}
|
||||
}
|
||||
|
||||
// Use fallback project ID
|
||||
onProgress?.("Using default project...");
|
||||
return DEFAULT_PROJECT_ID;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get user email from the access token
|
||||
*/
|
||||
async function getUserEmail(accessToken: string): Promise<string | undefined> {
|
||||
try {
|
||||
const response = await fetch("https://www.googleapis.com/oauth2/v1/userinfo?alt=json", {
|
||||
headers: {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
},
|
||||
});
|
||||
|
||||
if (response.ok) {
|
||||
const data = (await response.json()) as { email?: string };
|
||||
return data.email;
|
||||
}
|
||||
} catch {
|
||||
// Ignore errors, email is optional
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh Antigravity token
|
||||
*/
|
||||
export async function refreshAntigravityToken(refreshToken: string, projectId: string): Promise<OAuthCredentials> {
|
||||
const response = await fetch(TOKEN_URL, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/x-www-form-urlencoded" },
|
||||
body: new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
client_secret: CLIENT_SECRET,
|
||||
refresh_token: refreshToken,
|
||||
grant_type: "refresh_token",
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.text();
|
||||
throw new Error(`Antigravity token refresh failed: ${error}`);
|
||||
}
|
||||
|
||||
const data = (await response.json()) as {
|
||||
access_token: string;
|
||||
expires_in: number;
|
||||
refresh_token?: string;
|
||||
};
|
||||
|
||||
return {
|
||||
refresh: data.refresh_token || refreshToken,
|
||||
access: data.access_token,
|
||||
expires: Date.now() + data.expires_in * 1000 - 5 * 60 * 1000,
|
||||
projectId,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Login with Antigravity OAuth
|
||||
*
|
||||
* @param onAuth - Callback with URL and optional instructions
|
||||
* @param onProgress - Optional progress callback
|
||||
* @param onManualCodeInput - Optional promise that resolves with user-pasted redirect URL.
|
||||
* Races with browser callback - whichever completes first wins.
|
||||
*/
|
||||
export async function loginAntigravity(
|
||||
onAuth: (info: { url: string; instructions?: string }) => void,
|
||||
onProgress?: (message: string) => void,
|
||||
onManualCodeInput?: () => Promise<string>,
|
||||
): Promise<OAuthCredentials> {
|
||||
const { verifier, challenge } = await generatePKCE();
|
||||
|
||||
// Start local server for callback
|
||||
onProgress?.("Starting local server for OAuth callback...");
|
||||
const server = await startCallbackServer();
|
||||
|
||||
let code: string | undefined;
|
||||
|
||||
try {
|
||||
// Build authorization URL
|
||||
const authParams = new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
response_type: "code",
|
||||
redirect_uri: REDIRECT_URI,
|
||||
scope: SCOPES.join(" "),
|
||||
code_challenge: challenge,
|
||||
code_challenge_method: "S256",
|
||||
state: verifier,
|
||||
access_type: "offline",
|
||||
prompt: "consent",
|
||||
});
|
||||
|
||||
const authUrl = `${AUTH_URL}?${authParams.toString()}`;
|
||||
|
||||
// Notify caller with URL to open
|
||||
onAuth({
|
||||
url: authUrl,
|
||||
instructions: "Complete the sign-in in your browser.",
|
||||
});
|
||||
|
||||
// Wait for the callback, racing with manual input if provided
|
||||
onProgress?.("Waiting for OAuth callback...");
|
||||
|
||||
if (onManualCodeInput) {
|
||||
// Race between browser callback and manual input
|
||||
let manualInput: string | undefined;
|
||||
let manualError: Error | undefined;
|
||||
const manualPromise = onManualCodeInput()
|
||||
.then((input) => {
|
||||
manualInput = input;
|
||||
server.cancelWait();
|
||||
})
|
||||
.catch((err) => {
|
||||
manualError = err instanceof Error ? err : new Error(String(err));
|
||||
server.cancelWait();
|
||||
});
|
||||
|
||||
const result = await server.waitForCode();
|
||||
|
||||
// If manual input was cancelled, throw that error
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
|
||||
if (result?.code) {
|
||||
// Browser callback won - verify state
|
||||
if (result.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = result.code;
|
||||
} else if (manualInput) {
|
||||
// Manual input won
|
||||
const parsed = parseRedirectUrl(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
|
||||
// If still no code, wait for manual promise and try that
|
||||
if (!code) {
|
||||
await manualPromise;
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
if (manualInput) {
|
||||
const parsed = parseRedirectUrl(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Original flow: just wait for callback
|
||||
const result = await server.waitForCode();
|
||||
if (result?.code) {
|
||||
if (result.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = result.code;
|
||||
}
|
||||
}
|
||||
|
||||
if (!code) {
|
||||
throw new Error("No authorization code received");
|
||||
}
|
||||
|
||||
// Exchange code for tokens
|
||||
onProgress?.("Exchanging authorization code for tokens...");
|
||||
const tokenResponse = await fetch(TOKEN_URL, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
},
|
||||
body: new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
client_secret: CLIENT_SECRET,
|
||||
code,
|
||||
grant_type: "authorization_code",
|
||||
redirect_uri: REDIRECT_URI,
|
||||
code_verifier: verifier,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!tokenResponse.ok) {
|
||||
const error = await tokenResponse.text();
|
||||
throw new Error(`Token exchange failed: ${error}`);
|
||||
}
|
||||
|
||||
const tokenData = (await tokenResponse.json()) as {
|
||||
access_token: string;
|
||||
refresh_token: string;
|
||||
expires_in: number;
|
||||
};
|
||||
|
||||
if (!tokenData.refresh_token) {
|
||||
throw new Error("No refresh token received. Please try again.");
|
||||
}
|
||||
|
||||
// Get user email
|
||||
onProgress?.("Getting user info...");
|
||||
const email = await getUserEmail(tokenData.access_token);
|
||||
|
||||
// Discover project
|
||||
const projectId = await discoverProject(tokenData.access_token, onProgress);
|
||||
|
||||
// Calculate expiry time (current time + expires_in seconds - 5 min buffer)
|
||||
const expiresAt = Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000;
|
||||
|
||||
const credentials: OAuthCredentials = {
|
||||
refresh: tokenData.refresh_token,
|
||||
access: tokenData.access_token,
|
||||
expires: expiresAt,
|
||||
projectId,
|
||||
email,
|
||||
};
|
||||
|
||||
return credentials;
|
||||
} finally {
|
||||
server.server.close();
|
||||
}
|
||||
}
|
||||
|
||||
export const antigravityOAuthProvider: OAuthProviderInterface = {
|
||||
id: "google-antigravity",
|
||||
name: "Antigravity (Gemini 3, Claude, GPT-OSS)",
|
||||
usesCallbackServer: true,
|
||||
|
||||
async login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
|
||||
return loginAntigravity(callbacks.onAuth, callbacks.onProgress, callbacks.onManualCodeInput);
|
||||
},
|
||||
|
||||
async refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
|
||||
const creds = credentials as AntigravityCredentials;
|
||||
if (!creds.projectId) {
|
||||
throw new Error("Antigravity credentials missing projectId");
|
||||
}
|
||||
return refreshAntigravityToken(creds.refresh, creds.projectId);
|
||||
},
|
||||
|
||||
getApiKey(credentials: OAuthCredentials): string {
|
||||
const creds = credentials as AntigravityCredentials;
|
||||
return JSON.stringify({ token: creds.access, projectId: creds.projectId });
|
||||
},
|
||||
};
|
||||
@@ -1,595 +0,0 @@
|
||||
/**
|
||||
* Gemini CLI OAuth flow (Google Cloud Code Assist)
|
||||
* Standard Gemini models only (gemini-2.0-flash, gemini-2.5-*)
|
||||
*
|
||||
* NOTE: This module uses Node.js http.createServer for the OAuth callback.
|
||||
* It is only intended for CLI use, not browser environments.
|
||||
*/
|
||||
|
||||
import type { Server } from "node:http";
|
||||
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.js";
|
||||
import { generatePKCE } from "./pkce.js";
|
||||
import type { OAuthCredentials, OAuthLoginCallbacks, OAuthProviderInterface } from "./types.js";
|
||||
|
||||
type GeminiCredentials = OAuthCredentials & {
|
||||
projectId: string;
|
||||
};
|
||||
|
||||
let _createServer: typeof import("node:http").createServer | null = null;
|
||||
let _httpImportPromise: Promise<void> | null = null;
|
||||
if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) {
|
||||
_httpImportPromise = import("node:http").then((m) => {
|
||||
_createServer = m.createServer;
|
||||
});
|
||||
}
|
||||
|
||||
const decode = (s: string) => atob(s);
|
||||
const CLIENT_ID = decode(
|
||||
"NjgxMjU1ODA5Mzk1LW9vOGZ0Mm9wcmRybnA5ZTNhcWY2YXYzaG1kaWIxMzVqLmFwcHMuZ29vZ2xldXNlcmNvbnRlbnQuY29t",
|
||||
);
|
||||
const CLIENT_SECRET = decode("R09DU1BYLTR1SGdNUG0tMW83U2stZ2VWNkN1NWNsWEZzeGw=");
|
||||
const REDIRECT_URI = "http://localhost:8085/oauth2callback";
|
||||
const SCOPES = [
|
||||
"https://www.googleapis.com/auth/cloud-platform",
|
||||
"https://www.googleapis.com/auth/userinfo.email",
|
||||
"https://www.googleapis.com/auth/userinfo.profile",
|
||||
];
|
||||
const AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth";
|
||||
const TOKEN_URL = "https://oauth2.googleapis.com/token";
|
||||
const CODE_ASSIST_ENDPOINT = "https://cloudcode-pa.googleapis.com";
|
||||
|
||||
type CallbackServerInfo = {
|
||||
server: Server;
|
||||
cancelWait: () => void;
|
||||
waitForCode: () => Promise<{ code: string; state: string } | null>;
|
||||
};
|
||||
|
||||
/**
|
||||
* Start a local HTTP server to receive the OAuth callback
|
||||
*/
|
||||
async function getNodeCreateServer(): Promise<typeof import("node:http").createServer> {
|
||||
if (_createServer) return _createServer;
|
||||
if (_httpImportPromise) {
|
||||
await _httpImportPromise;
|
||||
}
|
||||
if (_createServer) return _createServer;
|
||||
throw new Error("Gemini CLI OAuth is only available in Node.js environments");
|
||||
}
|
||||
|
||||
async function startCallbackServer(): Promise<CallbackServerInfo> {
|
||||
const createServer = await getNodeCreateServer();
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
let settleWait: ((value: { code: string; state: string } | null) => void) | undefined;
|
||||
const waitForCodePromise = new Promise<{ code: string; state: string } | null>((resolveWait) => {
|
||||
let settled = false;
|
||||
settleWait = (value) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
resolveWait(value);
|
||||
};
|
||||
});
|
||||
|
||||
const server = createServer((req, res) => {
|
||||
const url = new URL(req.url || "", `http://localhost:8085`);
|
||||
|
||||
if (url.pathname === "/oauth2callback") {
|
||||
const code = url.searchParams.get("code");
|
||||
const state = url.searchParams.get("state");
|
||||
const error = url.searchParams.get("error");
|
||||
|
||||
if (error) {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Google authentication did not complete.", `Error: ${error}`));
|
||||
return;
|
||||
}
|
||||
|
||||
if (code && state) {
|
||||
res.writeHead(200, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthSuccessHtml("Google authentication completed. You can close this window."));
|
||||
settleWait?.({ code, state });
|
||||
} else {
|
||||
res.writeHead(400, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Missing code or state parameter."));
|
||||
}
|
||||
} else {
|
||||
res.writeHead(404, { "Content-Type": "text/html; charset=utf-8" });
|
||||
res.end(oauthErrorHtml("Callback route not found."));
|
||||
}
|
||||
});
|
||||
|
||||
server.on("error", (err) => {
|
||||
reject(err);
|
||||
});
|
||||
|
||||
server.listen(8085, "127.0.0.1", () => {
|
||||
resolve({
|
||||
server,
|
||||
cancelWait: () => {
|
||||
settleWait?.(null);
|
||||
},
|
||||
waitForCode: () => waitForCodePromise,
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse redirect URL to extract code and state
|
||||
*/
|
||||
function parseRedirectUrl(input: string): { code?: string; state?: string } {
|
||||
const value = input.trim();
|
||||
if (!value) return {};
|
||||
|
||||
try {
|
||||
const url = new URL(value);
|
||||
return {
|
||||
code: url.searchParams.get("code") ?? undefined,
|
||||
state: url.searchParams.get("state") ?? undefined,
|
||||
};
|
||||
} catch {
|
||||
// Not a URL, return empty
|
||||
return {};
|
||||
}
|
||||
}
|
||||
|
||||
interface LoadCodeAssistPayload {
|
||||
cloudaicompanionProject?: string;
|
||||
currentTier?: { id?: string };
|
||||
allowedTiers?: Array<{ id?: string; isDefault?: boolean }>;
|
||||
}
|
||||
|
||||
/**
|
||||
* Long-running operation response from onboardUser
|
||||
*/
|
||||
interface LongRunningOperationResponse {
|
||||
name?: string;
|
||||
done?: boolean;
|
||||
response?: {
|
||||
cloudaicompanionProject?: { id?: string };
|
||||
};
|
||||
}
|
||||
|
||||
// Tier IDs as used by the Cloud Code API
|
||||
const TIER_FREE = "free-tier";
|
||||
const TIER_LEGACY = "legacy-tier";
|
||||
const TIER_STANDARD = "standard-tier";
|
||||
|
||||
interface GoogleRpcErrorResponse {
|
||||
error?: {
|
||||
details?: Array<{ reason?: string }>;
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Wait helper for onboarding retries
|
||||
*/
|
||||
function wait(ms: number): Promise<void> {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
|
||||
/**
|
||||
* Get default tier from allowed tiers
|
||||
*/
|
||||
function getDefaultTier(allowedTiers?: Array<{ id?: string; isDefault?: boolean }>): { id?: string } {
|
||||
if (!allowedTiers || allowedTiers.length === 0) return { id: TIER_LEGACY };
|
||||
const defaultTier = allowedTiers.find((t) => t.isDefault);
|
||||
return defaultTier ?? { id: TIER_LEGACY };
|
||||
}
|
||||
|
||||
function isVpcScAffectedUser(payload: unknown): boolean {
|
||||
if (!payload || typeof payload !== "object") return false;
|
||||
if (!("error" in payload)) return false;
|
||||
const error = (payload as GoogleRpcErrorResponse).error;
|
||||
if (!error?.details || !Array.isArray(error.details)) return false;
|
||||
return error.details.some((detail) => detail.reason === "SECURITY_POLICY_VIOLATED");
|
||||
}
|
||||
|
||||
/**
|
||||
* Poll a long-running operation until completion
|
||||
*/
|
||||
async function pollOperation(
|
||||
operationName: string,
|
||||
headers: Record<string, string>,
|
||||
onProgress?: (message: string) => void,
|
||||
): Promise<LongRunningOperationResponse> {
|
||||
let attempt = 0;
|
||||
while (true) {
|
||||
if (attempt > 0) {
|
||||
onProgress?.(`Waiting for project provisioning (attempt ${attempt + 1})...`);
|
||||
await wait(5000);
|
||||
}
|
||||
|
||||
const response = await fetch(`${CODE_ASSIST_ENDPOINT}/v1internal/${operationName}`, {
|
||||
method: "GET",
|
||||
headers,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`Failed to poll operation: ${response.status} ${response.statusText}`);
|
||||
}
|
||||
|
||||
const data = (await response.json()) as LongRunningOperationResponse;
|
||||
if (data.done) {
|
||||
return data;
|
||||
}
|
||||
|
||||
attempt += 1;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Discover or provision a Google Cloud project for the user
|
||||
*/
|
||||
async function discoverProject(accessToken: string, onProgress?: (message: string) => void): Promise<string> {
|
||||
// Check for user-provided project ID via environment variable
|
||||
const envProjectId = process.env.GOOGLE_CLOUD_PROJECT || process.env.GOOGLE_CLOUD_PROJECT_ID;
|
||||
|
||||
const headers = {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": "google-api-nodejs-client/9.15.1",
|
||||
"X-Goog-Api-Client": "gl-node/22.17.0",
|
||||
};
|
||||
|
||||
// Try to load existing project via loadCodeAssist
|
||||
onProgress?.("Checking for existing Cloud Code Assist project...");
|
||||
const loadResponse = await fetch(`${CODE_ASSIST_ENDPOINT}/v1internal:loadCodeAssist`, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify({
|
||||
cloudaicompanionProject: envProjectId,
|
||||
metadata: {
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
duetProject: envProjectId,
|
||||
},
|
||||
}),
|
||||
});
|
||||
|
||||
let data: LoadCodeAssistPayload;
|
||||
|
||||
if (!loadResponse.ok) {
|
||||
let errorPayload: unknown;
|
||||
try {
|
||||
errorPayload = await loadResponse.clone().json();
|
||||
} catch {
|
||||
errorPayload = undefined;
|
||||
}
|
||||
|
||||
if (isVpcScAffectedUser(errorPayload)) {
|
||||
data = { currentTier: { id: TIER_STANDARD } };
|
||||
} else {
|
||||
const errorText = await loadResponse.text();
|
||||
throw new Error(`loadCodeAssist failed: ${loadResponse.status} ${loadResponse.statusText}: ${errorText}`);
|
||||
}
|
||||
} else {
|
||||
data = (await loadResponse.json()) as LoadCodeAssistPayload;
|
||||
}
|
||||
|
||||
// If user already has a current tier and project, use it
|
||||
if (data.currentTier) {
|
||||
if (data.cloudaicompanionProject) {
|
||||
return data.cloudaicompanionProject;
|
||||
}
|
||||
// User has a tier but no managed project - they need to provide one via env var
|
||||
if (envProjectId) {
|
||||
return envProjectId;
|
||||
}
|
||||
throw new Error(
|
||||
"This account requires setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID environment variable. " +
|
||||
"See https://goo.gle/gemini-cli-auth-docs#workspace-gca",
|
||||
);
|
||||
}
|
||||
|
||||
// User needs to be onboarded - get the default tier
|
||||
const tier = getDefaultTier(data.allowedTiers);
|
||||
const tierId = tier?.id ?? TIER_FREE;
|
||||
|
||||
if (tierId !== TIER_FREE && !envProjectId) {
|
||||
throw new Error(
|
||||
"This account requires setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID environment variable. " +
|
||||
"See https://goo.gle/gemini-cli-auth-docs#workspace-gca",
|
||||
);
|
||||
}
|
||||
|
||||
onProgress?.("Provisioning Cloud Code Assist project (this may take a moment)...");
|
||||
|
||||
// Build onboard request - for free tier, don't include project ID (Google provisions one)
|
||||
// For other tiers, include the user's project ID if available
|
||||
const onboardBody: Record<string, unknown> = {
|
||||
tierId,
|
||||
metadata: {
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
},
|
||||
};
|
||||
|
||||
if (tierId !== TIER_FREE && envProjectId) {
|
||||
onboardBody.cloudaicompanionProject = envProjectId;
|
||||
(onboardBody.metadata as Record<string, unknown>).duetProject = envProjectId;
|
||||
}
|
||||
|
||||
// Start onboarding - this returns a long-running operation
|
||||
const onboardResponse = await fetch(`${CODE_ASSIST_ENDPOINT}/v1internal:onboardUser`, {
|
||||
method: "POST",
|
||||
headers,
|
||||
body: JSON.stringify(onboardBody),
|
||||
});
|
||||
|
||||
if (!onboardResponse.ok) {
|
||||
const errorText = await onboardResponse.text();
|
||||
throw new Error(`onboardUser failed: ${onboardResponse.status} ${onboardResponse.statusText}: ${errorText}`);
|
||||
}
|
||||
|
||||
let lroData = (await onboardResponse.json()) as LongRunningOperationResponse;
|
||||
|
||||
// If the operation isn't done yet, poll until completion
|
||||
if (!lroData.done && lroData.name) {
|
||||
lroData = await pollOperation(lroData.name, headers, onProgress);
|
||||
}
|
||||
|
||||
// Try to get project ID from the response
|
||||
const projectId = lroData.response?.cloudaicompanionProject?.id;
|
||||
if (projectId) {
|
||||
return projectId;
|
||||
}
|
||||
|
||||
// If no project ID from onboarding, fall back to env var
|
||||
if (envProjectId) {
|
||||
return envProjectId;
|
||||
}
|
||||
|
||||
throw new Error(
|
||||
"Could not discover or provision a Google Cloud project. " +
|
||||
"Try setting the GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID environment variable. " +
|
||||
"See https://goo.gle/gemini-cli-auth-docs#workspace-gca",
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get user email from the access token
|
||||
*/
|
||||
async function getUserEmail(accessToken: string): Promise<string | undefined> {
|
||||
try {
|
||||
const response = await fetch("https://www.googleapis.com/oauth2/v1/userinfo?alt=json", {
|
||||
headers: {
|
||||
Authorization: `Bearer ${accessToken}`,
|
||||
},
|
||||
});
|
||||
|
||||
if (response.ok) {
|
||||
const data = (await response.json()) as { email?: string };
|
||||
return data.email;
|
||||
}
|
||||
} catch {
|
||||
// Ignore errors, email is optional
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh Google Cloud Code Assist token
|
||||
*/
|
||||
export async function refreshGoogleCloudToken(refreshToken: string, projectId: string): Promise<OAuthCredentials> {
|
||||
const response = await fetch(TOKEN_URL, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/x-www-form-urlencoded" },
|
||||
body: new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
client_secret: CLIENT_SECRET,
|
||||
refresh_token: refreshToken,
|
||||
grant_type: "refresh_token",
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const error = await response.text();
|
||||
throw new Error(`Google Cloud token refresh failed: ${error}`);
|
||||
}
|
||||
|
||||
const data = (await response.json()) as {
|
||||
access_token: string;
|
||||
expires_in: number;
|
||||
refresh_token?: string;
|
||||
};
|
||||
|
||||
return {
|
||||
refresh: data.refresh_token || refreshToken,
|
||||
access: data.access_token,
|
||||
expires: Date.now() + data.expires_in * 1000 - 5 * 60 * 1000,
|
||||
projectId,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Login with Gemini CLI (Google Cloud Code Assist) OAuth
|
||||
*
|
||||
* @param onAuth - Callback with URL and optional instructions
|
||||
* @param onProgress - Optional progress callback
|
||||
* @param onManualCodeInput - Optional promise that resolves with user-pasted redirect URL.
|
||||
* Races with browser callback - whichever completes first wins.
|
||||
*/
|
||||
export async function loginGeminiCli(
|
||||
onAuth: (info: { url: string; instructions?: string }) => void,
|
||||
onProgress?: (message: string) => void,
|
||||
onManualCodeInput?: () => Promise<string>,
|
||||
): Promise<OAuthCredentials> {
|
||||
const { verifier, challenge } = await generatePKCE();
|
||||
|
||||
// Start local server for callback
|
||||
onProgress?.("Starting local server for OAuth callback...");
|
||||
const server = await startCallbackServer();
|
||||
|
||||
let code: string | undefined;
|
||||
|
||||
try {
|
||||
// Build authorization URL
|
||||
const authParams = new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
response_type: "code",
|
||||
redirect_uri: REDIRECT_URI,
|
||||
scope: SCOPES.join(" "),
|
||||
code_challenge: challenge,
|
||||
code_challenge_method: "S256",
|
||||
state: verifier,
|
||||
access_type: "offline",
|
||||
prompt: "consent",
|
||||
});
|
||||
|
||||
const authUrl = `${AUTH_URL}?${authParams.toString()}`;
|
||||
|
||||
// Notify caller with URL to open
|
||||
onAuth({
|
||||
url: authUrl,
|
||||
instructions: "Complete the sign-in in your browser.",
|
||||
});
|
||||
|
||||
// Wait for the callback, racing with manual input if provided
|
||||
onProgress?.("Waiting for OAuth callback...");
|
||||
|
||||
if (onManualCodeInput) {
|
||||
// Race between browser callback and manual input
|
||||
let manualInput: string | undefined;
|
||||
let manualError: Error | undefined;
|
||||
const manualPromise = onManualCodeInput()
|
||||
.then((input) => {
|
||||
manualInput = input;
|
||||
server.cancelWait();
|
||||
})
|
||||
.catch((err) => {
|
||||
manualError = err instanceof Error ? err : new Error(String(err));
|
||||
server.cancelWait();
|
||||
});
|
||||
|
||||
const result = await server.waitForCode();
|
||||
|
||||
// If manual input was cancelled, throw that error
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
|
||||
if (result?.code) {
|
||||
// Browser callback won - verify state
|
||||
if (result.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = result.code;
|
||||
} else if (manualInput) {
|
||||
// Manual input won
|
||||
const parsed = parseRedirectUrl(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
|
||||
// If still no code, wait for manual promise and try that
|
||||
if (!code) {
|
||||
await manualPromise;
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
if (manualInput) {
|
||||
const parsed = parseRedirectUrl(manualInput);
|
||||
if (parsed.state && parsed.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Original flow: just wait for callback
|
||||
const result = await server.waitForCode();
|
||||
if (result?.code) {
|
||||
if (result.state !== verifier) {
|
||||
throw new Error("OAuth state mismatch - possible CSRF attack");
|
||||
}
|
||||
code = result.code;
|
||||
}
|
||||
}
|
||||
|
||||
if (!code) {
|
||||
throw new Error("No authorization code received");
|
||||
}
|
||||
|
||||
// Exchange code for tokens
|
||||
onProgress?.("Exchanging authorization code for tokens...");
|
||||
const tokenResponse = await fetch(TOKEN_URL, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
},
|
||||
body: new URLSearchParams({
|
||||
client_id: CLIENT_ID,
|
||||
client_secret: CLIENT_SECRET,
|
||||
code,
|
||||
grant_type: "authorization_code",
|
||||
redirect_uri: REDIRECT_URI,
|
||||
code_verifier: verifier,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!tokenResponse.ok) {
|
||||
const error = await tokenResponse.text();
|
||||
throw new Error(`Token exchange failed: ${error}`);
|
||||
}
|
||||
|
||||
const tokenData = (await tokenResponse.json()) as {
|
||||
access_token: string;
|
||||
refresh_token: string;
|
||||
expires_in: number;
|
||||
};
|
||||
|
||||
if (!tokenData.refresh_token) {
|
||||
throw new Error("No refresh token received. Please try again.");
|
||||
}
|
||||
|
||||
// Get user email
|
||||
onProgress?.("Getting user info...");
|
||||
const email = await getUserEmail(tokenData.access_token);
|
||||
|
||||
// Discover project
|
||||
const projectId = await discoverProject(tokenData.access_token, onProgress);
|
||||
|
||||
// Calculate expiry time (current time + expires_in seconds - 5 min buffer)
|
||||
const expiresAt = Date.now() + tokenData.expires_in * 1000 - 5 * 60 * 1000;
|
||||
|
||||
const credentials: OAuthCredentials = {
|
||||
refresh: tokenData.refresh_token,
|
||||
access: tokenData.access_token,
|
||||
expires: expiresAt,
|
||||
projectId,
|
||||
email,
|
||||
};
|
||||
|
||||
return credentials;
|
||||
} finally {
|
||||
server.server.close();
|
||||
}
|
||||
}
|
||||
|
||||
export const geminiCliOAuthProvider: OAuthProviderInterface = {
|
||||
id: "google-gemini-cli",
|
||||
name: "Google Cloud Code Assist (Gemini CLI)",
|
||||
usesCallbackServer: true,
|
||||
|
||||
async login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
|
||||
return loginGeminiCli(callbacks.onAuth, callbacks.onProgress, callbacks.onManualCodeInput);
|
||||
},
|
||||
|
||||
async refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
|
||||
const creds = credentials as GeminiCredentials;
|
||||
if (!creds.projectId) {
|
||||
throw new Error("Google Cloud credentials missing projectId");
|
||||
}
|
||||
return refreshGoogleCloudToken(creds.refresh, creds.projectId);
|
||||
},
|
||||
|
||||
getApiKey(credentials: OAuthCredentials): string {
|
||||
const creds = credentials as GeminiCredentials;
|
||||
return JSON.stringify({ token: creds.access, projectId: creds.projectId });
|
||||
},
|
||||
};
|
||||
@@ -1,162 +0,0 @@
|
||||
/**
|
||||
* OAuth credential management for AI providers.
|
||||
*
|
||||
* This module handles login, token refresh, and credential storage
|
||||
* for OAuth-based providers:
|
||||
* - Anthropic (Claude Pro/Max)
|
||||
* - GitHub Copilot
|
||||
* - Google Cloud Code Assist (Gemini CLI)
|
||||
* - Antigravity (Gemini 3, Claude, GPT-OSS via Google Cloud)
|
||||
*/
|
||||
|
||||
// Anthropic
|
||||
export { anthropicOAuthProvider, loginAnthropic, refreshAnthropicToken } from "./anthropic.js";
|
||||
// GitHub Copilot
|
||||
export {
|
||||
getGitHubCopilotBaseUrl,
|
||||
githubCopilotOAuthProvider,
|
||||
loginGitHubCopilot,
|
||||
normalizeDomain,
|
||||
refreshGitHubCopilotToken,
|
||||
} from "./github-copilot.js";
|
||||
// Google Antigravity
|
||||
export { antigravityOAuthProvider, loginAntigravity, refreshAntigravityToken } from "./google-antigravity.js";
|
||||
// Google Gemini CLI
|
||||
export { geminiCliOAuthProvider, loginGeminiCli, refreshGoogleCloudToken } from "./google-gemini-cli.js";
|
||||
// OpenAI Codex (ChatGPT OAuth)
|
||||
export { loginOpenAICodex, openaiCodexOAuthProvider, refreshOpenAICodexToken } from "./openai-codex.js";
|
||||
|
||||
export * from "./types.js";
|
||||
|
||||
// ============================================================================
|
||||
// Provider Registry
|
||||
// ============================================================================
|
||||
|
||||
import { anthropicOAuthProvider } from "./anthropic.js";
|
||||
import { githubCopilotOAuthProvider } from "./github-copilot.js";
|
||||
import { antigravityOAuthProvider } from "./google-antigravity.js";
|
||||
import { geminiCliOAuthProvider } from "./google-gemini-cli.js";
|
||||
import { openaiCodexOAuthProvider } from "./openai-codex.js";
|
||||
import type { OAuthCredentials, OAuthProviderId, OAuthProviderInfo, OAuthProviderInterface } from "./types.js";
|
||||
|
||||
const BUILT_IN_OAUTH_PROVIDERS: OAuthProviderInterface[] = [
|
||||
anthropicOAuthProvider,
|
||||
githubCopilotOAuthProvider,
|
||||
geminiCliOAuthProvider,
|
||||
antigravityOAuthProvider,
|
||||
openaiCodexOAuthProvider,
|
||||
];
|
||||
|
||||
const oauthProviderRegistry = new Map<string, OAuthProviderInterface>(
|
||||
BUILT_IN_OAUTH_PROVIDERS.map((provider) => [provider.id, provider]),
|
||||
);
|
||||
|
||||
/**
|
||||
* Get an OAuth provider by ID
|
||||
*/
|
||||
export function getOAuthProvider(id: OAuthProviderId): OAuthProviderInterface | undefined {
|
||||
return oauthProviderRegistry.get(id);
|
||||
}
|
||||
|
||||
/**
|
||||
* Register a custom OAuth provider
|
||||
*/
|
||||
export function registerOAuthProvider(provider: OAuthProviderInterface): void {
|
||||
oauthProviderRegistry.set(provider.id, provider);
|
||||
}
|
||||
|
||||
/**
|
||||
* Unregister an OAuth provider.
|
||||
*
|
||||
* If the provider is built-in, restores the built-in implementation.
|
||||
* Custom providers are removed completely.
|
||||
*/
|
||||
export function unregisterOAuthProvider(id: string): void {
|
||||
const builtInProvider = BUILT_IN_OAUTH_PROVIDERS.find((provider) => provider.id === id);
|
||||
if (builtInProvider) {
|
||||
oauthProviderRegistry.set(id, builtInProvider);
|
||||
return;
|
||||
}
|
||||
oauthProviderRegistry.delete(id);
|
||||
}
|
||||
|
||||
/**
|
||||
* Reset OAuth providers to built-ins.
|
||||
*/
|
||||
export function resetOAuthProviders(): void {
|
||||
oauthProviderRegistry.clear();
|
||||
for (const provider of BUILT_IN_OAUTH_PROVIDERS) {
|
||||
oauthProviderRegistry.set(provider.id, provider);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all registered OAuth providers
|
||||
*/
|
||||
export function getOAuthProviders(): OAuthProviderInterface[] {
|
||||
return Array.from(oauthProviderRegistry.values());
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated Use getOAuthProviders() which returns OAuthProviderInterface[]
|
||||
*/
|
||||
export function getOAuthProviderInfoList(): OAuthProviderInfo[] {
|
||||
return getOAuthProviders().map((p) => ({
|
||||
id: p.id,
|
||||
name: p.name,
|
||||
available: true,
|
||||
}));
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// High-level API (uses provider registry)
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* Refresh token for any OAuth provider.
|
||||
* @deprecated Use getOAuthProvider(id).refreshToken() instead
|
||||
*/
|
||||
export async function refreshOAuthToken(
|
||||
providerId: OAuthProviderId,
|
||||
credentials: OAuthCredentials,
|
||||
): Promise<OAuthCredentials> {
|
||||
const provider = getOAuthProvider(providerId);
|
||||
if (!provider) {
|
||||
throw new Error(`Unknown OAuth provider: ${providerId}`);
|
||||
}
|
||||
return provider.refreshToken(credentials);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get API key for a provider from OAuth credentials.
|
||||
* Automatically refreshes expired tokens.
|
||||
*
|
||||
* @returns API key string and updated credentials, or null if no credentials
|
||||
* @throws Error if refresh fails
|
||||
*/
|
||||
export async function getOAuthApiKey(
|
||||
providerId: OAuthProviderId,
|
||||
credentials: Record<string, OAuthCredentials>,
|
||||
): Promise<{ newCredentials: OAuthCredentials; apiKey: string } | null> {
|
||||
const provider = getOAuthProvider(providerId);
|
||||
if (!provider) {
|
||||
throw new Error(`Unknown OAuth provider: ${providerId}`);
|
||||
}
|
||||
|
||||
let creds = credentials[providerId];
|
||||
if (!creds) {
|
||||
return null;
|
||||
}
|
||||
|
||||
// Refresh if expired
|
||||
if (Date.now() >= creds.expires) {
|
||||
try {
|
||||
creds = await provider.refreshToken(creds);
|
||||
} catch (_error) {
|
||||
throw new Error(`Failed to refresh OAuth token for ${providerId}`);
|
||||
}
|
||||
}
|
||||
|
||||
const apiKey = provider.getApiKey(creds);
|
||||
return { newCredentials: creds, apiKey };
|
||||
}
|
||||
@@ -1,109 +0,0 @@
|
||||
const LOGO_SVG = `<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 800 800" aria-hidden="true"><path fill="#fff" fill-rule="evenodd" d="M165.29 165.29 H517.36 V400 H400 V517.36 H282.65 V634.72 H165.29 Z M282.65 282.65 V400 H400 V282.65 Z"/><path fill="#fff" d="M517.36 400 H634.72 V634.72 H517.36 Z"/></svg>`;
|
||||
|
||||
function escapeHtml(value: string): string {
|
||||
return value
|
||||
.replaceAll("&", "&")
|
||||
.replaceAll("<", "<")
|
||||
.replaceAll(">", ">")
|
||||
.replaceAll('"', """)
|
||||
.replaceAll("'", "'");
|
||||
}
|
||||
|
||||
function renderPage(options: { title: string; heading: string; message: string; details?: string }): string {
|
||||
const title = escapeHtml(options.title);
|
||||
const heading = escapeHtml(options.heading);
|
||||
const message = escapeHtml(options.message);
|
||||
const details = options.details ? escapeHtml(options.details) : undefined;
|
||||
|
||||
return `<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1" />
|
||||
<title>${title}</title>
|
||||
<style>
|
||||
:root {
|
||||
--text: #fafafa;
|
||||
--text-dim: #a1a1aa;
|
||||
--page-bg: #09090b;
|
||||
--font-sans: ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, "Noto Sans", sans-serif, "Apple Color Emoji", "Segoe UI Emoji", "Segoe UI Symbol", "Noto Color Emoji";
|
||||
--font-mono: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
html { color-scheme: dark; }
|
||||
body {
|
||||
margin: 0;
|
||||
min-height: 100vh;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
padding: 24px;
|
||||
background: var(--page-bg);
|
||||
color: var(--text);
|
||||
font-family: var(--font-sans);
|
||||
text-align: center;
|
||||
}
|
||||
main {
|
||||
width: 100%;
|
||||
max-width: 560px;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
}
|
||||
.logo {
|
||||
width: 72px;
|
||||
height: 72px;
|
||||
display: block;
|
||||
margin-bottom: 24px;
|
||||
}
|
||||
h1 {
|
||||
margin: 0 0 10px;
|
||||
font-size: 28px;
|
||||
line-height: 1.15;
|
||||
font-weight: 650;
|
||||
color: var(--text);
|
||||
}
|
||||
p {
|
||||
margin: 0;
|
||||
line-height: 1.7;
|
||||
color: var(--text-dim);
|
||||
font-size: 15px;
|
||||
}
|
||||
.details {
|
||||
margin-top: 16px;
|
||||
font-family: var(--font-mono);
|
||||
font-size: 13px;
|
||||
color: var(--text-dim);
|
||||
white-space: pre-wrap;
|
||||
word-break: break-word;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<main>
|
||||
<div class="logo">${LOGO_SVG}</div>
|
||||
<h1>${heading}</h1>
|
||||
<p>${message}</p>
|
||||
${details ? `<div class="details">${details}</div>` : ""}
|
||||
</main>
|
||||
</body>
|
||||
</html>`;
|
||||
}
|
||||
|
||||
export function oauthSuccessHtml(message: string): string {
|
||||
return renderPage({
|
||||
title: "Authentication successful",
|
||||
heading: "Authentication successful",
|
||||
message,
|
||||
});
|
||||
}
|
||||
|
||||
export function oauthErrorHtml(message: string, details?: string): string {
|
||||
return renderPage({
|
||||
title: "Authentication failed",
|
||||
heading: "Authentication failed",
|
||||
message,
|
||||
details,
|
||||
});
|
||||
}
|
||||
@@ -1,450 +0,0 @@
|
||||
/**
|
||||
* OpenAI Codex (ChatGPT OAuth) flow
|
||||
*
|
||||
* NOTE: This module uses Node.js crypto and http for the OAuth callback.
|
||||
* It is only intended for CLI use, not browser environments.
|
||||
*/
|
||||
|
||||
// NEVER convert to top-level imports - breaks browser/Vite builds (web-ui)
|
||||
let _randomBytes: typeof import("node:crypto").randomBytes | null = null;
|
||||
let _http: typeof import("node:http") | null = null;
|
||||
if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) {
|
||||
import("node:crypto").then((m) => {
|
||||
_randomBytes = m.randomBytes;
|
||||
});
|
||||
import("node:http").then((m) => {
|
||||
_http = m;
|
||||
});
|
||||
}
|
||||
|
||||
import { oauthErrorHtml, oauthSuccessHtml } from "./oauth-page.js";
|
||||
import { generatePKCE } from "./pkce.js";
|
||||
import type { OAuthCredentials, OAuthLoginCallbacks, OAuthPrompt, OAuthProviderInterface } from "./types.js";
|
||||
|
||||
const CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann";
|
||||
const AUTHORIZE_URL = "https://auth.openai.com/oauth/authorize";
|
||||
const TOKEN_URL = "https://auth.openai.com/oauth/token";
|
||||
const REDIRECT_URI = "http://localhost:1455/auth/callback";
|
||||
const SCOPE = "openid profile email offline_access";
|
||||
const JWT_CLAIM_PATH = "https://api.openai.com/auth";
|
||||
|
||||
type TokenSuccess = { type: "success"; access: string; refresh: string; expires: number };
|
||||
type TokenFailure = { type: "failed" };
|
||||
type TokenResult = TokenSuccess | TokenFailure;
|
||||
|
||||
type JwtPayload = {
|
||||
[JWT_CLAIM_PATH]?: {
|
||||
chatgpt_account_id?: string;
|
||||
};
|
||||
[key: string]: unknown;
|
||||
};
|
||||
|
||||
function createState(): string {
|
||||
if (!_randomBytes) {
|
||||
throw new Error("OpenAI Codex OAuth is only available in Node.js environments");
|
||||
}
|
||||
return _randomBytes(16).toString("hex");
|
||||
}
|
||||
|
||||
function parseAuthorizationInput(input: string): { code?: string; state?: string } {
|
||||
const value = input.trim();
|
||||
if (!value) return {};
|
||||
|
||||
try {
|
||||
const url = new URL(value);
|
||||
return {
|
||||
code: url.searchParams.get("code") ?? undefined,
|
||||
state: url.searchParams.get("state") ?? undefined,
|
||||
};
|
||||
} catch {
|
||||
// not a URL
|
||||
}
|
||||
|
||||
if (value.includes("#")) {
|
||||
const [code, state] = value.split("#", 2);
|
||||
return { code, state };
|
||||
}
|
||||
|
||||
if (value.includes("code=")) {
|
||||
const params = new URLSearchParams(value);
|
||||
return {
|
||||
code: params.get("code") ?? undefined,
|
||||
state: params.get("state") ?? undefined,
|
||||
};
|
||||
}
|
||||
|
||||
return { code: value };
|
||||
}
|
||||
|
||||
function decodeJwt(token: string): JwtPayload | null {
|
||||
try {
|
||||
const parts = token.split(".");
|
||||
if (parts.length !== 3) return null;
|
||||
const payload = parts[1] ?? "";
|
||||
const decoded = atob(payload);
|
||||
return JSON.parse(decoded) as JwtPayload;
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
async function exchangeAuthorizationCode(
|
||||
code: string,
|
||||
verifier: string,
|
||||
redirectUri: string = REDIRECT_URI,
|
||||
): Promise<TokenResult> {
|
||||
const response = await fetch(TOKEN_URL, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/x-www-form-urlencoded" },
|
||||
body: new URLSearchParams({
|
||||
grant_type: "authorization_code",
|
||||
client_id: CLIENT_ID,
|
||||
code,
|
||||
code_verifier: verifier,
|
||||
redirect_uri: redirectUri,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const text = await response.text().catch(() => "");
|
||||
console.error("[openai-codex] code->token failed:", response.status, text);
|
||||
return { type: "failed" };
|
||||
}
|
||||
|
||||
const json = (await response.json()) as {
|
||||
access_token?: string;
|
||||
refresh_token?: string;
|
||||
expires_in?: number;
|
||||
};
|
||||
|
||||
if (!json.access_token || !json.refresh_token || typeof json.expires_in !== "number") {
|
||||
console.error("[openai-codex] token response missing fields:", json);
|
||||
return { type: "failed" };
|
||||
}
|
||||
|
||||
return {
|
||||
type: "success",
|
||||
access: json.access_token,
|
||||
refresh: json.refresh_token,
|
||||
expires: Date.now() + json.expires_in * 1000,
|
||||
};
|
||||
}
|
||||
|
||||
async function refreshAccessToken(refreshToken: string): Promise<TokenResult> {
|
||||
try {
|
||||
const response = await fetch(TOKEN_URL, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/x-www-form-urlencoded" },
|
||||
body: new URLSearchParams({
|
||||
grant_type: "refresh_token",
|
||||
refresh_token: refreshToken,
|
||||
client_id: CLIENT_ID,
|
||||
}),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const text = await response.text().catch(() => "");
|
||||
console.error("[openai-codex] Token refresh failed:", response.status, text);
|
||||
return { type: "failed" };
|
||||
}
|
||||
|
||||
const json = (await response.json()) as {
|
||||
access_token?: string;
|
||||
refresh_token?: string;
|
||||
expires_in?: number;
|
||||
};
|
||||
|
||||
if (!json.access_token || !json.refresh_token || typeof json.expires_in !== "number") {
|
||||
console.error("[openai-codex] Token refresh response missing fields:", json);
|
||||
return { type: "failed" };
|
||||
}
|
||||
|
||||
return {
|
||||
type: "success",
|
||||
access: json.access_token,
|
||||
refresh: json.refresh_token,
|
||||
expires: Date.now() + json.expires_in * 1000,
|
||||
};
|
||||
} catch (error) {
|
||||
console.error("[openai-codex] Token refresh error:", error);
|
||||
return { type: "failed" };
|
||||
}
|
||||
}
|
||||
|
||||
async function createAuthorizationFlow(
|
||||
originator: string = "pi",
|
||||
): Promise<{ verifier: string; state: string; url: string }> {
|
||||
const { verifier, challenge } = await generatePKCE();
|
||||
const state = createState();
|
||||
|
||||
const url = new URL(AUTHORIZE_URL);
|
||||
url.searchParams.set("response_type", "code");
|
||||
url.searchParams.set("client_id", CLIENT_ID);
|
||||
url.searchParams.set("redirect_uri", REDIRECT_URI);
|
||||
url.searchParams.set("scope", SCOPE);
|
||||
url.searchParams.set("code_challenge", challenge);
|
||||
url.searchParams.set("code_challenge_method", "S256");
|
||||
url.searchParams.set("state", state);
|
||||
url.searchParams.set("id_token_add_organizations", "true");
|
||||
url.searchParams.set("codex_cli_simplified_flow", "true");
|
||||
url.searchParams.set("originator", originator);
|
||||
|
||||
return { verifier, state, url: url.toString() };
|
||||
}
|
||||
|
||||
type OAuthServerInfo = {
|
||||
close: () => void;
|
||||
cancelWait: () => void;
|
||||
waitForCode: () => Promise<{ code: string } | null>;
|
||||
};
|
||||
|
||||
function startLocalOAuthServer(state: string): Promise<OAuthServerInfo> {
|
||||
if (!_http) {
|
||||
throw new Error("OpenAI Codex OAuth is only available in Node.js environments");
|
||||
}
|
||||
|
||||
let settleWait: ((value: { code: string } | null) => void) | undefined;
|
||||
const waitForCodePromise = new Promise<{ code: string } | null>((resolve) => {
|
||||
let settled = false;
|
||||
settleWait = (value) => {
|
||||
if (settled) return;
|
||||
settled = true;
|
||||
resolve(value);
|
||||
};
|
||||
});
|
||||
|
||||
const server = _http.createServer((req, res) => {
|
||||
try {
|
||||
const url = new URL(req.url || "", "http://localhost");
|
||||
if (url.pathname !== "/auth/callback") {
|
||||
res.statusCode = 404;
|
||||
res.setHeader("Content-Type", "text/html; charset=utf-8");
|
||||
res.end(oauthErrorHtml("Callback route not found."));
|
||||
return;
|
||||
}
|
||||
if (url.searchParams.get("state") !== state) {
|
||||
res.statusCode = 400;
|
||||
res.setHeader("Content-Type", "text/html; charset=utf-8");
|
||||
res.end(oauthErrorHtml("State mismatch."));
|
||||
return;
|
||||
}
|
||||
const code = url.searchParams.get("code");
|
||||
if (!code) {
|
||||
res.statusCode = 400;
|
||||
res.setHeader("Content-Type", "text/html; charset=utf-8");
|
||||
res.end(oauthErrorHtml("Missing authorization code."));
|
||||
return;
|
||||
}
|
||||
res.statusCode = 200;
|
||||
res.setHeader("Content-Type", "text/html; charset=utf-8");
|
||||
res.end(oauthSuccessHtml("OpenAI authentication completed. You can close this window."));
|
||||
settleWait?.({ code });
|
||||
} catch {
|
||||
res.statusCode = 500;
|
||||
res.setHeader("Content-Type", "text/html; charset=utf-8");
|
||||
res.end(oauthErrorHtml("Internal error while processing OAuth callback."));
|
||||
}
|
||||
});
|
||||
|
||||
return new Promise((resolve) => {
|
||||
server
|
||||
.listen(1455, "127.0.0.1", () => {
|
||||
resolve({
|
||||
close: () => server.close(),
|
||||
cancelWait: () => {
|
||||
settleWait?.(null);
|
||||
},
|
||||
waitForCode: () => waitForCodePromise,
|
||||
});
|
||||
})
|
||||
.on("error", (err: NodeJS.ErrnoException) => {
|
||||
console.error(
|
||||
"[openai-codex] Failed to bind http://127.0.0.1:1455 (",
|
||||
err.code,
|
||||
") Falling back to manual paste.",
|
||||
);
|
||||
settleWait?.(null);
|
||||
resolve({
|
||||
close: () => {
|
||||
try {
|
||||
server.close();
|
||||
} catch {
|
||||
// ignore
|
||||
}
|
||||
},
|
||||
cancelWait: () => {},
|
||||
waitForCode: async () => null,
|
||||
});
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
function getAccountId(accessToken: string): string | null {
|
||||
const payload = decodeJwt(accessToken);
|
||||
const auth = payload?.[JWT_CLAIM_PATH];
|
||||
const accountId = auth?.chatgpt_account_id;
|
||||
return typeof accountId === "string" && accountId.length > 0 ? accountId : null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Login with OpenAI Codex OAuth
|
||||
*
|
||||
* @param options.onAuth - Called with URL and instructions when auth starts
|
||||
* @param options.onPrompt - Called to prompt user for manual code paste (fallback if no onManualCodeInput)
|
||||
* @param options.onProgress - Optional progress messages
|
||||
* @param options.onManualCodeInput - Optional promise that resolves with user-pasted code.
|
||||
* Races with browser callback - whichever completes first wins.
|
||||
* Useful for showing paste input immediately alongside browser flow.
|
||||
* @param options.originator - OAuth originator parameter (defaults to "pi")
|
||||
*/
|
||||
export async function loginOpenAICodex(options: {
|
||||
onAuth: (info: { url: string; instructions?: string }) => void;
|
||||
onPrompt: (prompt: OAuthPrompt) => Promise<string>;
|
||||
onProgress?: (message: string) => void;
|
||||
onManualCodeInput?: () => Promise<string>;
|
||||
originator?: string;
|
||||
}): Promise<OAuthCredentials> {
|
||||
const { verifier, state, url } = await createAuthorizationFlow(options.originator);
|
||||
const server = await startLocalOAuthServer(state);
|
||||
|
||||
options.onAuth({ url, instructions: "A browser window should open. Complete login to finish." });
|
||||
|
||||
let code: string | undefined;
|
||||
try {
|
||||
if (options.onManualCodeInput) {
|
||||
// Race between browser callback and manual input
|
||||
let manualCode: string | undefined;
|
||||
let manualError: Error | undefined;
|
||||
const manualPromise = options
|
||||
.onManualCodeInput()
|
||||
.then((input) => {
|
||||
manualCode = input;
|
||||
server.cancelWait();
|
||||
})
|
||||
.catch((err) => {
|
||||
manualError = err instanceof Error ? err : new Error(String(err));
|
||||
server.cancelWait();
|
||||
});
|
||||
|
||||
const result = await server.waitForCode();
|
||||
|
||||
// If manual input was cancelled, throw that error
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
|
||||
if (result?.code) {
|
||||
// Browser callback won
|
||||
code = result.code;
|
||||
} else if (manualCode) {
|
||||
// Manual input won (or callback timed out and user had entered code)
|
||||
const parsed = parseAuthorizationInput(manualCode);
|
||||
if (parsed.state && parsed.state !== state) {
|
||||
throw new Error("State mismatch");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
|
||||
// If still no code, wait for manual promise to complete and try that
|
||||
if (!code) {
|
||||
await manualPromise;
|
||||
if (manualError) {
|
||||
throw manualError;
|
||||
}
|
||||
if (manualCode) {
|
||||
const parsed = parseAuthorizationInput(manualCode);
|
||||
if (parsed.state && parsed.state !== state) {
|
||||
throw new Error("State mismatch");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Original flow: wait for callback, then prompt if needed
|
||||
const result = await server.waitForCode();
|
||||
if (result?.code) {
|
||||
code = result.code;
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to onPrompt if still no code
|
||||
if (!code) {
|
||||
const input = await options.onPrompt({
|
||||
message: "Paste the authorization code (or full redirect URL):",
|
||||
});
|
||||
const parsed = parseAuthorizationInput(input);
|
||||
if (parsed.state && parsed.state !== state) {
|
||||
throw new Error("State mismatch");
|
||||
}
|
||||
code = parsed.code;
|
||||
}
|
||||
|
||||
if (!code) {
|
||||
throw new Error("Missing authorization code");
|
||||
}
|
||||
|
||||
const tokenResult = await exchangeAuthorizationCode(code, verifier);
|
||||
if (tokenResult.type !== "success") {
|
||||
throw new Error("Token exchange failed");
|
||||
}
|
||||
|
||||
const accountId = getAccountId(tokenResult.access);
|
||||
if (!accountId) {
|
||||
throw new Error("Failed to extract accountId from token");
|
||||
}
|
||||
|
||||
return {
|
||||
access: tokenResult.access,
|
||||
refresh: tokenResult.refresh,
|
||||
expires: tokenResult.expires,
|
||||
accountId,
|
||||
};
|
||||
} finally {
|
||||
server.close();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Refresh OpenAI Codex OAuth token
|
||||
*/
|
||||
export async function refreshOpenAICodexToken(refreshToken: string): Promise<OAuthCredentials> {
|
||||
const result = await refreshAccessToken(refreshToken);
|
||||
if (result.type !== "success") {
|
||||
throw new Error("Failed to refresh OpenAI Codex token");
|
||||
}
|
||||
|
||||
const accountId = getAccountId(result.access);
|
||||
if (!accountId) {
|
||||
throw new Error("Failed to extract accountId from token");
|
||||
}
|
||||
|
||||
return {
|
||||
access: result.access,
|
||||
refresh: result.refresh,
|
||||
expires: result.expires,
|
||||
accountId,
|
||||
};
|
||||
}
|
||||
|
||||
export const openaiCodexOAuthProvider: OAuthProviderInterface = {
|
||||
id: "openai-codex",
|
||||
name: "ChatGPT Plus/Pro (Codex Subscription)",
|
||||
usesCallbackServer: true,
|
||||
|
||||
async login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
|
||||
return loginOpenAICodex({
|
||||
onAuth: callbacks.onAuth,
|
||||
onPrompt: callbacks.onPrompt,
|
||||
onProgress: callbacks.onProgress,
|
||||
onManualCodeInput: callbacks.onManualCodeInput,
|
||||
});
|
||||
},
|
||||
|
||||
async refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials> {
|
||||
return refreshOpenAICodexToken(credentials.refresh);
|
||||
},
|
||||
|
||||
getApiKey(credentials: OAuthCredentials): string {
|
||||
return credentials.access;
|
||||
},
|
||||
};
|
||||
@@ -1,34 +0,0 @@
|
||||
/**
|
||||
* PKCE utilities using Web Crypto API.
|
||||
* Works in both Node.js 20+ and browsers.
|
||||
*/
|
||||
|
||||
/**
|
||||
* Encode bytes as base64url string.
|
||||
*/
|
||||
function base64urlEncode(bytes: Uint8Array): string {
|
||||
let binary = "";
|
||||
for (const byte of bytes) {
|
||||
binary += String.fromCharCode(byte);
|
||||
}
|
||||
return btoa(binary).replace(/\+/g, "-").replace(/\//g, "_").replace(/=/g, "");
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate PKCE code verifier and challenge.
|
||||
* Uses Web Crypto API for cross-platform compatibility.
|
||||
*/
|
||||
export async function generatePKCE(): Promise<{ verifier: string; challenge: string }> {
|
||||
// Generate random verifier
|
||||
const verifierBytes = new Uint8Array(32);
|
||||
crypto.getRandomValues(verifierBytes);
|
||||
const verifier = base64urlEncode(verifierBytes);
|
||||
|
||||
// Compute SHA-256 challenge
|
||||
const encoder = new TextEncoder();
|
||||
const data = encoder.encode(verifier);
|
||||
const hashBuffer = await crypto.subtle.digest("SHA-256", data);
|
||||
const challenge = base64urlEncode(new Uint8Array(hashBuffer));
|
||||
|
||||
return { verifier, challenge };
|
||||
}
|
||||
@@ -1,59 +0,0 @@
|
||||
import type { Api, Model } from "../../types.js";
|
||||
|
||||
export type OAuthCredentials = {
|
||||
refresh: string;
|
||||
access: string;
|
||||
expires: number;
|
||||
[key: string]: unknown;
|
||||
};
|
||||
|
||||
export type OAuthProviderId = string;
|
||||
|
||||
/** @deprecated Use OAuthProviderId instead */
|
||||
export type OAuthProvider = OAuthProviderId;
|
||||
|
||||
export type OAuthPrompt = {
|
||||
message: string;
|
||||
placeholder?: string;
|
||||
allowEmpty?: boolean;
|
||||
};
|
||||
|
||||
export type OAuthAuthInfo = {
|
||||
url: string;
|
||||
instructions?: string;
|
||||
};
|
||||
|
||||
export interface OAuthLoginCallbacks {
|
||||
onAuth: (info: OAuthAuthInfo) => void;
|
||||
onPrompt: (prompt: OAuthPrompt) => Promise<string>;
|
||||
onProgress?: (message: string) => void;
|
||||
onManualCodeInput?: () => Promise<string>;
|
||||
signal?: AbortSignal;
|
||||
}
|
||||
|
||||
export interface OAuthProviderInterface {
|
||||
readonly id: OAuthProviderId;
|
||||
readonly name: string;
|
||||
|
||||
/** Run the login flow, return credentials to persist */
|
||||
login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials>;
|
||||
|
||||
/** Whether login uses a local callback server and supports manual code input. */
|
||||
usesCallbackServer?: boolean;
|
||||
|
||||
/** Refresh expired credentials, return updated credentials to persist */
|
||||
refreshToken(credentials: OAuthCredentials): Promise<OAuthCredentials>;
|
||||
|
||||
/** Convert credentials to API key string for the provider */
|
||||
getApiKey(credentials: OAuthCredentials): string;
|
||||
|
||||
/** Optional: modify models for this provider (e.g., update baseUrl) */
|
||||
modifyModels?(models: Model<Api>[], credentials: OAuthCredentials): Model<Api>[];
|
||||
}
|
||||
|
||||
/** @deprecated Use OAuthProviderInterface instead */
|
||||
export interface OAuthProviderInfo {
|
||||
id: OAuthProviderId;
|
||||
name: string;
|
||||
available: boolean;
|
||||
}
|
||||
@@ -1,138 +0,0 @@
|
||||
import type { AssistantMessage } from "../types.js";
|
||||
|
||||
/**
|
||||
* Regex patterns to detect context overflow errors from different providers.
|
||||
*
|
||||
* These patterns match error messages returned when the input exceeds
|
||||
* the model's context window.
|
||||
*
|
||||
* Provider-specific patterns (with example error messages):
|
||||
*
|
||||
* - Anthropic: "prompt is too long: 213462 tokens > 200000 maximum"
|
||||
* - Anthropic: "413 {\"error\":{\"type\":\"request_too_large\",\"message\":\"Request exceeds the maximum size\"}}"
|
||||
* - OpenAI: "Your input exceeds the context window of this model"
|
||||
* - Google: "The input token count (1196265) exceeds the maximum number of tokens allowed (1048575)"
|
||||
* - xAI: "This model's maximum prompt length is 131072 but the request contains 537812 tokens"
|
||||
* - Groq: "Please reduce the length of the messages or completion"
|
||||
* - OpenRouter: "This endpoint's maximum context length is X tokens. However, you requested about Y tokens"
|
||||
* - llama.cpp: "the request exceeds the available context size, try increasing it"
|
||||
* - LM Studio: "tokens to keep from the initial prompt is greater than the context length"
|
||||
* - GitHub Copilot: "prompt token count of X exceeds the limit of Y"
|
||||
* - MiniMax: "invalid params, context window exceeds limit"
|
||||
* - Kimi For Coding: "Your request exceeded model token limit: X (requested: Y)"
|
||||
* - Cerebras: "400/413 status code (no body)"
|
||||
* - Mistral: "Prompt contains X tokens ... too large for model with Y maximum context length"
|
||||
* - z.ai: Does NOT error, accepts overflow silently - handled via usage.input > contextWindow
|
||||
* - Ollama: Some deployments truncate silently, others return errors like "prompt too long; exceeded max context length by X tokens"
|
||||
*/
|
||||
const OVERFLOW_PATTERNS = [
|
||||
/prompt is too long/i, // Anthropic token overflow
|
||||
/request_too_large/i, // Anthropic request byte-size overflow (HTTP 413)
|
||||
/input is too long for requested model/i, // Amazon Bedrock
|
||||
/exceeds the context window/i, // OpenAI (Completions & Responses API)
|
||||
/input token count.*exceeds the maximum/i, // Google (Gemini)
|
||||
/maximum prompt length is \d+/i, // xAI (Grok)
|
||||
/reduce the length of the messages/i, // Groq
|
||||
/maximum context length is \d+ tokens/i, // OpenRouter (all backends)
|
||||
/exceeds the limit of \d+/i, // GitHub Copilot
|
||||
/exceeds the available context size/i, // llama.cpp server
|
||||
/greater than the context length/i, // LM Studio
|
||||
/context window exceeds limit/i, // MiniMax
|
||||
/exceeded model token limit/i, // Kimi For Coding
|
||||
/too large for model with \d+ maximum context length/i, // Mistral
|
||||
/model_context_window_exceeded/i, // z.ai non-standard finish_reason surfaced as error text
|
||||
/prompt too long; exceeded (?:max )?context length/i, // Ollama explicit overflow error
|
||||
/context[_ ]length[_ ]exceeded/i, // Generic fallback
|
||||
/too many tokens/i, // Generic fallback
|
||||
/token limit exceeded/i, // Generic fallback
|
||||
/^4(?:00|13)\s*(?:status code)?\s*\(no body\)/i, // Cerebras: 400/413 with no body
|
||||
];
|
||||
|
||||
/**
|
||||
* Patterns that indicate non-overflow errors (e.g. rate limiting, server errors).
|
||||
* Error messages matching any of these are excluded from overflow detection
|
||||
* even if they also match an OVERFLOW_PATTERN.
|
||||
*
|
||||
* Example: Bedrock formats throttling errors as "ThrottlingException: Too many tokens,
|
||||
* please wait before trying again." which would match the /too many tokens/i overflow
|
||||
* pattern without this exclusion.
|
||||
*/
|
||||
const NON_OVERFLOW_PATTERNS = [
|
||||
/^(Throttling error|Service unavailable):/i, // AWS Bedrock non-overflow errors (human-readable prefixes from formatBedrockError)
|
||||
/rate limit/i, // Generic rate limiting
|
||||
/too many requests/i, // Generic HTTP 429 style
|
||||
];
|
||||
|
||||
/**
|
||||
* Check if an assistant message represents a context overflow error.
|
||||
*
|
||||
* This handles two cases:
|
||||
* 1. Error-based overflow: Most providers return stopReason "error" with a
|
||||
* specific error message pattern.
|
||||
* 2. Silent overflow: Some providers accept overflow requests and return
|
||||
* successfully. For these, we check if usage.input exceeds the context window.
|
||||
*
|
||||
* ## Reliability by Provider
|
||||
*
|
||||
* **Reliable detection (returns error with detectable message):**
|
||||
* - Anthropic: "prompt is too long: X tokens > Y maximum" or "request_too_large"
|
||||
* - OpenAI (Completions & Responses): "exceeds the context window"
|
||||
* - Google Gemini: "input token count exceeds the maximum"
|
||||
* - xAI (Grok): "maximum prompt length is X but request contains Y"
|
||||
* - Groq: "reduce the length of the messages"
|
||||
* - Cerebras: 400/413 status code (no body)
|
||||
* - Mistral: "Prompt contains X tokens ... too large for model with Y maximum context length"
|
||||
* - OpenRouter (all backends): "maximum context length is X tokens"
|
||||
* - llama.cpp: "exceeds the available context size"
|
||||
* - LM Studio: "greater than the context length"
|
||||
* - Kimi For Coding: "exceeded model token limit: X (requested: Y)"
|
||||
*
|
||||
* **Unreliable detection:**
|
||||
* - z.ai: Sometimes accepts overflow silently (detectable via usage.input > contextWindow),
|
||||
* sometimes returns rate limit errors. Pass contextWindow param to detect silent overflow.
|
||||
* - Ollama: May truncate input silently for some setups, but may also return explicit
|
||||
* overflow errors that match the patterns above. Silent truncation still cannot be
|
||||
* detected here because we do not know the expected token count.
|
||||
*
|
||||
* ## Custom Providers
|
||||
*
|
||||
* If you've added custom models via settings.json, this function may not detect
|
||||
* overflow errors from those providers. To add support:
|
||||
*
|
||||
* 1. Send a request that exceeds the model's context window
|
||||
* 2. Check the errorMessage in the response
|
||||
* 3. Create a regex pattern that matches the error
|
||||
* 4. The pattern should be added to OVERFLOW_PATTERNS in this file, or
|
||||
* check the errorMessage yourself before calling this function
|
||||
*
|
||||
* @param message - The assistant message to check
|
||||
* @param contextWindow - Optional context window size for detecting silent overflow (z.ai)
|
||||
* @returns true if the message indicates a context overflow
|
||||
*/
|
||||
export function isContextOverflow(message: AssistantMessage, contextWindow?: number): boolean {
|
||||
// Case 1: Check error message patterns
|
||||
if (message.stopReason === "error" && message.errorMessage) {
|
||||
// Skip messages matching known non-overflow patterns (e.g. throttling / rate-limit)
|
||||
const isNonOverflow = NON_OVERFLOW_PATTERNS.some((p) => p.test(message.errorMessage!));
|
||||
if (!isNonOverflow && OVERFLOW_PATTERNS.some((p) => p.test(message.errorMessage!))) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
// Case 2: Silent overflow (z.ai style) - successful but usage exceeds context
|
||||
if (contextWindow && message.stopReason === "stop") {
|
||||
const inputTokens = message.usage.input + message.usage.cacheRead;
|
||||
if (inputTokens > contextWindow) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the overflow patterns for testing purposes.
|
||||
*/
|
||||
export function getOverflowPatterns(): RegExp[] {
|
||||
return [...OVERFLOW_PATTERNS];
|
||||
}
|
||||
@@ -1,25 +0,0 @@
|
||||
/**
|
||||
* Removes unpaired Unicode surrogate characters from a string.
|
||||
*
|
||||
* Unpaired surrogates (high surrogates 0xD800-0xDBFF without matching low surrogates 0xDC00-0xDFFF,
|
||||
* or vice versa) cause JSON serialization errors in many API providers.
|
||||
*
|
||||
* Valid emoji and other characters outside the Basic Multilingual Plane use properly paired
|
||||
* surrogates and will NOT be affected by this function.
|
||||
*
|
||||
* @param text - The text to sanitize
|
||||
* @returns The sanitized text with unpaired surrogates removed
|
||||
*
|
||||
* @example
|
||||
* // Valid emoji (properly paired surrogates) are preserved
|
||||
* sanitizeSurrogates("Hello 🙈 World") // => "Hello 🙈 World"
|
||||
*
|
||||
* // Unpaired high surrogate is removed
|
||||
* const unpaired = String.fromCharCode(0xD83D); // high surrogate without low
|
||||
* sanitizeSurrogates(`Text ${unpaired} here`) // => "Text here"
|
||||
*/
|
||||
export function sanitizeSurrogates(text: string): string {
|
||||
// Replace unpaired high surrogates (0xD800-0xDBFF not followed by low surrogate)
|
||||
// Replace unpaired low surrogates (0xDC00-0xDFFF not preceded by high surrogate)
|
||||
return text.replace(/[\uD800-\uDBFF](?![\uDC00-\uDFFF])|(?<![\uD800-\uDBFF])[\uDC00-\uDFFF]/g, "");
|
||||
}
|
||||
@@ -1,24 +0,0 @@
|
||||
import { type TUnsafe, Type } from "@sinclair/typebox";
|
||||
|
||||
/**
|
||||
* Creates a string enum schema compatible with Google's API and other providers
|
||||
* that don't support anyOf/const patterns.
|
||||
*
|
||||
* @example
|
||||
* const OperationSchema = StringEnum(["add", "subtract", "multiply", "divide"], {
|
||||
* description: "The operation to perform"
|
||||
* });
|
||||
*
|
||||
* type Operation = Static<typeof OperationSchema>; // "add" | "subtract" | "multiply" | "divide"
|
||||
*/
|
||||
export function StringEnum<T extends readonly string[]>(
|
||||
values: T,
|
||||
options?: { description?: string; default?: T[number] },
|
||||
): TUnsafe<T[number]> {
|
||||
return Type.Unsafe<T[number]>({
|
||||
type: "string",
|
||||
enum: values as any,
|
||||
...(options?.description && { description: options.description }),
|
||||
...(options?.default && { default: options.default }),
|
||||
});
|
||||
}
|
||||
@@ -1,93 +0,0 @@
|
||||
import AjvModule from "ajv";
|
||||
import addFormatsModule from "ajv-formats";
|
||||
|
||||
// Handle both default and named exports
|
||||
const Ajv = (AjvModule as any).default || AjvModule;
|
||||
const addFormats = (addFormatsModule as any).default || addFormatsModule;
|
||||
|
||||
import type { Tool, ToolCall } from "../types.js";
|
||||
|
||||
// Detect if we're in a browser extension environment with strict CSP
|
||||
// Chrome extensions with Manifest V3 don't allow eval/Function constructor
|
||||
const isBrowserExtension = typeof globalThis !== "undefined" && (globalThis as any).chrome?.runtime?.id !== undefined;
|
||||
|
||||
function canUseRuntimeCodegen(): boolean {
|
||||
if (isBrowserExtension) {
|
||||
return false;
|
||||
}
|
||||
|
||||
try {
|
||||
new Function("return true;");
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Create a singleton AJV instance with formats only when runtime code generation is available.
|
||||
let ajv: any = null;
|
||||
if (canUseRuntimeCodegen()) {
|
||||
try {
|
||||
ajv = new Ajv({
|
||||
allErrors: true,
|
||||
strict: false,
|
||||
coerceTypes: true,
|
||||
});
|
||||
addFormats(ajv);
|
||||
} catch (_e) {
|
||||
console.warn("AJV validation disabled due to CSP restrictions");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Finds a tool by name and validates the tool call arguments against its TypeBox schema
|
||||
* @param tools Array of tool definitions
|
||||
* @param toolCall The tool call from the LLM
|
||||
* @returns The validated arguments
|
||||
* @throws Error if tool is not found or validation fails
|
||||
*/
|
||||
export function validateToolCall(tools: Tool[], toolCall: ToolCall): any {
|
||||
const tool = tools.find((t) => t.name === toolCall.name);
|
||||
if (!tool) {
|
||||
throw new Error(`Tool "${toolCall.name}" not found`);
|
||||
}
|
||||
return validateToolArguments(tool, toolCall);
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates tool call arguments against the tool's TypeBox schema
|
||||
* @param tool The tool definition with TypeBox schema
|
||||
* @param toolCall The tool call from the LLM
|
||||
* @returns The validated (and potentially coerced) arguments
|
||||
* @throws Error with formatted message if validation fails
|
||||
*/
|
||||
export function validateToolArguments(tool: Tool, toolCall: ToolCall): any {
|
||||
// Skip validation in environments where runtime code generation is unavailable.
|
||||
if (!ajv || !canUseRuntimeCodegen()) {
|
||||
return toolCall.arguments;
|
||||
}
|
||||
|
||||
// Compile the schema.
|
||||
const validate = ajv.compile(tool.parameters);
|
||||
|
||||
// Clone arguments so AJV can safely mutate for type coercion
|
||||
const args = structuredClone(toolCall.arguments);
|
||||
|
||||
// Validate the arguments (AJV mutates args in-place for type coercion)
|
||||
if (validate(args)) {
|
||||
return args;
|
||||
}
|
||||
|
||||
// Format validation errors nicely
|
||||
const errors =
|
||||
validate.errors
|
||||
?.map((err: any) => {
|
||||
const path = err.instancePath ? err.instancePath.substring(1) : err.params.missingProperty || "root";
|
||||
return ` - ${path}: ${err.message}`;
|
||||
})
|
||||
.join("\n") || "Unknown validation error";
|
||||
|
||||
const errorMessage = `Validation failed for tool "${toolCall.name}":\n${errors}\n\nReceived arguments:\n${JSON.stringify(toolCall.arguments, null, 2)}`;
|
||||
|
||||
throw new Error(errorMessage);
|
||||
}
|
||||
@@ -1,260 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { complete, stream } from "../src/stream.js";
|
||||
import type { Api, Context, Model, StreamOptions } from "../src/types.js";
|
||||
|
||||
type StreamOptionsWithExtras = StreamOptions & Record<string, unknown>;
|
||||
|
||||
import { hasAzureOpenAICredentials, resolveAzureDeploymentName } from "./azure-utils.js";
|
||||
import { hasBedrockCredentials } from "./bedrock-utils.js";
|
||||
import { resolveApiKey } from "./oauth.js";
|
||||
|
||||
// Resolve OAuth tokens at module level (async, runs before tests)
|
||||
const [geminiCliToken, openaiCodexToken] = await Promise.all([
|
||||
resolveApiKey("google-gemini-cli"),
|
||||
resolveApiKey("openai-codex"),
|
||||
]);
|
||||
|
||||
async function testAbortSignal<TApi extends Api>(llm: Model<TApi>, options: StreamOptionsWithExtras = {}) {
|
||||
const context: Context = {
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "What is 15 + 27? Think step by step. Then list 50 first names.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
};
|
||||
|
||||
let abortFired = false;
|
||||
let text = "";
|
||||
const controller = new AbortController();
|
||||
const response = await stream(llm, context, { ...options, signal: controller.signal });
|
||||
for await (const event of response) {
|
||||
if (abortFired) return;
|
||||
if (event.type === "text_delta" || event.type === "thinking_delta") {
|
||||
text += event.delta;
|
||||
}
|
||||
if (text.length >= 50) {
|
||||
controller.abort();
|
||||
abortFired = true;
|
||||
}
|
||||
}
|
||||
const msg = await response.result();
|
||||
|
||||
// If we get here without throwing, the abort didn't work
|
||||
expect(msg.stopReason).toBe("aborted");
|
||||
expect(msg.content.length).toBeGreaterThan(0);
|
||||
|
||||
context.messages.push(msg);
|
||||
context.messages.push({
|
||||
role: "user",
|
||||
content: "Please continue, but only generate 5 names.",
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
|
||||
const followUp = await complete(llm, context, options);
|
||||
expect(followUp.stopReason).toBe("stop");
|
||||
expect(followUp.content.length).toBeGreaterThan(0);
|
||||
}
|
||||
|
||||
async function testImmediateAbort<TApi extends Api>(llm: Model<TApi>, options: StreamOptionsWithExtras = {}) {
|
||||
const controller = new AbortController();
|
||||
|
||||
controller.abort();
|
||||
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
const response = await complete(llm, context, { ...options, signal: controller.signal });
|
||||
expect(response.stopReason).toBe("aborted");
|
||||
}
|
||||
|
||||
async function testAbortThenNewMessage<TApi extends Api>(llm: Model<TApi>, options: StreamOptionsWithExtras = {}) {
|
||||
// First request: abort immediately before any response content arrives
|
||||
const controller = new AbortController();
|
||||
controller.abort();
|
||||
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "Hello, how are you?", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
const abortedResponse = await complete(llm, context, { ...options, signal: controller.signal });
|
||||
expect(abortedResponse.stopReason).toBe("aborted");
|
||||
// The aborted message has empty content since we aborted before anything arrived
|
||||
expect(abortedResponse.content.length).toBe(0);
|
||||
|
||||
// Add the aborted assistant message to context (this is what happens in the real coding agent)
|
||||
context.messages.push(abortedResponse);
|
||||
|
||||
// Second request: send a new message - this should work even with the aborted message in context
|
||||
context.messages.push({
|
||||
role: "user",
|
||||
content: "What is 2 + 2?",
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
|
||||
const followUp = await complete(llm, context, options);
|
||||
expect(followUp.stopReason).toBe("stop");
|
||||
expect(followUp.content.length).toBeGreaterThan(0);
|
||||
}
|
||||
|
||||
describe("AI Providers Abort Tests", () => {
|
||||
describe.skipIf(!process.env.GEMINI_API_KEY)("Google Provider Abort", () => {
|
||||
const llm = getModel("google", "gemini-2.5-flash");
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm, { thinking: { enabled: true } });
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm, { thinking: { enabled: true } });
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Completions Provider Abort", () => {
|
||||
const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!;
|
||||
void _compat;
|
||||
const llm: Model<"openai-completions"> = {
|
||||
...baseModel,
|
||||
api: "openai-completions",
|
||||
};
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm);
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Responses Provider Abort", () => {
|
||||
const llm = getModel("openai", "gpt-5-mini");
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm);
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!hasAzureOpenAICredentials())("Azure OpenAI Responses Provider Abort", () => {
|
||||
const llm = getModel("azure-openai-responses", "gpt-4o-mini");
|
||||
const azureDeploymentName = resolveAzureDeploymentName(llm.id);
|
||||
const azureOptions = azureDeploymentName ? { azureDeploymentName } : {};
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm, azureOptions);
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm, azureOptions);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("Anthropic Provider Abort", () => {
|
||||
const llm = getModel("anthropic", "claude-opus-4-1-20250805");
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm, { thinkingEnabled: true, thinkingBudgetTokens: 2048 });
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm, { thinkingEnabled: true, thinkingBudgetTokens: 2048 });
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.MISTRAL_API_KEY)("Mistral Provider Abort", () => {
|
||||
const llm = getModel("mistral", "devstral-medium-latest");
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm);
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.MINIMAX_API_KEY)("MiniMax Provider Abort", () => {
|
||||
const llm = getModel("minimax", "MiniMax-M2.7");
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm);
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider Abort", () => {
|
||||
const llm = getModel("kimi-coding", "kimi-k2-thinking");
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm);
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.AI_GATEWAY_API_KEY)("Vercel AI Gateway Provider Abort", () => {
|
||||
const llm = getModel("vercel-ai-gateway", "google/gemini-2.5-flash");
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm);
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm);
|
||||
});
|
||||
});
|
||||
|
||||
// Google Gemini CLI / Antigravity share the same provider, so one test covers both
|
||||
describe("Google Gemini CLI Provider Abort", () => {
|
||||
it.skipIf(!geminiCliToken)("should abort mid-stream", { retry: 3 }, async () => {
|
||||
const llm = getModel("google-gemini-cli", "gemini-2.5-flash");
|
||||
await testAbortSignal(llm, { apiKey: geminiCliToken });
|
||||
});
|
||||
|
||||
it.skipIf(!geminiCliToken)("should handle immediate abort", { retry: 3 }, async () => {
|
||||
const llm = getModel("google-gemini-cli", "gemini-2.5-flash");
|
||||
await testImmediateAbort(llm, { apiKey: geminiCliToken });
|
||||
});
|
||||
});
|
||||
|
||||
describe("OpenAI Codex Provider Abort", () => {
|
||||
it.skipIf(!openaiCodexToken)("should abort mid-stream", { retry: 3 }, async () => {
|
||||
const llm = getModel("openai-codex", "gpt-5.2-codex");
|
||||
await testAbortSignal(llm, { apiKey: openaiCodexToken });
|
||||
});
|
||||
|
||||
it.skipIf(!openaiCodexToken)("should handle immediate abort", { retry: 3 }, async () => {
|
||||
const llm = getModel("openai-codex", "gpt-5.2-codex");
|
||||
await testImmediateAbort(llm, { apiKey: openaiCodexToken });
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!hasBedrockCredentials())("Amazon Bedrock Provider Abort", () => {
|
||||
const llm = getModel("amazon-bedrock", "global.anthropic.claude-sonnet-4-5-20250929-v1:0");
|
||||
|
||||
it("should abort mid-stream", { retry: 3 }, async () => {
|
||||
await testAbortSignal(llm, { reasoning: "medium" });
|
||||
});
|
||||
|
||||
it("should handle immediate abort", { retry: 3 }, async () => {
|
||||
await testImmediateAbort(llm);
|
||||
});
|
||||
|
||||
it("should handle abort then new message", { retry: 3 }, async () => {
|
||||
await testAbortThenNewMessage(llm);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,99 +0,0 @@
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { loginAnthropic, refreshAnthropicToken } from "../src/utils/oauth/anthropic.js";
|
||||
|
||||
function jsonResponse(body: unknown, status: number = 200): Response {
|
||||
return new Response(JSON.stringify(body), {
|
||||
status,
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function getUrl(input: unknown): string {
|
||||
if (typeof input === "string") {
|
||||
return input;
|
||||
}
|
||||
if (input instanceof URL) {
|
||||
return input.toString();
|
||||
}
|
||||
if (input instanceof Request) {
|
||||
return input.url;
|
||||
}
|
||||
throw new Error(`Unsupported fetch input: ${String(input)}`);
|
||||
}
|
||||
|
||||
function getJsonBody(init?: RequestInit): Record<string, string> {
|
||||
if (typeof init?.body !== "string") {
|
||||
throw new Error(`Expected string request body, got ${typeof init?.body}`);
|
||||
}
|
||||
return JSON.parse(init.body) as Record<string, string>;
|
||||
}
|
||||
|
||||
describe.sequential("Anthropic OAuth", () => {
|
||||
afterEach(() => {
|
||||
vi.unstubAllGlobals();
|
||||
});
|
||||
|
||||
it("keeps the localhost redirect_uri for manual callback login", async () => {
|
||||
let authUrl = "";
|
||||
const fetchMock = vi.fn(async (input: unknown, init?: RequestInit): Promise<Response> => {
|
||||
expect(getUrl(input)).toBe("https://platform.claude.com/v1/oauth/token");
|
||||
expect(init?.method).toBe("POST");
|
||||
const body = getJsonBody(init);
|
||||
expect(body.grant_type).toBe("authorization_code");
|
||||
expect(body.code).toBe("manual-code");
|
||||
expect(body.redirect_uri).toBe("http://localhost:53692/callback");
|
||||
return jsonResponse({
|
||||
access_token: "access-token",
|
||||
refresh_token: "refresh-token",
|
||||
expires_in: 3600,
|
||||
});
|
||||
});
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
const credentials = await loginAnthropic({
|
||||
onAuth: (info) => {
|
||||
authUrl = info.url;
|
||||
},
|
||||
onPrompt: async () => "",
|
||||
onManualCodeInput: async () => {
|
||||
const url = new URL(authUrl);
|
||||
const state = url.searchParams.get("state");
|
||||
const redirectUri = url.searchParams.get("redirect_uri");
|
||||
if (!state || !redirectUri) {
|
||||
throw new Error("Missing OAuth state or redirect_uri in auth URL");
|
||||
}
|
||||
return `${redirectUri}?code=manual-code&state=${state}`;
|
||||
},
|
||||
});
|
||||
|
||||
expect(credentials.access).toBe("access-token");
|
||||
expect(credentials.refresh).toBe("refresh-token");
|
||||
expect(fetchMock).toHaveBeenCalledOnce();
|
||||
});
|
||||
|
||||
it("omits scope from refresh token requests", async () => {
|
||||
const fetchMock = vi.fn(async (input: unknown, init?: RequestInit): Promise<Response> => {
|
||||
expect(getUrl(input)).toBe("https://platform.claude.com/v1/oauth/token");
|
||||
expect(init?.method).toBe("POST");
|
||||
const body = getJsonBody(init);
|
||||
expect(body.grant_type).toBe("refresh_token");
|
||||
expect(body.client_id).toBeTruthy();
|
||||
expect(body.refresh_token).toBe("refresh-token");
|
||||
expect(body).not.toHaveProperty("scope");
|
||||
return jsonResponse({
|
||||
access_token: "new-access-token",
|
||||
refresh_token: "new-refresh-token",
|
||||
expires_in: 3600,
|
||||
});
|
||||
});
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
const credentials = await refreshAnthropicToken("refresh-token");
|
||||
|
||||
expect(credentials.access).toBe("new-access-token");
|
||||
expect(credentials.refresh).toBe("new-refresh-token");
|
||||
expect(fetchMock).toHaveBeenCalledOnce();
|
||||
});
|
||||
});
|
||||
@@ -1,72 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { streamSimple } from "../src/stream.js";
|
||||
import type { Context } from "../src/types.js";
|
||||
|
||||
interface AnthropicThinkingPayload {
|
||||
thinking?: { type: string };
|
||||
output_config?: { effort?: string };
|
||||
}
|
||||
|
||||
function makeContext(): Context {
|
||||
return {
|
||||
systemPrompt: "You are a precise assistant. Follow the user's instructions exactly.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content:
|
||||
"Compute 48291 * 7317 and 90844 - 17729, add the results, and determine whether the sum is divisible by 11. Reply with exactly this format and nothing else: sum=<sum>; divisibleBy11=<yes|no>",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic Opus 4.7 smoke", () => {
|
||||
it("streams Claude Opus 4.7 with reasoning enabled", { retry: 2, timeout: 30000 }, async () => {
|
||||
const model = getModel("anthropic", "claude-opus-4-7");
|
||||
let capturedPayload: AnthropicThinkingPayload | undefined;
|
||||
const s = streamSimple(model, makeContext(), {
|
||||
reasoning: "high",
|
||||
maxTokens: 1024,
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload as AnthropicThinkingPayload;
|
||||
return payload;
|
||||
},
|
||||
});
|
||||
|
||||
let sawThinking = false;
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "thinking_start" || event.type === "thinking_delta" || event.type === "thinking_end") {
|
||||
sawThinking = true;
|
||||
}
|
||||
}
|
||||
|
||||
const response = await s.result();
|
||||
expect(response.stopReason, response.errorMessage).toBe("stop");
|
||||
expect(response.errorMessage).toBeFalsy();
|
||||
expect(capturedPayload?.thinking).toEqual({ type: "adaptive" });
|
||||
expect(capturedPayload?.output_config).toEqual({ effort: "high" });
|
||||
expect(sawThinking).toBe(true);
|
||||
|
||||
const thinkingBlock = response.content.find((block) => block.type === "thinking");
|
||||
expect(thinkingBlock?.type).toBe("thinking");
|
||||
if (!thinkingBlock || thinkingBlock.type !== "thinking") {
|
||||
throw new Error("Expected thinking block from Claude Opus 4.7");
|
||||
}
|
||||
expect(typeof thinkingBlock.thinkingSignature).toBe("string");
|
||||
const thinkingSignature = thinkingBlock.thinkingSignature;
|
||||
if (!thinkingSignature) {
|
||||
throw new Error("Expected thinking signature from Claude Opus 4.7");
|
||||
}
|
||||
expect(thinkingSignature.length).toBeGreaterThan(0);
|
||||
|
||||
const text = response.content
|
||||
.filter((block) => block.type === "text")
|
||||
.map((block) => block.text)
|
||||
.join("")
|
||||
.trim();
|
||||
expect(text).toBe("sum=353418362; divisibleBy11=yes");
|
||||
});
|
||||
});
|
||||
@@ -1,152 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { streamSimple } from "../src/stream.js";
|
||||
import type { Context, Model, SimpleStreamOptions } from "../src/types.js";
|
||||
|
||||
interface AnthropicThinkingPayload {
|
||||
thinking?: { type: string; budget_tokens?: number; display?: string };
|
||||
output_config?: { effort?: string };
|
||||
}
|
||||
|
||||
function makePayloadCaptureContext(): Context {
|
||||
return {
|
||||
messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
|
||||
};
|
||||
}
|
||||
|
||||
async function capturePayload(
|
||||
model: Model<"anthropic-messages">,
|
||||
options?: SimpleStreamOptions,
|
||||
): Promise<AnthropicThinkingPayload> {
|
||||
let capturedPayload: AnthropicThinkingPayload | undefined;
|
||||
const payloadCaptureModel: Model<"anthropic-messages"> = {
|
||||
...model,
|
||||
baseUrl: "http://127.0.0.1:9",
|
||||
};
|
||||
|
||||
const s = streamSimple(payloadCaptureModel, makePayloadCaptureContext(), {
|
||||
...options,
|
||||
apiKey: "fake-key",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload as AnthropicThinkingPayload;
|
||||
return payload;
|
||||
},
|
||||
});
|
||||
|
||||
await s.result();
|
||||
|
||||
if (!capturedPayload) {
|
||||
throw new Error("Expected payload to be captured before request failure");
|
||||
}
|
||||
|
||||
return capturedPayload;
|
||||
}
|
||||
|
||||
interface RunResult {
|
||||
thinkingEventCount: number;
|
||||
thinkingCharCount: number;
|
||||
text: string;
|
||||
contentTypes: string[];
|
||||
}
|
||||
|
||||
function makeE2EContext(): Context {
|
||||
return {
|
||||
systemPrompt: "You are a precise assistant. Follow the requested output format exactly.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content:
|
||||
"Before replying, carefully solve 36863 * 5279 internally. Then reply with the word pong repeated exactly 40 times, separated by single spaces. Do not add any other text.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
function countPongs(text: string): number {
|
||||
return text.match(/\bpong\b/gi)?.length ?? 0;
|
||||
}
|
||||
|
||||
async function runWithoutReasoning(model: Model<"anthropic-messages">): Promise<RunResult> {
|
||||
const s = streamSimple(model, makeE2EContext(), {
|
||||
temperature: 0,
|
||||
maxTokens: 160,
|
||||
});
|
||||
|
||||
let thinkingEventCount = 0;
|
||||
let thinkingCharCount = 0;
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "thinking_start" || event.type === "thinking_end") {
|
||||
thinkingEventCount += 1;
|
||||
}
|
||||
if (event.type === "thinking_delta") {
|
||||
thinkingEventCount += 1;
|
||||
thinkingCharCount += event.delta.length;
|
||||
}
|
||||
}
|
||||
|
||||
const response = await s.result();
|
||||
expect(response.stopReason, response.errorMessage).toBe("stop");
|
||||
|
||||
const text = response.content
|
||||
.filter((block) => block.type === "text")
|
||||
.map((block) => block.text)
|
||||
.join("")
|
||||
.trim();
|
||||
|
||||
return {
|
||||
thinkingEventCount,
|
||||
thinkingCharCount,
|
||||
text,
|
||||
contentTypes: response.content.map((block) => block.type),
|
||||
};
|
||||
}
|
||||
|
||||
describe("Anthropic thinking disable payload", () => {
|
||||
it("sends thinking.type=disabled for budget-based reasoning models when thinking is off", async () => {
|
||||
const payload = await capturePayload(getModel("anthropic", "claude-sonnet-4-5"));
|
||||
|
||||
expect(payload.thinking).toEqual({ type: "disabled" });
|
||||
expect(payload.output_config).toBeUndefined();
|
||||
});
|
||||
|
||||
it("sends thinking.type=disabled for adaptive reasoning models when thinking is off", async () => {
|
||||
const payload = await capturePayload(getModel("anthropic", "claude-opus-4-6"));
|
||||
|
||||
expect(payload.thinking).toEqual({ type: "disabled" });
|
||||
expect(payload.output_config).toBeUndefined();
|
||||
});
|
||||
|
||||
it("sends thinking.type=disabled for Claude Opus 4.7 when thinking is off", async () => {
|
||||
const payload = await capturePayload(getModel("anthropic", "claude-opus-4-7"));
|
||||
|
||||
expect(payload.thinking).toEqual({ type: "disabled" });
|
||||
expect(payload.output_config).toBeUndefined();
|
||||
});
|
||||
|
||||
it("uses adaptive thinking for Claude Opus 4.7 when reasoning is enabled", async () => {
|
||||
const payload = await capturePayload(getModel("anthropic", "claude-opus-4-7"), { reasoning: "high" });
|
||||
|
||||
expect(payload.thinking).toEqual({ type: "adaptive", display: "summarized" });
|
||||
expect(payload.output_config).toEqual({ effort: "high" });
|
||||
});
|
||||
|
||||
it("maps xhigh reasoning to effort=xhigh for Claude Opus 4.7", async () => {
|
||||
const payload = await capturePayload(getModel("anthropic", "claude-opus-4-7"), { reasoning: "xhigh" });
|
||||
|
||||
expect(payload.thinking).toEqual({ type: "adaptive", display: "summarized" });
|
||||
expect(payload.output_config).toEqual({ effort: "xhigh" });
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic thinking disable E2E", () => {
|
||||
it("disables thinking for Claude reasoning models", { retry: 2, timeout: 30000 }, async () => {
|
||||
const result = await runWithoutReasoning(getModel("anthropic", "claude-sonnet-4-5"));
|
||||
|
||||
expect(result.thinkingEventCount).toBe(0);
|
||||
expect(result.thinkingCharCount).toBe(0);
|
||||
expect(result.contentTypes).not.toContain("thinking");
|
||||
expect(countPongs(result.text)).toBeGreaterThanOrEqual(35);
|
||||
});
|
||||
});
|
||||
@@ -1,205 +0,0 @@
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { stream } from "../src/stream.js";
|
||||
import type { Context, Tool } from "../src/types.js";
|
||||
import { resolveApiKey } from "./oauth.js";
|
||||
|
||||
const oauthToken = await resolveApiKey("anthropic");
|
||||
|
||||
/**
|
||||
* Tests for Anthropic OAuth tool name normalization.
|
||||
*
|
||||
* When using Claude Code OAuth, tool names must match CC's canonical casing.
|
||||
* The normalization should:
|
||||
* 1. Convert tool names that match CC tools (case-insensitive) to CC casing on outbound
|
||||
* 2. Convert tool names back to the original casing on inbound
|
||||
*
|
||||
* This is a simple case-insensitive lookup, NOT a mapping of different names.
|
||||
* e.g., "todowrite" -> "TodoWrite" -> "todowrite" (round-trip works)
|
||||
*
|
||||
* The old `find -> Glob` mapping was WRONG because:
|
||||
* - Outbound: "find" -> "Glob"
|
||||
* - Inbound: "Glob" -> ??? (no tool named "glob" in context.tools, only "find")
|
||||
* - Result: tool call has name "Glob" but no tool exists with that name
|
||||
*/
|
||||
describe.skipIf(!oauthToken)("Anthropic OAuth tool name normalization", () => {
|
||||
const model = getModel("anthropic", "claude-sonnet-4-20250514");
|
||||
|
||||
it("should normalize user-defined tool matching CC name (todowrite -> TodoWrite -> todowrite)", async () => {
|
||||
// User defines a tool named "todowrite" (lowercase)
|
||||
// CC has "TodoWrite" - this should round-trip correctly
|
||||
const todoTool: Tool = {
|
||||
name: "todowrite",
|
||||
description: "Write a todo item",
|
||||
parameters: Type.Object({
|
||||
task: Type.String({ description: "The task to add" }),
|
||||
}),
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant. Use the todowrite tool when asked to add todos.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "Add a todo: buy milk. Use the todowrite tool.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
tools: [todoTool],
|
||||
};
|
||||
|
||||
const s = stream(model, context, { apiKey: oauthToken });
|
||||
let toolCallName: string | undefined;
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "toolcall_end") {
|
||||
const toolCall = event.partial.content[event.contentIndex];
|
||||
if (toolCall.type === "toolCall") {
|
||||
toolCallName = toolCall.name;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const response = await s.result();
|
||||
expect(response.stopReason, `Error: ${response.errorMessage}`).toBe("toolUse");
|
||||
|
||||
// The tool call should come back with the ORIGINAL name "todowrite", not "TodoWrite"
|
||||
expect(toolCallName).toBe("todowrite");
|
||||
});
|
||||
|
||||
it("should handle pi's built-in tools (read, write, edit, bash)", async () => {
|
||||
// Pi's tools use lowercase names, CC uses PascalCase
|
||||
const readTool: Tool = {
|
||||
name: "read",
|
||||
description: "Read a file",
|
||||
parameters: Type.Object({
|
||||
path: Type.String({ description: "File path" }),
|
||||
}),
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant. Use the read tool to read files.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "Read the file /tmp/test.txt using the read tool.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
tools: [readTool],
|
||||
};
|
||||
|
||||
const s = stream(model, context, { apiKey: oauthToken });
|
||||
let toolCallName: string | undefined;
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "toolcall_end") {
|
||||
const toolCall = event.partial.content[event.contentIndex];
|
||||
if (toolCall.type === "toolCall") {
|
||||
toolCallName = toolCall.name;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const response = await s.result();
|
||||
expect(response.stopReason, `Error: ${response.errorMessage}`).toBe("toolUse");
|
||||
|
||||
// The tool call should come back with the ORIGINAL name "read", not "Read"
|
||||
expect(toolCallName).toBe("read");
|
||||
});
|
||||
|
||||
it("should NOT map find to Glob - find is not a CC tool name", async () => {
|
||||
// Pi has a "find" tool, CC has "Glob" - these are DIFFERENT tools
|
||||
// The old code incorrectly mapped find -> Glob, which broke the round-trip
|
||||
// because there's no tool named "glob" in context.tools
|
||||
const findTool: Tool = {
|
||||
name: "find",
|
||||
description: "Find files by pattern",
|
||||
parameters: Type.Object({
|
||||
pattern: Type.String({ description: "Glob pattern" }),
|
||||
}),
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant. Use the find tool to search for files.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "Find all .ts files using the find tool.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
tools: [findTool],
|
||||
};
|
||||
|
||||
const s = stream(model, context, { apiKey: oauthToken });
|
||||
let toolCallName: string | undefined;
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "toolcall_end") {
|
||||
const toolCall = event.partial.content[event.contentIndex];
|
||||
if (toolCall.type === "toolCall") {
|
||||
toolCallName = toolCall.name;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const response = await s.result();
|
||||
expect(response.stopReason, `Error: ${response.errorMessage}`).toBe("toolUse");
|
||||
|
||||
// With the BROKEN find -> Glob mapping:
|
||||
// - Sent as "Glob" to Anthropic
|
||||
// - Received back as "Glob"
|
||||
// - fromClaudeCodeName("Glob", tools) looks for tool.name.toLowerCase() === "glob"
|
||||
// - No match (tool is named "find"), returns "Glob"
|
||||
// - Test fails: toolCallName is "Glob" instead of "find"
|
||||
//
|
||||
// With the CORRECT implementation (no find->Glob mapping):
|
||||
// - Sent as "find" to Anthropic (no CC tool named "Find")
|
||||
// - Received back as "find"
|
||||
// - Test passes: toolCallName is "find"
|
||||
expect(toolCallName).toBe("find");
|
||||
});
|
||||
|
||||
it("should handle custom tools that don't match any CC tool names", async () => {
|
||||
// A completely custom tool should pass through unchanged
|
||||
const customTool: Tool = {
|
||||
name: "my_custom_tool",
|
||||
description: "A custom tool",
|
||||
parameters: Type.Object({
|
||||
input: Type.String({ description: "Input value" }),
|
||||
}),
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant. Use my_custom_tool when asked.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "Use my_custom_tool with input 'hello'.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
tools: [customTool],
|
||||
};
|
||||
|
||||
const s = stream(model, context, { apiKey: oauthToken });
|
||||
let toolCallName: string | undefined;
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "toolcall_end") {
|
||||
const toolCall = event.partial.content[event.contentIndex];
|
||||
if (toolCall.type === "toolCall") {
|
||||
toolCallName = toolCall.name;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const response = await s.result();
|
||||
expect(response.stopReason, `Error: ${response.errorMessage}`).toBe("toolUse");
|
||||
|
||||
// Custom tool names should pass through unchanged
|
||||
expect(toolCallName).toBe("my_custom_tool");
|
||||
});
|
||||
});
|
||||
@@ -1,28 +0,0 @@
|
||||
/**
|
||||
* Utility functions for Azure OpenAI tests
|
||||
*/
|
||||
|
||||
function parseDeploymentNameMap(value: string | undefined): Map<string, string> {
|
||||
const map = new Map<string, string>();
|
||||
if (!value) return map;
|
||||
for (const entry of value.split(",")) {
|
||||
const trimmed = entry.trim();
|
||||
if (!trimmed) continue;
|
||||
const [modelId, deploymentName] = trimmed.split("=", 2);
|
||||
if (!modelId || !deploymentName) continue;
|
||||
map.set(modelId.trim(), deploymentName.trim());
|
||||
}
|
||||
return map;
|
||||
}
|
||||
|
||||
export function hasAzureOpenAICredentials(): boolean {
|
||||
const hasKey = !!process.env.AZURE_OPENAI_API_KEY;
|
||||
const hasBaseUrl = !!(process.env.AZURE_OPENAI_BASE_URL || process.env.AZURE_OPENAI_RESOURCE_NAME);
|
||||
return hasKey && hasBaseUrl;
|
||||
}
|
||||
|
||||
export function resolveAzureDeploymentName(modelId: string): string | undefined {
|
||||
const mapValue = process.env.AZURE_OPENAI_DEPLOYMENT_NAME_MAP;
|
||||
if (!mapValue) return undefined;
|
||||
return parseDeploymentNameMap(mapValue).get(modelId);
|
||||
}
|
||||
@@ -1,66 +0,0 @@
|
||||
/**
|
||||
* A test suite to ensure all configured Amazon Bedrock models are usable.
|
||||
*
|
||||
* This is here to make sure we got correct model identifiers from models.dev and other sources.
|
||||
* Because Amazon Bedrock requires cross-region inference in some models,
|
||||
* plain model identifiers are not always usable and it requires tweaking of model identifiers to use cross-region inference.
|
||||
* See https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html#inference-profiles-support-system for more details.
|
||||
*
|
||||
* This test suite is not enabled by default unless AWS credentials and `BEDROCK_EXTENSIVE_MODEL_TEST` environment variables are set.
|
||||
* This test suite takes ~2 minutes to run. Because not all models are available in all regions,
|
||||
* it's recommended to use `us-west-2` region for best coverage for running this test suite.
|
||||
*
|
||||
* You can run this test suite with:
|
||||
* ```bash
|
||||
* $ AWS_REGION=us-west-2 BEDROCK_EXTENSIVE_MODEL_TEST=1 AWS_PROFILE=... npm test -- ./test/bedrock-models.test.ts
|
||||
* ```
|
||||
*/
|
||||
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModels } from "../src/models.js";
|
||||
import { complete } from "../src/stream.js";
|
||||
import type { Context } from "../src/types.js";
|
||||
import { hasBedrockCredentials } from "./bedrock-utils.js";
|
||||
|
||||
describe("Amazon Bedrock Models", () => {
|
||||
const models = getModels("amazon-bedrock");
|
||||
|
||||
it("should get all available Bedrock models", () => {
|
||||
expect(models.length).toBeGreaterThan(0);
|
||||
console.log(`Found ${models.length} Bedrock models`);
|
||||
});
|
||||
|
||||
if (hasBedrockCredentials() && process.env.BEDROCK_EXTENSIVE_MODEL_TEST) {
|
||||
for (const model of models) {
|
||||
it(`should make a simple request with ${model.id}`, { timeout: 10_000 }, async () => {
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant. Be extremely concise.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "Reply with exactly: 'OK'",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const response = await complete(model, context);
|
||||
|
||||
expect(response.role).toBe("assistant");
|
||||
expect(response.content).toBeTruthy();
|
||||
expect(response.content.length).toBeGreaterThan(0);
|
||||
expect(response.usage.input + response.usage.cacheRead).toBeGreaterThan(0);
|
||||
expect(response.usage.output).toBeGreaterThan(0);
|
||||
expect(response.errorMessage).toBeFalsy();
|
||||
|
||||
const textContent = response.content
|
||||
.filter((b) => b.type === "text")
|
||||
.map((b) => (b.type === "text" ? b.text : ""))
|
||||
.join("")
|
||||
.trim();
|
||||
expect(textContent).toBeTruthy();
|
||||
console.log(`${model.id}: ${textContent.substring(0, 100)}`);
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
@@ -1,107 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { type BedrockOptions, streamBedrock } from "../src/providers/amazon-bedrock.js";
|
||||
import type { Context, Model } from "../src/types.js";
|
||||
|
||||
interface BedrockThinkingPayload {
|
||||
additionalModelRequestFields?: {
|
||||
thinking?: { type: string; budget_tokens?: number; display?: string };
|
||||
output_config?: { effort?: string };
|
||||
anthropic_beta?: string[];
|
||||
};
|
||||
}
|
||||
|
||||
function makeContext(): Context {
|
||||
return {
|
||||
messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
|
||||
};
|
||||
}
|
||||
|
||||
async function capturePayload(
|
||||
model: Model<"bedrock-converse-stream">,
|
||||
options?: BedrockOptions,
|
||||
): Promise<BedrockThinkingPayload> {
|
||||
let capturedPayload: BedrockThinkingPayload | undefined;
|
||||
const s = streamBedrock(model, makeContext(), {
|
||||
...options,
|
||||
reasoning: options?.reasoning ?? "high",
|
||||
signal: AbortSignal.abort(),
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload as BedrockThinkingPayload;
|
||||
return payload;
|
||||
},
|
||||
});
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (!capturedPayload) {
|
||||
throw new Error("Expected Bedrock payload to be captured before request abort");
|
||||
}
|
||||
|
||||
return capturedPayload;
|
||||
}
|
||||
|
||||
describe("Bedrock thinking payload", () => {
|
||||
it("uses adaptive thinking for Claude Opus 4.7 when reasoning is enabled", async () => {
|
||||
const baseModel = getModel("amazon-bedrock", "global.anthropic.claude-opus-4-6-v1");
|
||||
const model: Model<"bedrock-converse-stream"> = {
|
||||
...baseModel,
|
||||
id: "global.anthropic.claude-opus-4-7-v1",
|
||||
name: "Claude Opus 4.7 (Global)",
|
||||
};
|
||||
|
||||
const payload = await capturePayload(model);
|
||||
|
||||
expect(payload.additionalModelRequestFields?.thinking).toEqual({ type: "adaptive", display: "summarized" });
|
||||
expect(payload.additionalModelRequestFields?.output_config).toEqual({ effort: "high" });
|
||||
expect(payload.additionalModelRequestFields?.anthropic_beta).toBeUndefined();
|
||||
});
|
||||
|
||||
it("maps xhigh reasoning to effort=xhigh for Claude Opus 4.7", async () => {
|
||||
const baseModel = getModel("amazon-bedrock", "global.anthropic.claude-opus-4-6-v1");
|
||||
const model: Model<"bedrock-converse-stream"> = {
|
||||
...baseModel,
|
||||
id: "global.anthropic.claude-opus-4-7-v1",
|
||||
name: "Claude Opus 4.7 (Global)",
|
||||
};
|
||||
|
||||
const payload = await capturePayload(model, { reasoning: "xhigh" });
|
||||
|
||||
expect(payload.additionalModelRequestFields?.thinking).toEqual({ type: "adaptive", display: "summarized" });
|
||||
expect(payload.additionalModelRequestFields?.output_config).toEqual({ effort: "xhigh" });
|
||||
expect(payload.additionalModelRequestFields?.anthropic_beta).toBeUndefined();
|
||||
});
|
||||
|
||||
it("omits display for GovCloud model ids on non-adaptive Claude thinking", async () => {
|
||||
const baseModel = getModel("amazon-bedrock", "us.anthropic.claude-sonnet-4-5-20250929-v1:0");
|
||||
const model: Model<"bedrock-converse-stream"> = {
|
||||
...baseModel,
|
||||
id: "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
name: "Claude Sonnet 4.5 (GovCloud)",
|
||||
};
|
||||
|
||||
const payload = await capturePayload(model);
|
||||
|
||||
expect(payload.additionalModelRequestFields?.thinking).toEqual({ type: "enabled", budget_tokens: 16384 });
|
||||
expect(payload.additionalModelRequestFields?.anthropic_beta).toEqual(["interleaved-thinking-2025-05-14"]);
|
||||
});
|
||||
|
||||
it("omits display for GovCloud regions on adaptive Claude thinking", async () => {
|
||||
const baseModel = getModel("amazon-bedrock", "global.anthropic.claude-opus-4-6-v1");
|
||||
const model: Model<"bedrock-converse-stream"> = {
|
||||
...baseModel,
|
||||
id: "global.anthropic.claude-opus-4-7-v1",
|
||||
name: "Claude Opus 4.7 (Global)",
|
||||
};
|
||||
|
||||
const payload = await capturePayload(model, { region: "us-gov-west-1" });
|
||||
|
||||
expect(payload.additionalModelRequestFields?.thinking).toEqual({ type: "adaptive" });
|
||||
expect(payload.additionalModelRequestFields?.output_config).toEqual({ effort: "high" });
|
||||
expect(payload.additionalModelRequestFields?.anthropic_beta).toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -1,18 +0,0 @@
|
||||
/**
|
||||
* Utility functions for Amazon Bedrock tests
|
||||
*/
|
||||
|
||||
/**
|
||||
* Check if any valid AWS credentials are configured for Bedrock.
|
||||
* Returns true if any of the following are set:
|
||||
* - AWS_PROFILE (named profile from ~/.aws/credentials)
|
||||
* - AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY (IAM keys)
|
||||
* - AWS_BEARER_TOKEN_BEDROCK (Bedrock API key)
|
||||
*/
|
||||
export function hasBedrockCredentials(): boolean {
|
||||
return !!(
|
||||
process.env.AWS_PROFILE ||
|
||||
(process.env.AWS_ACCESS_KEY_ID && process.env.AWS_SECRET_ACCESS_KEY) ||
|
||||
process.env.AWS_BEARER_TOKEN_BEDROCK
|
||||
);
|
||||
}
|
||||
@@ -1,335 +0,0 @@
|
||||
import { afterEach, beforeEach, describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { stream } from "../src/stream.js";
|
||||
import type { Context } from "../src/types.js";
|
||||
|
||||
describe("Cache Retention (PI_CACHE_RETENTION)", () => {
|
||||
const originalEnv = process.env.PI_CACHE_RETENTION;
|
||||
|
||||
beforeEach(() => {
|
||||
delete process.env.PI_CACHE_RETENTION;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalEnv !== undefined) {
|
||||
process.env.PI_CACHE_RETENTION = originalEnv;
|
||||
} else {
|
||||
delete process.env.PI_CACHE_RETENTION;
|
||||
}
|
||||
});
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
describe("Anthropic Provider", () => {
|
||||
it.skipIf(!process.env.ANTHROPIC_API_KEY)(
|
||||
"should use default cache TTL (no ttl field) when PI_CACHE_RETENTION is not set",
|
||||
async () => {
|
||||
const model = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const s = stream(model, context, {
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
// Consume the stream to trigger the request
|
||||
for await (const _ of s) {
|
||||
// Just consume
|
||||
}
|
||||
|
||||
expect(capturedPayload).not.toBeNull();
|
||||
// System prompt should have cache_control without ttl
|
||||
expect(capturedPayload.system).toBeDefined();
|
||||
expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral" });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!process.env.ANTHROPIC_API_KEY)("should use 1h cache TTL when PI_CACHE_RETENTION=long", async () => {
|
||||
process.env.PI_CACHE_RETENTION = "long";
|
||||
const model = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const s = stream(model, context, {
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
// Consume the stream to trigger the request
|
||||
for await (const _ of s) {
|
||||
// Just consume
|
||||
}
|
||||
|
||||
expect(capturedPayload).not.toBeNull();
|
||||
// System prompt should have cache_control with ttl: "1h"
|
||||
expect(capturedPayload.system).toBeDefined();
|
||||
expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral", ttl: "1h" });
|
||||
});
|
||||
|
||||
it("should not add ttl when baseUrl is not api.anthropic.com", async () => {
|
||||
process.env.PI_CACHE_RETENTION = "long";
|
||||
|
||||
// Create a model with a different baseUrl (simulating a proxy)
|
||||
const baseModel = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
const proxyModel = {
|
||||
...baseModel,
|
||||
baseUrl: "https://my-proxy.example.com/v1",
|
||||
};
|
||||
|
||||
let capturedPayload: any = null;
|
||||
|
||||
// We can't actually make the request (no proxy), but we can verify the payload
|
||||
// by using a mock or checking the logic directly
|
||||
// For this test, we'll import the helper directly
|
||||
|
||||
// Since we can't easily test this without mocking, we'll skip the actual API call
|
||||
// and just verify the helper logic works correctly
|
||||
const { streamAnthropic } = await import("../src/providers/anthropic.js");
|
||||
|
||||
try {
|
||||
const s = streamAnthropic(proxyModel, context, {
|
||||
apiKey: "fake-key",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
// This will fail since we're using a fake key and fake proxy, but the payload should be captured
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") break;
|
||||
}
|
||||
} catch {
|
||||
// Expected to fail
|
||||
}
|
||||
|
||||
// The payload should have been captured before the error
|
||||
if (capturedPayload) {
|
||||
// System prompt should have cache_control WITHOUT ttl (proxy URL)
|
||||
expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral" });
|
||||
}
|
||||
});
|
||||
|
||||
it("should omit cache_control when cacheRetention is none", async () => {
|
||||
const baseModel = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const { streamAnthropic } = await import("../src/providers/anthropic.js");
|
||||
|
||||
try {
|
||||
const s = streamAnthropic(baseModel, context, {
|
||||
apiKey: "fake-key",
|
||||
cacheRetention: "none",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") break;
|
||||
}
|
||||
} catch {
|
||||
// Expected to fail
|
||||
}
|
||||
|
||||
expect(capturedPayload).not.toBeNull();
|
||||
expect(capturedPayload.system[0].cache_control).toBeUndefined();
|
||||
});
|
||||
|
||||
it("should add cache_control to string user messages", async () => {
|
||||
const baseModel = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const { streamAnthropic } = await import("../src/providers/anthropic.js");
|
||||
|
||||
try {
|
||||
const s = streamAnthropic(baseModel, context, {
|
||||
apiKey: "fake-key",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") break;
|
||||
}
|
||||
} catch {
|
||||
// Expected to fail
|
||||
}
|
||||
|
||||
expect(capturedPayload).not.toBeNull();
|
||||
const lastMessage = capturedPayload.messages[capturedPayload.messages.length - 1];
|
||||
expect(Array.isArray(lastMessage.content)).toBe(true);
|
||||
const lastBlock = lastMessage.content[lastMessage.content.length - 1];
|
||||
expect(lastBlock.cache_control).toEqual({ type: "ephemeral" });
|
||||
});
|
||||
|
||||
it("should set 1h cache TTL when cacheRetention is long", async () => {
|
||||
const baseModel = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const { streamAnthropic } = await import("../src/providers/anthropic.js");
|
||||
|
||||
try {
|
||||
const s = streamAnthropic(baseModel, context, {
|
||||
apiKey: "fake-key",
|
||||
cacheRetention: "long",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") break;
|
||||
}
|
||||
} catch {
|
||||
// Expected to fail
|
||||
}
|
||||
|
||||
expect(capturedPayload).not.toBeNull();
|
||||
expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral", ttl: "1h" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("OpenAI Responses Provider", () => {
|
||||
it.skipIf(!process.env.OPENAI_API_KEY)(
|
||||
"should not set prompt_cache_retention when PI_CACHE_RETENTION is not set",
|
||||
async () => {
|
||||
const model = getModel("openai", "gpt-4o-mini");
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const s = stream(model, context, {
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
// Consume the stream to trigger the request
|
||||
for await (const _ of s) {
|
||||
// Just consume
|
||||
}
|
||||
|
||||
expect(capturedPayload).not.toBeNull();
|
||||
expect(capturedPayload.prompt_cache_retention).toBeUndefined();
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!process.env.OPENAI_API_KEY)(
|
||||
"should set prompt_cache_retention to 24h when PI_CACHE_RETENTION=long",
|
||||
async () => {
|
||||
process.env.PI_CACHE_RETENTION = "long";
|
||||
const model = getModel("openai", "gpt-4o-mini");
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const s = stream(model, context, {
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
// Consume the stream to trigger the request
|
||||
for await (const _ of s) {
|
||||
// Just consume
|
||||
}
|
||||
|
||||
expect(capturedPayload).not.toBeNull();
|
||||
expect(capturedPayload.prompt_cache_retention).toBe("24h");
|
||||
},
|
||||
);
|
||||
|
||||
it("should not set prompt_cache_retention when baseUrl is not api.openai.com", async () => {
|
||||
process.env.PI_CACHE_RETENTION = "long";
|
||||
|
||||
// Create a model with a different baseUrl (simulating a proxy)
|
||||
const baseModel = getModel("openai", "gpt-4o-mini");
|
||||
const proxyModel = {
|
||||
...baseModel,
|
||||
baseUrl: "https://my-proxy.example.com/v1",
|
||||
};
|
||||
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");
|
||||
|
||||
try {
|
||||
const s = streamOpenAIResponses(proxyModel, context, {
|
||||
apiKey: "fake-key",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
// This will fail since we're using a fake key and fake proxy, but the payload should be captured
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") break;
|
||||
}
|
||||
} catch {
|
||||
// Expected to fail
|
||||
}
|
||||
|
||||
// The payload should have been captured before the error
|
||||
if (capturedPayload) {
|
||||
expect(capturedPayload.prompt_cache_retention).toBeUndefined();
|
||||
}
|
||||
});
|
||||
|
||||
it("should omit prompt_cache_key when cacheRetention is none", async () => {
|
||||
const model = getModel("openai", "gpt-4o-mini");
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");
|
||||
|
||||
try {
|
||||
const s = streamOpenAIResponses(model, context, {
|
||||
apiKey: "fake-key",
|
||||
cacheRetention: "none",
|
||||
sessionId: "session-1",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") break;
|
||||
}
|
||||
} catch {
|
||||
// Expected to fail
|
||||
}
|
||||
|
||||
expect(capturedPayload).not.toBeNull();
|
||||
expect(capturedPayload.prompt_cache_key).toBeUndefined();
|
||||
expect(capturedPayload.prompt_cache_retention).toBeUndefined();
|
||||
});
|
||||
|
||||
it("should set prompt_cache_retention when cacheRetention is long", async () => {
|
||||
const model = getModel("openai", "gpt-4o-mini");
|
||||
let capturedPayload: any = null;
|
||||
|
||||
const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");
|
||||
|
||||
try {
|
||||
const s = streamOpenAIResponses(model, context, {
|
||||
apiKey: "fake-key",
|
||||
cacheRetention: "long",
|
||||
sessionId: "session-2",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload;
|
||||
},
|
||||
});
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") break;
|
||||
}
|
||||
} catch {
|
||||
// Expected to fail
|
||||
}
|
||||
|
||||
expect(capturedPayload).not.toBeNull();
|
||||
expect(capturedPayload.prompt_cache_key).toBe("session-2");
|
||||
expect(capturedPayload.prompt_cache_retention).toBe("24h");
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,714 +0,0 @@
|
||||
/**
|
||||
* Test context overflow error handling across providers.
|
||||
*
|
||||
* Context overflow occurs when the input (prompt + history) exceeds
|
||||
* the model's context window. This is different from output token limits.
|
||||
*
|
||||
* Expected behavior: All providers should return stopReason: "error"
|
||||
* with an errorMessage that indicates the context was too large,
|
||||
* OR (for z.ai) return successfully with usage.input > contextWindow.
|
||||
*
|
||||
* The isContextOverflow() function must return true for all providers.
|
||||
*/
|
||||
|
||||
import type { ChildProcess } from "child_process";
|
||||
import { execSync, spawn } from "child_process";
|
||||
import { afterAll, beforeAll, describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { complete } from "../src/stream.js";
|
||||
import type { AssistantMessage, Context, Model, Usage } from "../src/types.js";
|
||||
import { isContextOverflow } from "../src/utils/overflow.js";
|
||||
import { hasAzureOpenAICredentials } from "./azure-utils.js";
|
||||
import { hasBedrockCredentials } from "./bedrock-utils.js";
|
||||
import { resolveApiKey } from "./oauth.js";
|
||||
|
||||
// Resolve OAuth tokens at module level (async, runs before tests)
|
||||
const oauthTokens = await Promise.all([
|
||||
resolveApiKey("github-copilot"),
|
||||
resolveApiKey("google-gemini-cli"),
|
||||
resolveApiKey("google-antigravity"),
|
||||
resolveApiKey("openai-codex"),
|
||||
]);
|
||||
const [githubCopilotToken, geminiCliToken, antigravityToken, openaiCodexToken] = oauthTokens;
|
||||
|
||||
// Lorem ipsum paragraph for realistic token estimation
|
||||
const LOREM_IPSUM = `Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. `;
|
||||
|
||||
// Generate a string that will exceed the context window
|
||||
// Using chars/4 as token estimate (works better with varied text than repeated chars)
|
||||
function generateOverflowContent(contextWindow: number): string {
|
||||
const targetTokens = contextWindow + 10000; // Exceed by 10k tokens
|
||||
const targetChars = targetTokens * 4 * 1.5;
|
||||
const repetitions = Math.ceil(targetChars / LOREM_IPSUM.length);
|
||||
return LOREM_IPSUM.repeat(repetitions);
|
||||
}
|
||||
|
||||
interface OverflowResult {
|
||||
provider: string;
|
||||
model: string;
|
||||
contextWindow: number;
|
||||
stopReason: string;
|
||||
errorMessage: string | undefined;
|
||||
usage: Usage;
|
||||
hasUsageData: boolean;
|
||||
response: AssistantMessage;
|
||||
}
|
||||
|
||||
async function testContextOverflow(model: Model<any>, apiKey: string): Promise<OverflowResult> {
|
||||
const overflowContent = generateOverflowContent(model.contextWindow);
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: overflowContent,
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const response = await complete(model, context, { apiKey });
|
||||
|
||||
const hasUsageData = response.usage.input > 0 || response.usage.cacheRead > 0;
|
||||
|
||||
return {
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
contextWindow: model.contextWindow,
|
||||
stopReason: response.stopReason,
|
||||
errorMessage: response.errorMessage,
|
||||
usage: response.usage,
|
||||
hasUsageData,
|
||||
response,
|
||||
};
|
||||
}
|
||||
|
||||
function logResult(result: OverflowResult) {
|
||||
console.log(`\n${result.provider} / ${result.model}:`);
|
||||
console.log(` contextWindow: ${result.contextWindow}`);
|
||||
console.log(` stopReason: ${result.stopReason}`);
|
||||
console.log(` errorMessage: ${result.errorMessage}`);
|
||||
console.log(` usage: ${JSON.stringify(result.usage)}`);
|
||||
console.log(` hasUsageData: ${result.hasUsageData}`);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Anthropic
|
||||
// Expected pattern: "prompt is too long: X tokens > Y maximum"
|
||||
// =============================================================================
|
||||
|
||||
describe("Context overflow error handling", () => {
|
||||
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic (API Key)", () => {
|
||||
it("claude-3-5-haiku - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
const result = await testContextOverflow(model, process.env.ANTHROPIC_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/prompt is too long/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("Anthropic (OAuth)", () => {
|
||||
it("claude-sonnet-4 - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("anthropic", "claude-sonnet-4-20250514");
|
||||
const result = await testContextOverflow(model, process.env.ANTHROPIC_OAUTH_TOKEN!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/prompt is too long/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// GitHub Copilot (OAuth)
|
||||
// Tests both OpenAI and Anthropic models via Copilot
|
||||
// =============================================================================
|
||||
|
||||
describe("GitHub Copilot (OAuth)", () => {
|
||||
// OpenAI model via Copilot
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"gpt-4o - should detect overflow via isContextOverflow",
|
||||
async () => {
|
||||
const model = getModel("github-copilot", "gpt-4o");
|
||||
const result = await testContextOverflow(model, githubCopilotToken!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/exceeds the limit of \d+/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
},
|
||||
120000,
|
||||
);
|
||||
|
||||
// Anthropic model via Copilot
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"claude-sonnet-4 - should detect overflow via isContextOverflow",
|
||||
async () => {
|
||||
const model = getModel("github-copilot", "claude-sonnet-4");
|
||||
const result = await testContextOverflow(model, githubCopilotToken!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/exceeds the limit of \d+|input is too long/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
},
|
||||
120000,
|
||||
);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// OpenAI
|
||||
// Expected pattern: "exceeds the context window"
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Completions", () => {
|
||||
it("gpt-4o-mini - should detect overflow via isContextOverflow", async () => {
|
||||
const model = { ...getModel("openai", "gpt-4o-mini") };
|
||||
model.api = "openai-completions" as any;
|
||||
const result = await testContextOverflow(model, process.env.OPENAI_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/maximum context length/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Responses", () => {
|
||||
it("gpt-4o - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("openai", "gpt-4o");
|
||||
const result = await testContextOverflow(model, process.env.OPENAI_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/exceeds the context window/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
describe.skipIf(!hasAzureOpenAICredentials())("Azure OpenAI Responses", () => {
|
||||
it("gpt-4o-mini - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("azure-openai-responses", "gpt-4o-mini");
|
||||
const result = await testContextOverflow(model, process.env.AZURE_OPENAI_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/context|maximum/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Google
|
||||
// Expected pattern: "input token count (X) exceeds the maximum"
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.GEMINI_API_KEY)("Google", () => {
|
||||
it("gemini-2.0-flash - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("google", "gemini-2.0-flash");
|
||||
const result = await testContextOverflow(model, process.env.GEMINI_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/input token count.*exceeds the maximum/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Google Gemini CLI (OAuth)
|
||||
// Uses same API as Google, expects same error pattern
|
||||
// =============================================================================
|
||||
|
||||
describe("Google Gemini CLI (OAuth)", () => {
|
||||
it.skipIf(!geminiCliToken)(
|
||||
"gemini-2.5-flash - should detect overflow via isContextOverflow",
|
||||
async () => {
|
||||
const model = getModel("google-gemini-cli", "gemini-2.5-flash");
|
||||
const result = await testContextOverflow(model, geminiCliToken!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/input token count.*exceeds the maximum/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
},
|
||||
120000,
|
||||
);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Google Antigravity (OAuth)
|
||||
// Tests both Gemini and Anthropic models via Antigravity
|
||||
// =============================================================================
|
||||
|
||||
describe("Google Antigravity (OAuth)", () => {
|
||||
// Gemini model
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gemini-3-flash - should detect overflow via isContextOverflow",
|
||||
async () => {
|
||||
const model = getModel("google-antigravity", "gemini-3-flash");
|
||||
const result = await testContextOverflow(model, antigravityToken!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/input token count.*exceeds the maximum/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
},
|
||||
120000,
|
||||
);
|
||||
|
||||
// Anthropic model via Antigravity
|
||||
it.skipIf(!antigravityToken)(
|
||||
"claude-sonnet-4-5 - should detect overflow via isContextOverflow",
|
||||
async () => {
|
||||
const model = getModel("google-antigravity", "claude-sonnet-4-5");
|
||||
const result = await testContextOverflow(model, antigravityToken!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
// Anthropic models return "prompt is too long" pattern
|
||||
expect(result.errorMessage).toMatch(/prompt is too long/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
},
|
||||
120000,
|
||||
);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// OpenAI Codex (OAuth)
|
||||
// Uses ChatGPT Plus/Pro subscription via OAuth
|
||||
// =============================================================================
|
||||
|
||||
describe("OpenAI Codex (OAuth)", () => {
|
||||
it.skipIf(!openaiCodexToken)(
|
||||
"gpt-5.2-codex - should detect overflow via isContextOverflow",
|
||||
async () => {
|
||||
const model = getModel("openai-codex", "gpt-5.2-codex");
|
||||
const result = await testContextOverflow(model, openaiCodexToken!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
},
|
||||
120000,
|
||||
);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Amazon Bedrock
|
||||
// Expected pattern: "Input is too long for requested model"
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!hasBedrockCredentials())("Amazon Bedrock", () => {
|
||||
it("claude-sonnet-4-5 - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("amazon-bedrock", "global.anthropic.claude-sonnet-4-5-20250929-v1:0");
|
||||
const result = await testContextOverflow(model, "");
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// xAI
|
||||
// Expected pattern: "maximum prompt length is X but the request contains Y"
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.XAI_API_KEY)("xAI", () => {
|
||||
it("grok-3-fast - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("xai", "grok-3-fast");
|
||||
const result = await testContextOverflow(model, process.env.XAI_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/maximum prompt length is \d+/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Groq
|
||||
// Expected pattern: "reduce the length of the messages"
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.GROQ_API_KEY)("Groq", () => {
|
||||
it("llama-3.3-70b-versatile - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("groq", "llama-3.3-70b-versatile");
|
||||
const result = await testContextOverflow(model, process.env.GROQ_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/reduce the length of the messages/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Cerebras
|
||||
// Expected: 400/413 status code with no body
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.CEREBRAS_API_KEY)("Cerebras", () => {
|
||||
it("qwen-3-235b - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("cerebras", "qwen-3-235b-a22b-instruct-2507");
|
||||
const result = await testContextOverflow(model, process.env.CEREBRAS_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
// Cerebras returns status code with no body (400, 413, or 429 for token rate limit)
|
||||
expect(result.errorMessage).toMatch(/4(00|13|29).*\(no body\)/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Hugging Face
|
||||
// Uses OpenAI-compatible Inference Router
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.HF_TOKEN)("Hugging Face", () => {
|
||||
it("Kimi-K2.5 - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("huggingface", "moonshotai/Kimi-K2.5");
|
||||
const result = await testContextOverflow(model, process.env.HF_TOKEN!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// z.ai
|
||||
// Special case: may return explicit overflow error text, may accept overflow silently,
|
||||
// or may rate limit instead
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.ZAI_API_KEY)("z.ai", () => {
|
||||
it("glm-4.5-flash - should detect overflow via isContextOverflow when z.ai reports it", async () => {
|
||||
const model = getModel("zai", "glm-4.5-flash");
|
||||
const result = await testContextOverflow(model, process.env.ZAI_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
// z.ai behavior is inconsistent:
|
||||
// - Sometimes returns explicit overflow error text via non-standard finish_reason handling
|
||||
// - Sometimes accepts overflow and returns successfully with usage.input > contextWindow
|
||||
// - Sometimes returns rate limit error
|
||||
if (result.stopReason === "error") {
|
||||
if (result.errorMessage?.match(/model_context_window_exceeded/i)) {
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
} else {
|
||||
console.log(" z.ai returned non-overflow error (possibly rate limited), skipping overflow detection");
|
||||
}
|
||||
} else if (result.stopReason === "stop") {
|
||||
if (result.hasUsageData && result.usage.input > model.contextWindow) {
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
} else {
|
||||
console.log(" z.ai returned stop without overflow usage data, skipping overflow detection");
|
||||
}
|
||||
}
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Mistral
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.MISTRAL_API_KEY)("Mistral", () => {
|
||||
it("devstral-medium-latest - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("mistral", "devstral-medium-latest");
|
||||
const result = await testContextOverflow(model, process.env.MISTRAL_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/too large for model with \d+ maximum context length/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// MiniMax
|
||||
// Expected pattern: TBD - need to test actual error message
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.MINIMAX_API_KEY)("MiniMax", () => {
|
||||
it("MiniMax-M2.7 - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("minimax", "MiniMax-M2.7");
|
||||
const result = await testContextOverflow(model, process.env.MINIMAX_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Kimi For Coding
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding", () => {
|
||||
it("kimi-k2-thinking - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("kimi-coding", "kimi-k2-thinking");
|
||||
const result = await testContextOverflow(model, process.env.KIMI_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Vercel AI Gateway - Unified API for multiple providers
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.AI_GATEWAY_API_KEY)("Vercel AI Gateway", () => {
|
||||
it("google/gemini-2.5-flash via AI Gateway - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("vercel-ai-gateway", "google/gemini-2.5-flash");
|
||||
const result = await testContextOverflow(model, process.env.AI_GATEWAY_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// OpenRouter - Multiple backend providers
|
||||
// Expected pattern: "maximum context length is X tokens"
|
||||
// =============================================================================
|
||||
|
||||
describe.skipIf(!process.env.OPENROUTER_API_KEY)("OpenRouter", () => {
|
||||
// Anthropic backend
|
||||
it("anthropic/claude-sonnet-4 via OpenRouter - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("openrouter", "anthropic/claude-sonnet-4");
|
||||
const result = await testContextOverflow(model, process.env.OPENROUTER_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/maximum context length is \d+ tokens/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
|
||||
// DeepSeek backend
|
||||
it("deepseek/deepseek-v3.2 via OpenRouter - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("openrouter", "deepseek/deepseek-v3.2");
|
||||
const result = await testContextOverflow(model, process.env.OPENROUTER_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/maximum context length is \d+ tokens/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
|
||||
// Mistral backend
|
||||
it("mistralai/mistral-large-2512 via OpenRouter - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("openrouter", "mistralai/mistral-large-2512");
|
||||
const result = await testContextOverflow(model, process.env.OPENROUTER_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/maximum context length is \d+ tokens/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
|
||||
// Google backend
|
||||
it("google/gemini-2.5-flash via OpenRouter - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("openrouter", "google/gemini-2.5-flash");
|
||||
const result = await testContextOverflow(model, process.env.OPENROUTER_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/maximum context length is \d+ tokens/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
|
||||
// Meta/Llama backend
|
||||
it("meta-llama/llama-4-maverick via OpenRouter - should detect overflow via isContextOverflow", async () => {
|
||||
const model = getModel("openrouter", "meta-llama/llama-4-maverick");
|
||||
const result = await testContextOverflow(model, process.env.OPENROUTER_API_KEY!);
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(result.errorMessage).toMatch(/maximum context length is \d+ tokens/i);
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Ollama (local)
|
||||
// =============================================================================
|
||||
|
||||
// Check if ollama is installed and local LLM tests are enabled
|
||||
let ollamaInstalled = false;
|
||||
if (!process.env.PI_NO_LOCAL_LLM) {
|
||||
try {
|
||||
execSync("which ollama", { stdio: "ignore" });
|
||||
ollamaInstalled = true;
|
||||
} catch {
|
||||
ollamaInstalled = false;
|
||||
}
|
||||
}
|
||||
|
||||
describe.skipIf(!ollamaInstalled)("Ollama (local)", () => {
|
||||
let ollamaProcess: ChildProcess | null = null;
|
||||
let model: Model<"openai-completions">;
|
||||
|
||||
beforeAll(async () => {
|
||||
// Check if model is available, if not pull it
|
||||
try {
|
||||
execSync("ollama list | grep -q 'gpt-oss:20b'", { stdio: "ignore" });
|
||||
} catch {
|
||||
console.log("Pulling gpt-oss:20b model for Ollama overflow tests...");
|
||||
try {
|
||||
execSync("ollama pull gpt-oss:20b", { stdio: "inherit" });
|
||||
} catch (_e) {
|
||||
console.warn("Failed to pull gpt-oss:20b model, tests will be skipped");
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Start ollama server
|
||||
ollamaProcess = spawn("ollama", ["serve"], {
|
||||
detached: false,
|
||||
stdio: "ignore",
|
||||
});
|
||||
|
||||
// Wait for server to be ready
|
||||
await new Promise<void>((resolve) => {
|
||||
const checkServer = async () => {
|
||||
try {
|
||||
const response = await fetch("http://localhost:11434/api/tags");
|
||||
if (response.ok) {
|
||||
resolve();
|
||||
} else {
|
||||
setTimeout(checkServer, 500);
|
||||
}
|
||||
} catch {
|
||||
setTimeout(checkServer, 500);
|
||||
}
|
||||
};
|
||||
setTimeout(checkServer, 1000);
|
||||
});
|
||||
|
||||
model = {
|
||||
id: "gpt-oss:20b",
|
||||
api: "openai-completions",
|
||||
provider: "ollama",
|
||||
baseUrl: "http://localhost:11434/v1",
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
contextWindow: 128000,
|
||||
maxTokens: 16000,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
name: "Ollama GPT-OSS 20B",
|
||||
};
|
||||
}, 60000);
|
||||
|
||||
afterAll(() => {
|
||||
if (ollamaProcess) {
|
||||
ollamaProcess.kill("SIGTERM");
|
||||
ollamaProcess = null;
|
||||
}
|
||||
});
|
||||
|
||||
it("gpt-oss:20b - should detect overflow via isContextOverflow (ollama silently truncates)", async () => {
|
||||
const result = await testContextOverflow(model, "ollama");
|
||||
logResult(result);
|
||||
|
||||
// Ollama silently truncates input instead of erroring
|
||||
// It returns stopReason "stop" with truncated usage
|
||||
// We cannot detect overflow via error message, only via usage comparison
|
||||
if (result.stopReason === "stop" && result.hasUsageData) {
|
||||
// Ollama truncated - check if reported usage is less than what we sent
|
||||
// This is a "silent overflow" - we can detect it if we know expected input size
|
||||
console.log(" Ollama silently truncated input to", result.usage.input, "tokens");
|
||||
// For now, we accept this behavior - Ollama doesn't give us a way to detect overflow
|
||||
} else if (result.stopReason === "error") {
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}
|
||||
}, 300000); // 5 min timeout for local model
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// LM Studio (local) - Skip if not running or local LLM tests disabled
|
||||
// =============================================================================
|
||||
|
||||
let lmStudioRunning = false;
|
||||
if (!process.env.PI_NO_LOCAL_LLM) {
|
||||
try {
|
||||
execSync("curl -s --max-time 1 http://localhost:1234/v1/models > /dev/null", { stdio: "ignore" });
|
||||
lmStudioRunning = true;
|
||||
} catch {
|
||||
lmStudioRunning = false;
|
||||
}
|
||||
}
|
||||
|
||||
describe.skipIf(!lmStudioRunning)("LM Studio (local)", () => {
|
||||
it("should detect overflow via isContextOverflow", async () => {
|
||||
const model: Model<"openai-completions"> = {
|
||||
id: "local-model",
|
||||
api: "openai-completions",
|
||||
provider: "lm-studio",
|
||||
baseUrl: "http://localhost:1234/v1",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
contextWindow: 8192,
|
||||
maxTokens: 2048,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
name: "LM Studio Local Model",
|
||||
};
|
||||
|
||||
const result = await testContextOverflow(model, "lm-studio");
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// llama.cpp server (local) - Skip if not running or not exposing /v1/completions
|
||||
// =============================================================================
|
||||
|
||||
let llamaCppRunning = false;
|
||||
if (!process.env.PI_NO_LOCAL_LLM) {
|
||||
try {
|
||||
execSync("curl -s --max-time 1 http://localhost:8081/health > /dev/null", { stdio: "ignore" });
|
||||
const probeStatus = execSync(
|
||||
'curl -s --max-time 1 -o /dev/null -w \'%{http_code}\' -X POST http://localhost:8081/v1/completions -H \'content-type: application/json\' -d \'{"model":"local-model","prompt":"ping","max_tokens":1}\'',
|
||||
{ encoding: "utf8" },
|
||||
).trim();
|
||||
llamaCppRunning = probeStatus !== "404" && probeStatus !== "405" && probeStatus !== "000";
|
||||
} catch {
|
||||
llamaCppRunning = false;
|
||||
}
|
||||
}
|
||||
|
||||
describe.skipIf(!llamaCppRunning)("llama.cpp (local)", () => {
|
||||
it("should detect overflow via isContextOverflow", async () => {
|
||||
// Using small context (4096) to match server --ctx-size setting
|
||||
const model: Model<"openai-completions"> = {
|
||||
id: "local-model",
|
||||
api: "openai-completions",
|
||||
provider: "llama.cpp",
|
||||
baseUrl: "http://localhost:8081/v1",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
contextWindow: 4096,
|
||||
maxTokens: 2048,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
name: "llama.cpp Local Model",
|
||||
};
|
||||
|
||||
const result = await testContextOverflow(model, "llama.cpp");
|
||||
logResult(result);
|
||||
|
||||
expect(result.stopReason).toBe("error");
|
||||
expect(isContextOverflow(result.response, model.contextWindow)).toBe(true);
|
||||
}, 120000);
|
||||
});
|
||||
});
|
||||
@@ -1,458 +0,0 @@
|
||||
/**
|
||||
* Cross-Provider Handoff Test
|
||||
*
|
||||
* Tests that contexts generated by one provider/model can be consumed by another.
|
||||
* This catches issues like:
|
||||
* - Tool call ID format incompatibilities (e.g., OpenAI Codex pipe characters)
|
||||
* - Thinking block transformation issues
|
||||
* - Message format incompatibilities
|
||||
*
|
||||
* Strategy:
|
||||
* 1. beforeAll: For each provider/model, generate a "small context" (if not cached):
|
||||
* - User message asking to use a tool
|
||||
* - Assistant response with thinking + tool call
|
||||
* - Tool result
|
||||
* - Final assistant response
|
||||
*
|
||||
* 2. Test: For each target provider/model:
|
||||
* - Concatenate ALL other contexts into one
|
||||
* - Ask the model to "say hi"
|
||||
* - If it fails, there's a compatibility issue
|
||||
*
|
||||
* Fixtures are generated fresh on each run.
|
||||
*/
|
||||
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { writeFileSync } from "fs";
|
||||
import { beforeAll, describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { completeSimple, getEnvApiKey } from "../src/stream.js";
|
||||
import type { Api, AssistantMessage, Message, Model, Tool, ToolResultMessage } from "../src/types.js";
|
||||
import { hasAzureOpenAICredentials } from "./azure-utils.js";
|
||||
import { resolveApiKey } from "./oauth.js";
|
||||
|
||||
// Simple tool for testing
|
||||
const testToolSchema = Type.Object({
|
||||
value: Type.Number({ description: "A number to double" }),
|
||||
});
|
||||
|
||||
const testTool: Tool<typeof testToolSchema> = {
|
||||
name: "double_number",
|
||||
description: "Doubles a number and returns the result",
|
||||
parameters: testToolSchema,
|
||||
};
|
||||
|
||||
// Provider/model pairs to test
|
||||
interface ProviderModelPair {
|
||||
provider: string;
|
||||
model: string;
|
||||
label: string;
|
||||
apiOverride?: Api;
|
||||
}
|
||||
|
||||
const PROVIDER_MODEL_PAIRS: ProviderModelPair[] = [
|
||||
// Anthropic
|
||||
{ provider: "anthropic", model: "claude-sonnet-4-5", label: "anthropic-claude-sonnet-4-5" },
|
||||
// Google
|
||||
{ provider: "google", model: "gemini-3-flash-preview", label: "google-gemini-3-flash-preview" },
|
||||
// OpenAI
|
||||
{
|
||||
provider: "openai",
|
||||
model: "gpt-4o-mini",
|
||||
label: "openai-completions-gpt-4o-mini",
|
||||
apiOverride: "openai-completions",
|
||||
},
|
||||
{ provider: "openai", model: "gpt-5-mini", label: "openai-responses-gpt-5-mini" },
|
||||
{ provider: "azure-openai-responses", model: "gpt-4o-mini", label: "azure-openai-responses-gpt-4o-mini" },
|
||||
// OpenAI Codex
|
||||
{ provider: "openai-codex", model: "gpt-5.2-codex", label: "openai-codex-gpt-5.2-codex" },
|
||||
// Google Antigravity
|
||||
{ provider: "google-antigravity", model: "gemini-3-flash", label: "antigravity-gemini-3-flash" },
|
||||
{ provider: "google-antigravity", model: "claude-sonnet-4-5", label: "antigravity-claude-sonnet-4-5" },
|
||||
// GitHub Copilot
|
||||
{ provider: "github-copilot", model: "claude-sonnet-4.5", label: "copilot-claude-sonnet-4.5" },
|
||||
{ provider: "github-copilot", model: "gpt-5.1-codex", label: "copilot-gpt-5.1-codex" },
|
||||
{ provider: "github-copilot", model: "gemini-3-flash-preview", label: "copilot-gemini-3-flash-preview" },
|
||||
{ provider: "github-copilot", model: "grok-code-fast-1", label: "copilot-grok-code-fast-1" },
|
||||
// Amazon Bedrock
|
||||
{
|
||||
provider: "amazon-bedrock",
|
||||
model: "global.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
||||
label: "bedrock-claude-sonnet-4-5",
|
||||
},
|
||||
// xAI
|
||||
{ provider: "xai", model: "grok-code-fast-1", label: "xai-grok-code-fast-1" },
|
||||
// Cerebras
|
||||
{ provider: "cerebras", model: "zai-glm-4.7", label: "cerebras-zai-glm-4.7" },
|
||||
// Groq
|
||||
{ provider: "groq", model: "openai/gpt-oss-120b", label: "groq-gpt-oss-120b" },
|
||||
// Hugging Face
|
||||
{ provider: "huggingface", model: "moonshotai/Kimi-K2.5", label: "huggingface-kimi-k2.5" },
|
||||
// Kimi For Coding
|
||||
{ provider: "kimi-coding", model: "kimi-k2-thinking", label: "kimi-coding-k2-thinking" },
|
||||
// Mistral
|
||||
{ provider: "mistral", model: "devstral-medium-latest", label: "mistral-devstral-medium" },
|
||||
// MiniMax
|
||||
{ provider: "minimax", model: "MiniMax-M2.7", label: "minimax-m2.7" },
|
||||
// OpenCode Zen
|
||||
{ provider: "opencode", model: "big-pickle", label: "zen-big-pickle" },
|
||||
{ provider: "opencode", model: "claude-sonnet-4-5", label: "zen-claude-sonnet-4-5" },
|
||||
{ provider: "opencode", model: "gemini-3-flash", label: "zen-gemini-3-flash" },
|
||||
{ provider: "opencode", model: "glm-4.7-free", label: "zen-glm-4.7-free" },
|
||||
{ provider: "opencode", model: "gpt-5.2-codex", label: "zen-gpt-5.2-codex" },
|
||||
{ provider: "opencode", model: "minimax-m2.1-free", label: "zen-minimax-m2.1-free" },
|
||||
// OpenCode Go
|
||||
{ provider: "opencode-go", model: "kimi-k2.5", label: "go-kimi-k2.5" },
|
||||
{ provider: "opencode-go", model: "minimax-m2.5", label: "go-minimax-m2.5" },
|
||||
];
|
||||
|
||||
// Cached context structure
|
||||
interface CachedContext {
|
||||
label: string;
|
||||
provider: string;
|
||||
model: string;
|
||||
api: Api;
|
||||
messages: Message[];
|
||||
generatedAt: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get API key for provider - checks OAuth storage first, then env vars
|
||||
*/
|
||||
async function getApiKey(provider: string): Promise<string | undefined> {
|
||||
const oauthKey = await resolveApiKey(provider);
|
||||
if (oauthKey) return oauthKey;
|
||||
return getEnvApiKey(provider);
|
||||
}
|
||||
|
||||
/**
|
||||
* Synchronous check for API key availability (env vars only, for skipIf)
|
||||
*/
|
||||
function hasApiKey(provider: string): boolean {
|
||||
if (provider === "azure-openai-responses") {
|
||||
return hasAzureOpenAICredentials();
|
||||
}
|
||||
return !!getEnvApiKey(provider);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if any provider has API keys available (for skipIf at describe level)
|
||||
*/
|
||||
function hasAnyApiKey(): boolean {
|
||||
return PROVIDER_MODEL_PAIRS.some((pair) => hasApiKey(pair.provider));
|
||||
}
|
||||
|
||||
function dumpFailurePayload(params: { label: string; error: string; payload?: unknown; messages: Message[] }): void {
|
||||
const filename = `/tmp/pi-handoff-${params.label}-${Date.now()}.json`;
|
||||
const body = {
|
||||
label: params.label,
|
||||
error: params.error,
|
||||
payload: params.payload,
|
||||
messages: params.messages,
|
||||
};
|
||||
writeFileSync(filename, JSON.stringify(body, null, 2));
|
||||
console.log(`Wrote failure payload to ${filename}`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a context from a provider/model pair.
|
||||
* Makes a real API call to get authentic tool call IDs and thinking blocks.
|
||||
*/
|
||||
async function generateContext(
|
||||
pair: ProviderModelPair,
|
||||
apiKey: string,
|
||||
): Promise<{ messages: Message[]; api: Api } | null> {
|
||||
const baseModel = (getModel as (p: string, m: string) => Model<Api> | undefined)(pair.provider, pair.model);
|
||||
if (!baseModel) {
|
||||
console.log(` Model not found: ${pair.provider}/${pair.model}`);
|
||||
return null;
|
||||
}
|
||||
|
||||
const model: Model<Api> = pair.apiOverride ? { ...baseModel, api: pair.apiOverride } : baseModel;
|
||||
|
||||
const userMessage: Message = {
|
||||
role: "user",
|
||||
content: "Please double the number 21 using the double_number tool.",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const supportsReasoning = model.reasoning === true;
|
||||
let lastPayload: unknown;
|
||||
let assistantResponse: AssistantMessage;
|
||||
try {
|
||||
assistantResponse = await completeSimple(
|
||||
model,
|
||||
{
|
||||
systemPrompt: "You are a helpful assistant. Use the provided tool to complete the task.",
|
||||
messages: [userMessage],
|
||||
tools: [testTool],
|
||||
},
|
||||
{
|
||||
apiKey,
|
||||
reasoning: supportsReasoning ? "high" : undefined,
|
||||
onPayload: (payload) => {
|
||||
lastPayload = payload;
|
||||
},
|
||||
},
|
||||
);
|
||||
} catch (error) {
|
||||
const msg = error instanceof Error ? error.message : String(error);
|
||||
console.log(` Initial request failed: ${msg}`);
|
||||
dumpFailurePayload({
|
||||
label: `${pair.label}-initial`,
|
||||
error: msg,
|
||||
payload: lastPayload,
|
||||
messages: [userMessage],
|
||||
});
|
||||
return null;
|
||||
}
|
||||
|
||||
if (assistantResponse.stopReason === "error") {
|
||||
console.log(` Initial request error: ${assistantResponse.errorMessage}`);
|
||||
dumpFailurePayload({
|
||||
label: `${pair.label}-initial`,
|
||||
error: assistantResponse.errorMessage || "Unknown error",
|
||||
payload: lastPayload,
|
||||
messages: [userMessage],
|
||||
});
|
||||
return null;
|
||||
}
|
||||
|
||||
const toolCall = assistantResponse.content.find((c) => c.type === "toolCall");
|
||||
if (!toolCall || toolCall.type !== "toolCall") {
|
||||
console.log(` No tool call in response (stopReason: ${assistantResponse.stopReason})`);
|
||||
return {
|
||||
messages: [userMessage, assistantResponse],
|
||||
api: model.api,
|
||||
};
|
||||
}
|
||||
|
||||
console.log(` Tool call ID: ${toolCall.id}`);
|
||||
|
||||
const toolResult: ToolResultMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
content: [{ type: "text", text: "42" }],
|
||||
isError: false,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
let finalResponse: AssistantMessage;
|
||||
const messagesForFinal = [userMessage, assistantResponse, toolResult];
|
||||
try {
|
||||
finalResponse = await completeSimple(
|
||||
model,
|
||||
{
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
messages: messagesForFinal,
|
||||
tools: [testTool],
|
||||
},
|
||||
{
|
||||
apiKey,
|
||||
reasoning: supportsReasoning ? "high" : undefined,
|
||||
onPayload: (payload) => {
|
||||
lastPayload = payload;
|
||||
},
|
||||
},
|
||||
);
|
||||
} catch (error) {
|
||||
const msg = error instanceof Error ? error.message : String(error);
|
||||
console.log(` Final request failed: ${msg}`);
|
||||
dumpFailurePayload({
|
||||
label: `${pair.label}-final`,
|
||||
error: msg,
|
||||
payload: lastPayload,
|
||||
messages: messagesForFinal,
|
||||
});
|
||||
return null;
|
||||
}
|
||||
|
||||
if (finalResponse.stopReason === "error") {
|
||||
console.log(` Final request error: ${finalResponse.errorMessage}`);
|
||||
dumpFailurePayload({
|
||||
label: `${pair.label}-final`,
|
||||
error: finalResponse.errorMessage || "Unknown error",
|
||||
payload: lastPayload,
|
||||
messages: messagesForFinal,
|
||||
});
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
messages: [userMessage, assistantResponse, toolResult, finalResponse],
|
||||
api: model.api,
|
||||
};
|
||||
}
|
||||
|
||||
describe.skipIf(!hasAnyApiKey())("Cross-Provider Handoff", () => {
|
||||
let contexts: Record<string, CachedContext>;
|
||||
let availablePairs: ProviderModelPair[];
|
||||
|
||||
beforeAll(async () => {
|
||||
contexts = {};
|
||||
availablePairs = [];
|
||||
|
||||
console.log("\n=== Generating Fixtures ===\n");
|
||||
|
||||
for (const pair of PROVIDER_MODEL_PAIRS) {
|
||||
const apiKey = await getApiKey(pair.provider);
|
||||
if (!apiKey) {
|
||||
console.log(`[${pair.label}] Skipping - no auth for ${pair.provider}`);
|
||||
continue;
|
||||
}
|
||||
|
||||
console.log(`[${pair.label}] Generating fixture...`);
|
||||
const result = await generateContext(pair, apiKey);
|
||||
|
||||
if (!result || result.messages.length < 4) {
|
||||
console.log(`[${pair.label}] Failed to generate fixture, skipping`);
|
||||
continue;
|
||||
}
|
||||
|
||||
contexts[pair.label] = {
|
||||
label: pair.label,
|
||||
provider: pair.provider,
|
||||
model: pair.model,
|
||||
api: result.api,
|
||||
messages: result.messages,
|
||||
generatedAt: new Date().toISOString(),
|
||||
};
|
||||
availablePairs.push(pair);
|
||||
console.log(`[${pair.label}] Generated ${result.messages.length} messages`);
|
||||
}
|
||||
|
||||
console.log(`\n=== ${availablePairs.length}/${PROVIDER_MODEL_PAIRS.length} contexts available ===\n`);
|
||||
}, 300000);
|
||||
|
||||
it.skipIf(!hasAnyApiKey())("should have at least 2 fixtures to test handoffs", () => {
|
||||
expect(Object.keys(contexts).length).toBeGreaterThanOrEqual(2);
|
||||
});
|
||||
|
||||
it.skipIf(!hasAnyApiKey())(
|
||||
"should handle cross-provider handoffs for each target",
|
||||
async () => {
|
||||
const contextLabels = Object.keys(contexts);
|
||||
|
||||
if (contextLabels.length < 2) {
|
||||
console.log("Not enough fixtures for handoff test, skipping");
|
||||
return;
|
||||
}
|
||||
|
||||
console.log("\n=== Testing Cross-Provider Handoffs ===\n");
|
||||
|
||||
const results: { target: string; success: boolean; error?: string }[] = [];
|
||||
|
||||
for (const targetPair of availablePairs) {
|
||||
const apiKey = await getApiKey(targetPair.provider);
|
||||
if (!apiKey) {
|
||||
console.log(`[Target: ${targetPair.label}] Skipping - no auth`);
|
||||
continue;
|
||||
}
|
||||
|
||||
// Collect messages from ALL OTHER contexts
|
||||
const otherMessages: Message[] = [];
|
||||
for (const [label, ctx] of Object.entries(contexts)) {
|
||||
if (label === targetPair.label) continue;
|
||||
otherMessages.push(...ctx.messages);
|
||||
}
|
||||
|
||||
if (otherMessages.length === 0) {
|
||||
console.log(`[Target: ${targetPair.label}] Skipping - no other contexts`);
|
||||
continue;
|
||||
}
|
||||
|
||||
const allMessages: Message[] = [
|
||||
...otherMessages,
|
||||
{
|
||||
role: "user",
|
||||
content:
|
||||
"Great, thanks for all that help! Now just say 'Hello, handoff successful!' to confirm you received everything.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
];
|
||||
|
||||
const baseModel = (getModel as (p: string, m: string) => Model<Api> | undefined)(
|
||||
targetPair.provider,
|
||||
targetPair.model,
|
||||
);
|
||||
if (!baseModel) {
|
||||
console.log(`[Target: ${targetPair.label}] Model not found`);
|
||||
continue;
|
||||
}
|
||||
|
||||
const model: Model<Api> = targetPair.apiOverride
|
||||
? { ...baseModel, api: targetPair.apiOverride }
|
||||
: baseModel;
|
||||
const supportsReasoning = model.reasoning === true;
|
||||
|
||||
console.log(
|
||||
`[Target: ${targetPair.label}] Testing with ${otherMessages.length} messages from other providers...`,
|
||||
);
|
||||
|
||||
let lastPayload: unknown;
|
||||
try {
|
||||
const response = await completeSimple(
|
||||
model,
|
||||
{
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
messages: allMessages,
|
||||
tools: [testTool],
|
||||
},
|
||||
{
|
||||
apiKey,
|
||||
reasoning: supportsReasoning ? "high" : undefined,
|
||||
onPayload: (payload) => {
|
||||
lastPayload = payload;
|
||||
},
|
||||
},
|
||||
);
|
||||
|
||||
if (response.stopReason === "error") {
|
||||
console.log(`[Target: ${targetPair.label}] FAILED: ${response.errorMessage}`);
|
||||
dumpFailurePayload({
|
||||
label: targetPair.label,
|
||||
error: response.errorMessage || "Unknown error",
|
||||
payload: lastPayload,
|
||||
messages: allMessages,
|
||||
});
|
||||
results.push({ target: targetPair.label, success: false, error: response.errorMessage });
|
||||
} else {
|
||||
const text = response.content
|
||||
.filter((c) => c.type === "text")
|
||||
.map((c) => c.text)
|
||||
.join(" ");
|
||||
const preview = text.slice(0, 100).replace(/\n/g, " ");
|
||||
console.log(`[Target: ${targetPair.label}] SUCCESS: ${preview}...`);
|
||||
results.push({ target: targetPair.label, success: true });
|
||||
}
|
||||
} catch (error) {
|
||||
const msg = error instanceof Error ? error.message : String(error);
|
||||
console.log(`[Target: ${targetPair.label}] EXCEPTION: ${msg}`);
|
||||
dumpFailurePayload({
|
||||
label: targetPair.label,
|
||||
error: msg,
|
||||
payload: lastPayload,
|
||||
messages: allMessages,
|
||||
});
|
||||
results.push({ target: targetPair.label, success: false, error: msg });
|
||||
}
|
||||
}
|
||||
|
||||
console.log("\n=== Results Summary ===\n");
|
||||
const successes = results.filter((r) => r.success);
|
||||
const failures = results.filter((r) => !r.success);
|
||||
|
||||
console.log(`Passed: ${successes.length}/${results.length}`);
|
||||
if (failures.length > 0) {
|
||||
console.log("\nFailures:");
|
||||
for (const f of failures) {
|
||||
console.log(` - ${f.target}: ${f.error}`);
|
||||
}
|
||||
}
|
||||
|
||||
expect(failures.length).toBe(0);
|
||||
},
|
||||
600000,
|
||||
);
|
||||
});
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 2.5 KiB |
@@ -1,742 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { complete } from "../src/stream.js";
|
||||
import type { Api, AssistantMessage, Context, Model, StreamOptions, UserMessage } from "../src/types.js";
|
||||
|
||||
type StreamOptionsWithExtras = StreamOptions & Record<string, unknown>;
|
||||
|
||||
import { hasAzureOpenAICredentials, resolveAzureDeploymentName } from "./azure-utils.js";
|
||||
import { hasBedrockCredentials } from "./bedrock-utils.js";
|
||||
import { resolveApiKey } from "./oauth.js";
|
||||
|
||||
// Resolve OAuth tokens at module level (async, runs before tests)
|
||||
const oauthTokens = await Promise.all([
|
||||
resolveApiKey("anthropic"),
|
||||
resolveApiKey("github-copilot"),
|
||||
resolveApiKey("google-gemini-cli"),
|
||||
resolveApiKey("google-antigravity"),
|
||||
resolveApiKey("openai-codex"),
|
||||
]);
|
||||
const [anthropicOAuthToken, githubCopilotToken, geminiCliToken, antigravityToken, openaiCodexToken] = oauthTokens;
|
||||
|
||||
async function testEmptyMessage<TApi extends Api>(llm: Model<TApi>, options: StreamOptionsWithExtras = {}) {
|
||||
// Test with completely empty content array
|
||||
const emptyMessage: UserMessage = {
|
||||
role: "user",
|
||||
content: [],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
messages: [emptyMessage],
|
||||
};
|
||||
|
||||
const response = await complete(llm, context, options);
|
||||
|
||||
// Should either handle gracefully or return an error
|
||||
expect(response).toBeDefined();
|
||||
expect(response.role).toBe("assistant");
|
||||
// Should handle empty string gracefully
|
||||
if (response.stopReason === "error") {
|
||||
expect(response.errorMessage).toBeDefined();
|
||||
} else {
|
||||
expect(response.content).toBeDefined();
|
||||
}
|
||||
}
|
||||
|
||||
async function testEmptyStringMessage<TApi extends Api>(llm: Model<TApi>, options: StreamOptionsWithExtras = {}) {
|
||||
// Test with empty string content
|
||||
const context: Context = {
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const response = await complete(llm, context, options);
|
||||
|
||||
expect(response).toBeDefined();
|
||||
expect(response.role).toBe("assistant");
|
||||
|
||||
// Should handle empty string gracefully
|
||||
if (response.stopReason === "error") {
|
||||
expect(response.errorMessage).toBeDefined();
|
||||
} else {
|
||||
expect(response.content).toBeDefined();
|
||||
}
|
||||
}
|
||||
|
||||
async function testWhitespaceOnlyMessage<TApi extends Api>(llm: Model<TApi>, options: StreamOptionsWithExtras = {}) {
|
||||
// Test with whitespace-only content
|
||||
const context: Context = {
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: " \n\t ",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const response = await complete(llm, context, options);
|
||||
|
||||
expect(response).toBeDefined();
|
||||
expect(response.role).toBe("assistant");
|
||||
|
||||
// Should handle whitespace-only gracefully
|
||||
if (response.stopReason === "error") {
|
||||
expect(response.errorMessage).toBeDefined();
|
||||
} else {
|
||||
expect(response.content).toBeDefined();
|
||||
}
|
||||
}
|
||||
|
||||
async function testEmptyAssistantMessage<TApi extends Api>(llm: Model<TApi>, options: StreamOptionsWithExtras = {}) {
|
||||
// Test with empty assistant message in conversation flow
|
||||
// User -> Empty Assistant -> User
|
||||
const emptyAssistant: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [],
|
||||
api: llm.api,
|
||||
provider: llm.provider,
|
||||
model: llm.id,
|
||||
usage: {
|
||||
input: 10,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 10,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "Hello, how are you?",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
emptyAssistant,
|
||||
{
|
||||
role: "user",
|
||||
content: "Please respond this time.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const response = await complete(llm, context, options);
|
||||
|
||||
expect(response).toBeDefined();
|
||||
expect(response.role).toBe("assistant");
|
||||
|
||||
// Should handle empty assistant message in context gracefully
|
||||
if (response.stopReason === "error") {
|
||||
expect(response.errorMessage).toBeDefined();
|
||||
} else {
|
||||
expect(response.content).toBeDefined();
|
||||
expect(response.content.length).toBeGreaterThan(0);
|
||||
}
|
||||
}
|
||||
|
||||
describe("AI Providers Empty Message Tests", () => {
|
||||
describe.skipIf(!process.env.GEMINI_API_KEY)("Google Provider Empty Messages", () => {
|
||||
const llm = getModel("google", "gemini-2.5-flash");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Completions Provider Empty Messages", () => {
|
||||
const llm = getModel("openai", "gpt-4o-mini");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Responses Provider Empty Messages", () => {
|
||||
const llm = getModel("openai", "gpt-5-mini");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!hasAzureOpenAICredentials())("Azure OpenAI Responses Provider Empty Messages", () => {
|
||||
const llm = getModel("azure-openai-responses", "gpt-4o-mini");
|
||||
const azureDeploymentName = resolveAzureDeploymentName(llm.id);
|
||||
const azureOptions = azureDeploymentName ? { azureDeploymentName } : {};
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm, azureOptions);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm, azureOptions);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm, azureOptions);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm, azureOptions);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic Provider Empty Messages", () => {
|
||||
const llm = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.XAI_API_KEY)("xAI Provider Empty Messages", () => {
|
||||
const llm = getModel("xai", "grok-3");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.GROQ_API_KEY)("Groq Provider Empty Messages", () => {
|
||||
const llm = getModel("groq", "openai/gpt-oss-20b");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.CEREBRAS_API_KEY)("Cerebras Provider Empty Messages", () => {
|
||||
const llm = getModel("cerebras", "gpt-oss-120b");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.HF_TOKEN)("Hugging Face Provider Empty Messages", () => {
|
||||
const llm = getModel("huggingface", "moonshotai/Kimi-K2.5");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.ZAI_API_KEY)("zAI Provider Empty Messages", () => {
|
||||
const llm = getModel("zai", "glm-4.5-air");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.MISTRAL_API_KEY)("Mistral Provider Empty Messages", () => {
|
||||
const llm = getModel("mistral", "devstral-medium-latest");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.MINIMAX_API_KEY)("MiniMax Provider Empty Messages", () => {
|
||||
const llm = getModel("minimax", "MiniMax-M2.7");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider Empty Messages", () => {
|
||||
const llm = getModel("kimi-coding", "kimi-k2-thinking");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.AI_GATEWAY_API_KEY)("Vercel AI Gateway Provider Empty Messages", () => {
|
||||
const llm = getModel("vercel-ai-gateway", "google/gemini-2.5-flash");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!hasBedrockCredentials())("Amazon Bedrock Provider Empty Messages", () => {
|
||||
const llm = getModel("amazon-bedrock", "global.anthropic.claude-sonnet-4-5-20250929-v1:0");
|
||||
|
||||
it("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle whitespace-only content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testWhitespaceOnlyMessage(llm);
|
||||
});
|
||||
|
||||
it("should handle empty assistant message in conversation", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyAssistantMessage(llm);
|
||||
});
|
||||
});
|
||||
|
||||
// =========================================================================
|
||||
// OAuth-based providers (credentials from ~/.pi/agent/oauth.json)
|
||||
// =========================================================================
|
||||
|
||||
describe("Anthropic OAuth Provider Empty Messages", () => {
|
||||
const llm = getModel("anthropic", "claude-3-5-haiku-20241022");
|
||||
|
||||
it.skipIf(!anthropicOAuthToken)("should handle empty content array", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyMessage(llm, { apiKey: anthropicOAuthToken });
|
||||
});
|
||||
|
||||
it.skipIf(!anthropicOAuthToken)("should handle empty string content", { retry: 3, timeout: 30000 }, async () => {
|
||||
await testEmptyStringMessage(llm, { apiKey: anthropicOAuthToken });
|
||||
});
|
||||
|
||||
it.skipIf(!anthropicOAuthToken)(
|
||||
"should handle whitespace-only content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
await testWhitespaceOnlyMessage(llm, { apiKey: anthropicOAuthToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!anthropicOAuthToken)(
|
||||
"should handle empty assistant message in conversation",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
await testEmptyAssistantMessage(llm, { apiKey: anthropicOAuthToken });
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
describe("GitHub Copilot Provider Empty Messages", () => {
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"gpt-4o - should handle empty content array",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "gpt-4o");
|
||||
await testEmptyMessage(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"gpt-4o - should handle empty string content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "gpt-4o");
|
||||
await testEmptyStringMessage(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"gpt-4o - should handle whitespace-only content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "gpt-4o");
|
||||
await testWhitespaceOnlyMessage(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"gpt-4o - should handle empty assistant message in conversation",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "gpt-4o");
|
||||
await testEmptyAssistantMessage(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"claude-sonnet-4 - should handle empty content array",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "claude-sonnet-4");
|
||||
await testEmptyMessage(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"claude-sonnet-4 - should handle empty string content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "claude-sonnet-4");
|
||||
await testEmptyStringMessage(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"claude-sonnet-4 - should handle whitespace-only content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "claude-sonnet-4");
|
||||
await testWhitespaceOnlyMessage(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"claude-sonnet-4 - should handle empty assistant message in conversation",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "claude-sonnet-4");
|
||||
await testEmptyAssistantMessage(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
describe("Google Gemini CLI Provider Empty Messages", () => {
|
||||
it.skipIf(!geminiCliToken)(
|
||||
"gemini-2.5-flash - should handle empty content array",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-gemini-cli", "gemini-2.5-flash");
|
||||
await testEmptyMessage(llm, { apiKey: geminiCliToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!geminiCliToken)(
|
||||
"gemini-2.5-flash - should handle empty string content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-gemini-cli", "gemini-2.5-flash");
|
||||
await testEmptyStringMessage(llm, { apiKey: geminiCliToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!geminiCliToken)(
|
||||
"gemini-2.5-flash - should handle whitespace-only content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-gemini-cli", "gemini-2.5-flash");
|
||||
await testWhitespaceOnlyMessage(llm, { apiKey: geminiCliToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!geminiCliToken)(
|
||||
"gemini-2.5-flash - should handle empty assistant message in conversation",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-gemini-cli", "gemini-2.5-flash");
|
||||
await testEmptyAssistantMessage(llm, { apiKey: geminiCliToken });
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
describe("Google Antigravity Provider Empty Messages", () => {
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gemini-3-flash - should handle empty content array",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gemini-3-flash");
|
||||
await testEmptyMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gemini-3-flash - should handle empty string content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gemini-3-flash");
|
||||
await testEmptyStringMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gemini-3-flash - should handle whitespace-only content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gemini-3-flash");
|
||||
await testWhitespaceOnlyMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gemini-3-flash - should handle empty assistant message in conversation",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gemini-3-flash");
|
||||
await testEmptyAssistantMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"claude-sonnet-4-5 - should handle empty content array",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "claude-sonnet-4-5");
|
||||
await testEmptyMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"claude-sonnet-4-5 - should handle empty string content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "claude-sonnet-4-5");
|
||||
await testEmptyStringMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"claude-sonnet-4-5 - should handle whitespace-only content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "claude-sonnet-4-5");
|
||||
await testWhitespaceOnlyMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"claude-sonnet-4-5 - should handle empty assistant message in conversation",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "claude-sonnet-4-5");
|
||||
await testEmptyAssistantMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gpt-oss-120b-medium - should handle empty content array",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gpt-oss-120b-medium");
|
||||
await testEmptyMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gpt-oss-120b-medium - should handle empty string content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gpt-oss-120b-medium");
|
||||
await testEmptyStringMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gpt-oss-120b-medium - should handle whitespace-only content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gpt-oss-120b-medium");
|
||||
await testWhitespaceOnlyMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gpt-oss-120b-medium - should handle empty assistant message in conversation",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gpt-oss-120b-medium");
|
||||
await testEmptyAssistantMessage(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
describe("OpenAI Codex Provider Empty Messages", () => {
|
||||
it.skipIf(!openaiCodexToken)(
|
||||
"gpt-5.2-codex - should handle empty content array",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("openai-codex", "gpt-5.2-codex");
|
||||
await testEmptyMessage(llm, { apiKey: openaiCodexToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!openaiCodexToken)(
|
||||
"gpt-5.2-codex - should handle empty string content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("openai-codex", "gpt-5.2-codex");
|
||||
await testEmptyStringMessage(llm, { apiKey: openaiCodexToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!openaiCodexToken)(
|
||||
"gpt-5.2-codex - should handle whitespace-only content",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("openai-codex", "gpt-5.2-codex");
|
||||
await testWhitespaceOnlyMessage(llm, { apiKey: openaiCodexToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!openaiCodexToken)(
|
||||
"gpt-5.2-codex - should handle empty assistant message in conversation",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("openai-codex", "gpt-5.2-codex");
|
||||
await testEmptyAssistantMessage(llm, { apiKey: openaiCodexToken });
|
||||
},
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -1,597 +0,0 @@
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import {
|
||||
complete,
|
||||
fauxAssistantMessage,
|
||||
fauxText,
|
||||
fauxThinking,
|
||||
fauxToolCall,
|
||||
registerFauxProvider,
|
||||
stream,
|
||||
Type,
|
||||
} from "../src/index.js";
|
||||
import type { AssistantMessageEvent, Context } from "../src/types.js";
|
||||
|
||||
async function collectEvents(streamResult: ReturnType<typeof stream>): Promise<AssistantMessageEvent[]> {
|
||||
const events: AssistantMessageEvent[] = [];
|
||||
for await (const event of streamResult) {
|
||||
events.push(event);
|
||||
}
|
||||
return events;
|
||||
}
|
||||
|
||||
const registrations: Array<{ unregister: () => void }> = [];
|
||||
|
||||
afterEach(() => {
|
||||
for (const registration of registrations.splice(0)) {
|
||||
registration.unregister();
|
||||
}
|
||||
});
|
||||
|
||||
describe("faux provider", () => {
|
||||
it("registers a custom provider and estimates usage", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([fauxAssistantMessage("hello world")]);
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "Be concise.",
|
||||
messages: [{ role: "user", content: "hi there", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
const response = await complete(registration.getModel(), context);
|
||||
expect(response.content).toEqual([{ type: "text", text: "hello world" }]);
|
||||
expect(response.usage.input).toBeGreaterThan(0);
|
||||
expect(response.usage.output).toBeGreaterThan(0);
|
||||
expect(response.usage.totalTokens).toBe(response.usage.input + response.usage.output);
|
||||
expect(registration.state.callCount).toBe(1);
|
||||
});
|
||||
|
||||
it("supports helper blocks for text, thinking, and tool calls", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
fauxAssistantMessage([fauxThinking("think"), fauxToolCall("echo", { text: "hi" }), fauxText("done")], {
|
||||
stopReason: "toolUse",
|
||||
}),
|
||||
]);
|
||||
|
||||
const response = await complete(registration.getModel(), {
|
||||
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
|
||||
});
|
||||
|
||||
expect(response.content).toEqual([
|
||||
{ type: "thinking", thinking: "think" },
|
||||
{ type: "toolCall", id: expect.any(String), name: "echo", arguments: { text: "hi" } },
|
||||
{ type: "text", text: "done" },
|
||||
]);
|
||||
expect(response.stopReason).toBe("toolUse");
|
||||
});
|
||||
|
||||
it("supports multiple models with per-model reasoning and model-aware factories", async () => {
|
||||
const registration = registerFauxProvider({
|
||||
models: [
|
||||
{ id: "faux-fast", name: "Faux Fast", reasoning: false },
|
||||
{ id: "faux-thinker", name: "Faux Thinker", reasoning: true },
|
||||
],
|
||||
});
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
(_context, _options, _state, model) => fauxAssistantMessage(`${model.id}:${String(model.reasoning)}`),
|
||||
(_context, _options, _state, model) => fauxAssistantMessage(`${model.id}:${String(model.reasoning)}`),
|
||||
]);
|
||||
|
||||
expect(registration.models.map((model) => model.id)).toEqual(["faux-fast", "faux-thinker"]);
|
||||
expect(registration.getModel()).toBe(registration.models[0]);
|
||||
expect(registration.getModel("faux-fast")?.reasoning).toBe(false);
|
||||
expect(registration.getModel("faux-thinker")?.reasoning).toBe(true);
|
||||
|
||||
const fast = await complete(registration.getModel("faux-fast")!, {
|
||||
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
|
||||
});
|
||||
const thinker = await complete(registration.getModel("faux-thinker")!, {
|
||||
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
|
||||
});
|
||||
|
||||
expect(fast.content).toEqual([{ type: "text", text: "faux-fast:false" }]);
|
||||
expect(thinker.content).toEqual([{ type: "text", text: "faux-thinker:true" }]);
|
||||
});
|
||||
|
||||
it("rewrites api, provider, and model on returned messages", async () => {
|
||||
const registration = registerFauxProvider({
|
||||
api: "faux:test",
|
||||
provider: "faux-provider",
|
||||
models: [{ id: "faux-model" }],
|
||||
});
|
||||
registrations.push(registration);
|
||||
registration.setResponses([fauxAssistantMessage("hello")]);
|
||||
|
||||
const response = await complete(registration.getModel(), {
|
||||
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
|
||||
});
|
||||
|
||||
expect(response.api).toBe("faux:test");
|
||||
expect(response.provider).toBe("faux-provider");
|
||||
expect(response.model).toBe("faux-model");
|
||||
});
|
||||
|
||||
it("consumes queued responses in order and errors when exhausted", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
|
||||
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
const first = await complete(registration.getModel(), context);
|
||||
const second = await complete(registration.getModel(), context);
|
||||
const exhausted = await complete(registration.getModel(), context);
|
||||
|
||||
expect(first.content).toEqual([{ type: "text", text: "first" }]);
|
||||
expect(second.content).toEqual([{ type: "text", text: "second" }]);
|
||||
expect(exhausted.stopReason).toBe("error");
|
||||
expect(exhausted.errorMessage).toBe("No more faux responses queued");
|
||||
expect(registration.getPendingResponseCount()).toBe(0);
|
||||
expect(registration.state.callCount).toBe(3);
|
||||
});
|
||||
|
||||
it("can replace and append queued responses", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([fauxAssistantMessage("first")]);
|
||||
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
expect((await complete(registration.getModel(), context)).content).toEqual([{ type: "text", text: "first" }]);
|
||||
expect(registration.getPendingResponseCount()).toBe(0);
|
||||
|
||||
registration.setResponses([fauxAssistantMessage("second")]);
|
||||
expect(registration.getPendingResponseCount()).toBe(1);
|
||||
expect((await complete(registration.getModel(), context)).content).toEqual([{ type: "text", text: "second" }]);
|
||||
|
||||
registration.appendResponses([fauxAssistantMessage("third"), fauxAssistantMessage("fourth")]);
|
||||
expect(registration.getPendingResponseCount()).toBe(2);
|
||||
expect((await complete(registration.getModel(), context)).content).toEqual([{ type: "text", text: "third" }]);
|
||||
expect((await complete(registration.getModel(), context)).content).toEqual([{ type: "text", text: "fourth" }]);
|
||||
expect(registration.getPendingResponseCount()).toBe(0);
|
||||
});
|
||||
|
||||
it("supports async response factories", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
async (context, _options, state) => fauxAssistantMessage(`${context.messages.length}:${state.callCount}`),
|
||||
]);
|
||||
|
||||
const response = await complete(registration.getModel(), {
|
||||
messages: [{ role: "user", content: "hi", timestamp: Date.now() }],
|
||||
});
|
||||
|
||||
expect(response.content).toEqual([{ type: "text", text: "1:1" }]);
|
||||
});
|
||||
|
||||
it("emits an error when a response factory throws", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
() => {
|
||||
throw new Error("boom");
|
||||
},
|
||||
]);
|
||||
|
||||
const events = await collectEvents(
|
||||
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
|
||||
);
|
||||
|
||||
expect(events).toHaveLength(1);
|
||||
expect(events[0].type).toBe("error");
|
||||
if (events[0].type === "error") {
|
||||
expect(events[0].error.stopReason).toBe("error");
|
||||
expect(events[0].error.errorMessage).toBe("boom");
|
||||
}
|
||||
});
|
||||
|
||||
it("estimates prompt and output tokens from serialized context", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([fauxAssistantMessage("done")]);
|
||||
|
||||
const tool = {
|
||||
name: "echo",
|
||||
description: "Echo back text",
|
||||
parameters: Type.Object({ text: Type.String() }),
|
||||
};
|
||||
const context: Context = {
|
||||
systemPrompt: "sys",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: [
|
||||
{ type: "text", text: "hello" },
|
||||
{ type: "image", mimeType: "image/png", data: "abcd" },
|
||||
],
|
||||
timestamp: 1,
|
||||
},
|
||||
fauxAssistantMessage("prior"),
|
||||
{
|
||||
role: "toolResult",
|
||||
toolCallId: "tool-1",
|
||||
toolName: "echo",
|
||||
content: [{ type: "text", text: "tool out" }],
|
||||
isError: false,
|
||||
timestamp: 2,
|
||||
},
|
||||
],
|
||||
tools: [tool],
|
||||
};
|
||||
|
||||
const response = await complete(registration.getModel(), context);
|
||||
const promptText = [
|
||||
"system:sys",
|
||||
"user:hello\n[image:image/png:4]",
|
||||
"assistant:prior",
|
||||
"toolResult:echo\ntool out",
|
||||
`tools:${JSON.stringify([tool])}`,
|
||||
].join("\n\n");
|
||||
const expectedPromptTokens = Math.ceil(promptText.length / 4);
|
||||
const expectedOutputTokens = Math.ceil("done".length / 4);
|
||||
|
||||
expect(response.usage.input).toBe(expectedPromptTokens);
|
||||
expect(response.usage.output).toBe(expectedOutputTokens);
|
||||
expect(response.usage.cacheRead).toBe(0);
|
||||
expect(response.usage.cacheWrite).toBe(0);
|
||||
expect(response.usage.totalTokens).toBe(expectedPromptTokens + expectedOutputTokens);
|
||||
});
|
||||
|
||||
it("does not share cache across sessions or requests without sessionId", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
fauxAssistantMessage("first"),
|
||||
fauxAssistantMessage("second"),
|
||||
fauxAssistantMessage("third"),
|
||||
]);
|
||||
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
const first = await complete(registration.getModel(), context, {
|
||||
sessionId: "session-1",
|
||||
cacheRetention: "short",
|
||||
});
|
||||
expect(first.usage.cacheWrite).toBeGreaterThan(0);
|
||||
context.messages.push(first);
|
||||
context.messages.push({ role: "user", content: "follow up", timestamp: Date.now() + 1 });
|
||||
|
||||
const second = await complete(registration.getModel(), context, {
|
||||
sessionId: "session-2",
|
||||
cacheRetention: "short",
|
||||
});
|
||||
expect(second.usage.cacheRead).toBe(0);
|
||||
expect(second.usage.cacheWrite).toBeGreaterThan(0);
|
||||
|
||||
const third = await complete(registration.getModel(), context);
|
||||
expect(third.usage.cacheRead).toBe(0);
|
||||
expect(third.usage.cacheWrite).toBe(0);
|
||||
});
|
||||
|
||||
it("simulates prompt caching per sessionId", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "Be concise.",
|
||||
messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
const first = await complete(registration.getModel(), context, {
|
||||
sessionId: "session-1",
|
||||
cacheRetention: "short",
|
||||
});
|
||||
expect(first.usage.cacheRead).toBe(0);
|
||||
expect(first.usage.cacheWrite).toBeGreaterThan(0);
|
||||
|
||||
context.messages.push(first);
|
||||
context.messages.push({ role: "user", content: "follow up", timestamp: Date.now() + 1 });
|
||||
|
||||
const second = await complete(registration.getModel(), context, {
|
||||
sessionId: "session-1",
|
||||
cacheRetention: "short",
|
||||
});
|
||||
expect(second.usage.cacheRead).toBeGreaterThan(0);
|
||||
expect(second.usage.input + second.usage.cacheRead).toBeGreaterThan(second.usage.input);
|
||||
});
|
||||
|
||||
it("does not simulate caching when cacheRetention is none", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([fauxAssistantMessage("first"), fauxAssistantMessage("second")]);
|
||||
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
await complete(registration.getModel(), context, { sessionId: "session-1", cacheRetention: "none" });
|
||||
context.messages.push(fauxAssistantMessage("first"));
|
||||
context.messages.push({ role: "user", content: "follow up", timestamp: Date.now() + 1 });
|
||||
const second = await complete(registration.getModel(), context, {
|
||||
sessionId: "session-1",
|
||||
cacheRetention: "none",
|
||||
});
|
||||
expect(second.usage.cacheRead).toBe(0);
|
||||
expect(second.usage.cacheWrite).toBe(0);
|
||||
});
|
||||
|
||||
it("streams thinking, text, and partial tool call deltas", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
fauxAssistantMessage(
|
||||
[
|
||||
fauxThinking("thinking text"),
|
||||
fauxText("answer text"),
|
||||
fauxToolCall("echo", { text: "hi", count: 12 }, { id: "tool-1" }),
|
||||
],
|
||||
{ stopReason: "toolUse" },
|
||||
),
|
||||
]);
|
||||
|
||||
const events: string[] = [];
|
||||
const toolCallDeltas: string[] = [];
|
||||
const s = stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] });
|
||||
for await (const event of s) {
|
||||
events.push(event.type);
|
||||
if (event.type === "toolcall_delta") {
|
||||
toolCallDeltas.push(event.delta);
|
||||
}
|
||||
}
|
||||
|
||||
expect(events).toContain("thinking_start");
|
||||
expect(events).toContain("thinking_delta");
|
||||
expect(events).toContain("text_start");
|
||||
expect(events).toContain("text_delta");
|
||||
expect(events).toContain("toolcall_start");
|
||||
expect(events).toContain("toolcall_delta");
|
||||
expect(events).toContain("toolcall_end");
|
||||
expect(toolCallDeltas.length).toBeGreaterThan(1);
|
||||
expect(JSON.parse(toolCallDeltas.join(""))).toEqual({ text: "hi", count: 12 });
|
||||
});
|
||||
|
||||
it("streams an exact event order for fixed-size chunks", async () => {
|
||||
const registration = registerFauxProvider({ tokenSize: { min: 1, max: 1 } });
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
fauxAssistantMessage([fauxThinking("go"), fauxText("ok"), fauxToolCall("echo", {}, { id: "tool-1" })], {
|
||||
stopReason: "toolUse",
|
||||
}),
|
||||
]);
|
||||
|
||||
const events = await collectEvents(
|
||||
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
|
||||
);
|
||||
|
||||
expect(events.map((event) => event.type)).toEqual([
|
||||
"start",
|
||||
"thinking_start",
|
||||
"thinking_delta",
|
||||
"thinking_end",
|
||||
"text_start",
|
||||
"text_delta",
|
||||
"text_end",
|
||||
"toolcall_start",
|
||||
"toolcall_delta",
|
||||
"toolcall_end",
|
||||
"done",
|
||||
]);
|
||||
});
|
||||
|
||||
it("streams multiple tool calls in one message", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
fauxAssistantMessage(
|
||||
[
|
||||
fauxToolCall("echo", { text: "one" }, { id: "tool-1" }),
|
||||
fauxToolCall("echo", { text: "two" }, { id: "tool-2" }),
|
||||
],
|
||||
{ stopReason: "toolUse" },
|
||||
),
|
||||
]);
|
||||
|
||||
const events = await collectEvents(
|
||||
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
|
||||
);
|
||||
|
||||
expect(events.filter((event) => event.type === "toolcall_start")).toHaveLength(2);
|
||||
expect(events.filter((event) => event.type === "toolcall_end")).toHaveLength(2);
|
||||
});
|
||||
|
||||
it("streams an explicit assistant error message as a terminal error", async () => {
|
||||
const registration = registerFauxProvider({ tokenSize: { min: 2, max: 2 } });
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
{
|
||||
...fauxAssistantMessage("partial"),
|
||||
stopReason: "error",
|
||||
errorMessage: "upstream failed",
|
||||
},
|
||||
]);
|
||||
|
||||
const events = await collectEvents(
|
||||
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
|
||||
);
|
||||
|
||||
expect(events.map((event) => event.type)).toEqual(["start", "text_start", "text_delta", "text_end", "error"]);
|
||||
const terminal = events[events.length - 1];
|
||||
expect(terminal.type).toBe("error");
|
||||
if (terminal.type === "error") {
|
||||
expect(terminal.reason).toBe("error");
|
||||
expect(terminal.error.stopReason).toBe("error");
|
||||
expect(terminal.error.errorMessage).toBe("upstream failed");
|
||||
}
|
||||
});
|
||||
|
||||
it("streams an explicit assistant aborted message as a terminal error", async () => {
|
||||
const registration = registerFauxProvider({ tokenSize: { min: 2, max: 2 } });
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
{
|
||||
...fauxAssistantMessage("partial"),
|
||||
stopReason: "aborted",
|
||||
errorMessage: "Request was aborted",
|
||||
},
|
||||
]);
|
||||
|
||||
const events = await collectEvents(
|
||||
stream(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
|
||||
);
|
||||
|
||||
expect(events.map((event) => event.type)).toEqual(["start", "text_start", "text_delta", "text_end", "error"]);
|
||||
const terminal = events[events.length - 1];
|
||||
expect(terminal.type).toBe("error");
|
||||
if (terminal.type === "error") {
|
||||
expect(terminal.reason).toBe("aborted");
|
||||
expect(terminal.error.stopReason).toBe("aborted");
|
||||
expect(terminal.error.errorMessage).toBe("Request was aborted");
|
||||
}
|
||||
});
|
||||
|
||||
it("supports aborting before the first chunk", async () => {
|
||||
const registration = registerFauxProvider({ tokensPerSecond: 50, tokenSize: { min: 3, max: 3 } });
|
||||
registrations.push(registration);
|
||||
registration.setResponses([fauxAssistantMessage("abcdefghijklmnopqrstuvwxyz")]);
|
||||
|
||||
const controller = new AbortController();
|
||||
controller.abort();
|
||||
const events = await collectEvents(
|
||||
stream(
|
||||
registration.getModel(),
|
||||
{ messages: [{ role: "user", content: "hi", timestamp: Date.now() }] },
|
||||
{ signal: controller.signal },
|
||||
),
|
||||
);
|
||||
|
||||
expect(events).toHaveLength(1);
|
||||
expect(events[0].type).toBe("error");
|
||||
if (events[0].type === "error") {
|
||||
expect(events[0].reason).toBe("aborted");
|
||||
expect(events[0].error.stopReason).toBe("aborted");
|
||||
}
|
||||
});
|
||||
|
||||
it("supports aborting mid-text stream when paced", async () => {
|
||||
const registration = registerFauxProvider({ tokensPerSecond: 100, tokenSize: { min: 3, max: 3 } });
|
||||
registrations.push(registration);
|
||||
registration.setResponses([fauxAssistantMessage("abcdefghijklmnopqrstuvwxyz")]);
|
||||
|
||||
const controller = new AbortController();
|
||||
const events: string[] = [];
|
||||
let textDeltaCount = 0;
|
||||
const s = stream(
|
||||
registration.getModel(),
|
||||
{ messages: [{ role: "user", content: "hi", timestamp: Date.now() }] },
|
||||
{ signal: controller.signal },
|
||||
);
|
||||
for await (const event of s) {
|
||||
events.push(event.type);
|
||||
if (event.type === "text_delta") {
|
||||
textDeltaCount++;
|
||||
controller.abort();
|
||||
}
|
||||
}
|
||||
|
||||
expect(textDeltaCount).toBe(1);
|
||||
expect(events).toContain("text_start");
|
||||
expect(events).toContain("text_delta");
|
||||
expect(events).toContain("error");
|
||||
expect(events).not.toContain("text_end");
|
||||
});
|
||||
|
||||
it("supports aborting mid-thinking stream when paced", async () => {
|
||||
const registration = registerFauxProvider({ tokensPerSecond: 100, tokenSize: { min: 3, max: 3 } });
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
{
|
||||
...fauxAssistantMessage("ignored"),
|
||||
content: [{ type: "thinking", thinking: "abcdefghijklmnopqrstuvwxyz" }],
|
||||
},
|
||||
]);
|
||||
|
||||
const controller = new AbortController();
|
||||
const events: string[] = [];
|
||||
let thinkingDeltaCount = 0;
|
||||
const s = stream(
|
||||
registration.getModel(),
|
||||
{ messages: [{ role: "user", content: "hi", timestamp: Date.now() }] },
|
||||
{ signal: controller.signal },
|
||||
);
|
||||
for await (const event of s) {
|
||||
events.push(event.type);
|
||||
if (event.type === "thinking_delta") {
|
||||
thinkingDeltaCount++;
|
||||
controller.abort();
|
||||
}
|
||||
}
|
||||
|
||||
expect(thinkingDeltaCount).toBe(1);
|
||||
expect(events).toContain("thinking_start");
|
||||
expect(events).toContain("thinking_delta");
|
||||
expect(events).toContain("error");
|
||||
expect(events).not.toContain("thinking_end");
|
||||
});
|
||||
|
||||
it("supports aborting mid-toolcall stream when paced", async () => {
|
||||
const registration = registerFauxProvider({ tokensPerSecond: 100, tokenSize: { min: 3, max: 3 } });
|
||||
registrations.push(registration);
|
||||
registration.setResponses([
|
||||
{
|
||||
...fauxAssistantMessage("done"),
|
||||
content: [
|
||||
{
|
||||
type: "toolCall",
|
||||
id: "tool-1",
|
||||
name: "echo",
|
||||
arguments: { text: "abcdefghijklmnopqrstuvwxyz", count: 123456789 },
|
||||
},
|
||||
],
|
||||
stopReason: "toolUse",
|
||||
},
|
||||
]);
|
||||
|
||||
const controller = new AbortController();
|
||||
const events: string[] = [];
|
||||
let toolCallDeltaCount = 0;
|
||||
const s = stream(
|
||||
registration.getModel(),
|
||||
{ messages: [{ role: "user", content: "hi", timestamp: Date.now() }] },
|
||||
{ signal: controller.signal },
|
||||
);
|
||||
for await (const event of s) {
|
||||
events.push(event.type);
|
||||
if (event.type === "toolcall_delta") {
|
||||
toolCallDeltaCount++;
|
||||
controller.abort();
|
||||
}
|
||||
}
|
||||
|
||||
expect(toolCallDeltaCount).toBe(1);
|
||||
expect(events).toContain("toolcall_start");
|
||||
expect(events).toContain("toolcall_delta");
|
||||
expect(events).toContain("error");
|
||||
expect(events).not.toContain("toolcall_end");
|
||||
});
|
||||
|
||||
it("unregisters the provider", async () => {
|
||||
const registration = registerFauxProvider();
|
||||
registration.setResponses([fauxAssistantMessage("hello")]);
|
||||
registration.unregister();
|
||||
|
||||
await expect(
|
||||
complete(registration.getModel(), { messages: [{ role: "user", content: "hi", timestamp: Date.now() }] }),
|
||||
).rejects.toThrow(`No API provider registered for api: ${registration.api}`);
|
||||
});
|
||||
});
|
||||
@@ -1,103 +0,0 @@
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import type { Context } from "../src/types.js";
|
||||
|
||||
const mockState = vi.hoisted(() => ({
|
||||
constructorOpts: undefined as Record<string, unknown> | undefined,
|
||||
streamParams: undefined as Record<string, unknown> | undefined,
|
||||
}));
|
||||
|
||||
vi.mock("@anthropic-ai/sdk", () => {
|
||||
const fakeStream = {
|
||||
async *[Symbol.asyncIterator]() {
|
||||
yield {
|
||||
type: "message_start",
|
||||
message: {
|
||||
usage: { input_tokens: 10, output_tokens: 0 },
|
||||
},
|
||||
};
|
||||
yield {
|
||||
type: "message_delta",
|
||||
delta: { stop_reason: "end_turn" },
|
||||
usage: { output_tokens: 5 },
|
||||
};
|
||||
},
|
||||
finalMessage: async () => ({
|
||||
usage: { input_tokens: 10, output_tokens: 5, cache_creation_input_tokens: 0, cache_read_input_tokens: 0 },
|
||||
}),
|
||||
};
|
||||
|
||||
class FakeAnthropic {
|
||||
constructor(opts: Record<string, unknown>) {
|
||||
mockState.constructorOpts = opts;
|
||||
}
|
||||
messages = {
|
||||
stream: (params: Record<string, unknown>) => {
|
||||
mockState.streamParams = params;
|
||||
return fakeStream;
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
return { default: FakeAnthropic };
|
||||
});
|
||||
|
||||
describe("Copilot Claude via Anthropic Messages", () => {
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
it("uses Bearer auth, Copilot headers, and valid Anthropic Messages payload", async () => {
|
||||
const model = getModel("github-copilot", "claude-sonnet-4");
|
||||
expect(model.api).toBe("anthropic-messages");
|
||||
|
||||
const { streamAnthropic } = await import("../src/providers/anthropic.js");
|
||||
const s = streamAnthropic(model, context, { apiKey: "tid_copilot_session_test_token" });
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") break;
|
||||
}
|
||||
|
||||
const opts = mockState.constructorOpts!;
|
||||
expect(opts).toBeDefined();
|
||||
|
||||
// Auth: apiKey null, authToken for Bearer
|
||||
expect(opts.apiKey).toBeNull();
|
||||
expect(opts.authToken).toBe("tid_copilot_session_test_token");
|
||||
const headers = opts.defaultHeaders as Record<string, string>;
|
||||
|
||||
// Copilot static headers from model.headers
|
||||
expect(headers["User-Agent"]).toContain("GitHubCopilotChat");
|
||||
expect(headers["Copilot-Integration-Id"]).toBe("vscode-chat");
|
||||
|
||||
// Dynamic headers
|
||||
expect(headers["X-Initiator"]).toBe("user");
|
||||
expect(headers["Openai-Intent"]).toBe("conversation-edits");
|
||||
|
||||
// No fine-grained-tool-streaming (Copilot doesn't support it)
|
||||
const beta = headers["anthropic-beta"] ?? "";
|
||||
expect(beta).not.toContain("fine-grained-tool-streaming");
|
||||
|
||||
// Payload is valid Anthropic Messages format
|
||||
const params = mockState.streamParams!;
|
||||
expect(params.model).toBe("claude-sonnet-4");
|
||||
expect(params.stream).toBe(true);
|
||||
expect(params.max_tokens).toBeGreaterThan(0);
|
||||
expect(Array.isArray(params.messages)).toBe(true);
|
||||
});
|
||||
|
||||
it("includes interleaved-thinking beta when reasoning is enabled", async () => {
|
||||
const model = getModel("github-copilot", "claude-sonnet-4");
|
||||
const { streamAnthropic } = await import("../src/providers/anthropic.js");
|
||||
const s = streamAnthropic(model, context, {
|
||||
apiKey: "tid_copilot_session_test_token",
|
||||
interleavedThinking: true,
|
||||
});
|
||||
for await (const event of s) {
|
||||
if (event.type === "error") break;
|
||||
}
|
||||
|
||||
const headers = mockState.constructorOpts!.defaultHeaders as Record<string, string>;
|
||||
expect(headers["anthropic-beta"]).toContain("interleaved-thinking-2025-05-14");
|
||||
});
|
||||
});
|
||||
@@ -1,196 +0,0 @@
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { loginGitHubCopilot } from "../src/utils/oauth/github-copilot.js";
|
||||
|
||||
function jsonResponse(body: unknown, status: number = 200): Response {
|
||||
return new Response(JSON.stringify(body), {
|
||||
status,
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function getUrl(input: unknown): string {
|
||||
if (typeof input === "string") {
|
||||
return input;
|
||||
}
|
||||
if (input instanceof URL) {
|
||||
return input.toString();
|
||||
}
|
||||
if (input instanceof Request) {
|
||||
return input.url;
|
||||
}
|
||||
throw new Error(`Unsupported fetch input: ${String(input)}`);
|
||||
}
|
||||
|
||||
describe("GitHub Copilot OAuth device flow", () => {
|
||||
afterEach(() => {
|
||||
vi.unstubAllGlobals();
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it("waits before the first poll and increases the safety margin after slow_down", async () => {
|
||||
vi.useFakeTimers();
|
||||
const startTime = new Date("2026-03-09T00:00:00Z");
|
||||
vi.setSystemTime(startTime);
|
||||
|
||||
const accessTokenPollTimes: number[] = [];
|
||||
const accessTokenResponses = [
|
||||
jsonResponse({ error: "authorization_pending", error_description: "pending" }),
|
||||
jsonResponse({ error: "slow_down", error_description: "slow down", interval: 10 }),
|
||||
jsonResponse({ access_token: "ghu_refresh_token" }),
|
||||
];
|
||||
|
||||
const fetchMock = vi.fn(async (input: unknown, init?: RequestInit): Promise<Response> => {
|
||||
const url = getUrl(input);
|
||||
|
||||
if (url.endsWith("/login/device/code")) {
|
||||
expect(init?.method).toBe("POST");
|
||||
expect(init?.headers).toMatchObject({
|
||||
Accept: "application/json",
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
});
|
||||
expect(String(init?.body)).toContain("client_id=");
|
||||
expect(String(init?.body)).toContain("scope=read%3Auser");
|
||||
return jsonResponse({
|
||||
device_code: "device-code",
|
||||
user_code: "ABCD-EFGH",
|
||||
verification_uri: "https://github.com/login/device",
|
||||
interval: 5,
|
||||
expires_in: 900,
|
||||
});
|
||||
}
|
||||
|
||||
if (url.endsWith("/login/oauth/access_token")) {
|
||||
accessTokenPollTimes.push(Date.now());
|
||||
expect(init?.method).toBe("POST");
|
||||
expect(init?.headers).toMatchObject({
|
||||
Accept: "application/json",
|
||||
"Content-Type": "application/x-www-form-urlencoded",
|
||||
});
|
||||
expect(String(init?.body)).toContain("client_id=");
|
||||
expect(String(init?.body)).toContain("device_code=device-code");
|
||||
expect(String(init?.body)).toContain("grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code");
|
||||
const response = accessTokenResponses.shift();
|
||||
if (!response) {
|
||||
throw new Error("Unexpected extra access token poll");
|
||||
}
|
||||
return response;
|
||||
}
|
||||
|
||||
if (url.includes("/copilot_internal/v2/token")) {
|
||||
return jsonResponse({
|
||||
token: "tid=test;exp=9999999999;proxy-ep=proxy.individual.githubcopilot.com;",
|
||||
expires_at: 9999999999,
|
||||
});
|
||||
}
|
||||
|
||||
if (url.includes("/models/") && url.endsWith("/policy")) {
|
||||
return new Response("", { status: 200 });
|
||||
}
|
||||
|
||||
throw new Error(`Unexpected fetch URL: ${url}`);
|
||||
});
|
||||
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
const loginPromise = loginGitHubCopilot({
|
||||
onAuth: () => {},
|
||||
onPrompt: async () => "",
|
||||
onProgress: () => {},
|
||||
});
|
||||
|
||||
await vi.advanceTimersByTimeAsync(0);
|
||||
expect(accessTokenPollTimes).toHaveLength(0);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(5999);
|
||||
expect(accessTokenPollTimes).toHaveLength(0);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(1);
|
||||
expect(accessTokenPollTimes).toHaveLength(1);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(5999);
|
||||
expect(accessTokenPollTimes).toHaveLength(1);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(1);
|
||||
expect(accessTokenPollTimes).toHaveLength(2);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(13999);
|
||||
expect(accessTokenPollTimes).toHaveLength(2);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(1);
|
||||
await loginPromise;
|
||||
|
||||
expect(accessTokenPollTimes).toEqual([
|
||||
startTime.getTime() + 6000,
|
||||
startTime.getTime() + 12000,
|
||||
startTime.getTime() + 26000,
|
||||
]);
|
||||
});
|
||||
|
||||
it("uses the remaining lifetime for a final poll before timing out after repeated slow_down responses", async () => {
|
||||
vi.useFakeTimers();
|
||||
const startTime = new Date("2026-03-09T00:00:00Z");
|
||||
vi.setSystemTime(startTime);
|
||||
|
||||
const accessTokenPollTimes: number[] = [];
|
||||
const accessTokenResponses = [
|
||||
jsonResponse({ error: "slow_down", error_description: "slow down", interval: 10 }),
|
||||
jsonResponse({ error: "slow_down", error_description: "still too fast", interval: 15 }),
|
||||
jsonResponse({ error: "authorization_pending", error_description: "pending" }),
|
||||
];
|
||||
|
||||
const fetchMock = vi.fn(async (input: unknown): Promise<Response> => {
|
||||
const url = getUrl(input);
|
||||
|
||||
if (url.endsWith("/login/device/code")) {
|
||||
return jsonResponse({
|
||||
device_code: "device-code",
|
||||
user_code: "ABCD-EFGH",
|
||||
verification_uri: "https://github.com/login/device",
|
||||
interval: 5,
|
||||
expires_in: 25,
|
||||
});
|
||||
}
|
||||
|
||||
if (url.endsWith("/login/oauth/access_token")) {
|
||||
accessTokenPollTimes.push(Date.now());
|
||||
const response = accessTokenResponses.shift();
|
||||
if (!response) {
|
||||
throw new Error("Unexpected extra access token poll");
|
||||
}
|
||||
return response;
|
||||
}
|
||||
|
||||
throw new Error(`Unexpected fetch URL: ${url}`);
|
||||
});
|
||||
|
||||
vi.stubGlobal("fetch", fetchMock);
|
||||
|
||||
const loginPromise = loginGitHubCopilot({
|
||||
onAuth: () => {},
|
||||
onPrompt: async () => "",
|
||||
});
|
||||
const rejection = expect(loginPromise).rejects.toThrow(
|
||||
/Device flow timed out after one or more slow_down responses/,
|
||||
);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(6000);
|
||||
expect(accessTokenPollTimes).toEqual([startTime.getTime() + 6000]);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(14000);
|
||||
expect(accessTokenPollTimes).toEqual([startTime.getTime() + 6000, startTime.getTime() + 20000]);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(4999);
|
||||
expect(accessTokenPollTimes).toEqual([startTime.getTime() + 6000, startTime.getTime() + 20000]);
|
||||
|
||||
await vi.advanceTimersByTimeAsync(1);
|
||||
await rejection;
|
||||
|
||||
expect(accessTokenPollTimes).toEqual([
|
||||
startTime.getTime() + 6000,
|
||||
startTime.getTime() + 20000,
|
||||
startTime.getTime() + 25000,
|
||||
]);
|
||||
});
|
||||
});
|
||||
@@ -1,103 +0,0 @@
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { streamGoogleGeminiCli } from "../src/providers/google-gemini-cli.js";
|
||||
import type { Context, Model } from "../src/types.js";
|
||||
|
||||
const originalFetch = global.fetch;
|
||||
const apiKey = JSON.stringify({ token: "token", projectId: "project" });
|
||||
|
||||
const createSseResponse = () => {
|
||||
const sse = `${[
|
||||
`data: ${JSON.stringify({
|
||||
response: {
|
||||
candidates: [
|
||||
{
|
||||
content: { role: "model", parts: [{ text: "Hello" }] },
|
||||
finishReason: "STOP",
|
||||
},
|
||||
],
|
||||
},
|
||||
})}`,
|
||||
].join("\n\n")}\n\n`;
|
||||
|
||||
const encoder = new TextEncoder();
|
||||
const stream = new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
controller.enqueue(encoder.encode(sse));
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
|
||||
return new Response(stream, {
|
||||
status: 200,
|
||||
headers: { "content-type": "text/event-stream" },
|
||||
});
|
||||
};
|
||||
|
||||
afterEach(() => {
|
||||
global.fetch = originalFetch;
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("google-gemini-cli Claude thinking header", () => {
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
it("adds anthropic-beta for Claude thinking models", async () => {
|
||||
const fetchMock = vi.fn(async (_input: string | URL, init?: RequestInit) => {
|
||||
const headers = new Headers(init?.headers);
|
||||
expect(headers.get("anthropic-beta")).toBe("interleaved-thinking-2025-05-14");
|
||||
return createSseResponse();
|
||||
});
|
||||
|
||||
global.fetch = fetchMock as typeof fetch;
|
||||
|
||||
const model: Model<"google-gemini-cli"> = {
|
||||
id: "claude-opus-4-5-thinking",
|
||||
name: "Claude Opus 4.5 Thinking",
|
||||
api: "google-gemini-cli",
|
||||
provider: "google-antigravity",
|
||||
baseUrl: "https://cloudcode-pa.googleapis.com",
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
};
|
||||
|
||||
const stream = streamGoogleGeminiCli(model, context, { apiKey });
|
||||
for await (const _event of stream) {
|
||||
// exhaust stream
|
||||
}
|
||||
await stream.result();
|
||||
});
|
||||
|
||||
it("does not add anthropic-beta for Gemini models", async () => {
|
||||
const fetchMock = vi.fn(async (_input: string | URL, init?: RequestInit) => {
|
||||
const headers = new Headers(init?.headers);
|
||||
expect(headers.has("anthropic-beta")).toBe(false);
|
||||
return createSseResponse();
|
||||
});
|
||||
|
||||
global.fetch = fetchMock as typeof fetch;
|
||||
|
||||
const model: Model<"google-gemini-cli"> = {
|
||||
id: "gemini-2.5-flash",
|
||||
name: "Gemini 2.5 Flash",
|
||||
api: "google-gemini-cli",
|
||||
provider: "google-gemini-cli",
|
||||
baseUrl: "https://cloudcode-pa.googleapis.com",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
};
|
||||
|
||||
const stream = streamGoogleGeminiCli(model, context, { apiKey });
|
||||
for await (const _event of stream) {
|
||||
// exhaust stream
|
||||
}
|
||||
await stream.result();
|
||||
});
|
||||
});
|
||||
@@ -1,108 +0,0 @@
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { streamGoogleGeminiCli } from "../src/providers/google-gemini-cli.js";
|
||||
import type { Context, Model } from "../src/types.js";
|
||||
|
||||
const originalFetch = global.fetch;
|
||||
|
||||
afterEach(() => {
|
||||
global.fetch = originalFetch;
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("google-gemini-cli empty stream retry", () => {
|
||||
it("retries empty SSE responses without duplicate start", async () => {
|
||||
const emptyStream = new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
|
||||
const sse = `${[
|
||||
`data: ${JSON.stringify({
|
||||
response: {
|
||||
candidates: [
|
||||
{
|
||||
content: { role: "model", parts: [{ text: "Hello" }] },
|
||||
finishReason: "STOP",
|
||||
},
|
||||
],
|
||||
usageMetadata: {
|
||||
promptTokenCount: 1,
|
||||
candidatesTokenCount: 1,
|
||||
totalTokenCount: 2,
|
||||
},
|
||||
},
|
||||
})}`,
|
||||
].join("\n\n")}\n\n`;
|
||||
|
||||
const encoder = new TextEncoder();
|
||||
const dataStream = new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
controller.enqueue(encoder.encode(sse));
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
|
||||
let callCount = 0;
|
||||
const fetchMock = vi.fn(async () => {
|
||||
callCount += 1;
|
||||
if (callCount === 1) {
|
||||
return new Response(emptyStream, {
|
||||
status: 200,
|
||||
headers: { "content-type": "text/event-stream" },
|
||||
});
|
||||
}
|
||||
return new Response(dataStream, {
|
||||
status: 200,
|
||||
headers: { "content-type": "text/event-stream" },
|
||||
});
|
||||
});
|
||||
|
||||
global.fetch = fetchMock as typeof fetch;
|
||||
|
||||
const model: Model<"google-gemini-cli"> = {
|
||||
id: "gemini-2.5-flash",
|
||||
name: "Gemini 2.5 Flash",
|
||||
api: "google-gemini-cli",
|
||||
provider: "google-gemini-cli",
|
||||
baseUrl: "https://cloudcode-pa.googleapis.com",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
const stream = streamGoogleGeminiCli(model, context, {
|
||||
apiKey: JSON.stringify({ token: "token", projectId: "project" }),
|
||||
});
|
||||
|
||||
let startCount = 0;
|
||||
let doneCount = 0;
|
||||
let text = "";
|
||||
|
||||
for await (const event of stream) {
|
||||
if (event.type === "start") {
|
||||
startCount += 1;
|
||||
}
|
||||
if (event.type === "done") {
|
||||
doneCount += 1;
|
||||
}
|
||||
if (event.type === "text_delta") {
|
||||
text += event.delta;
|
||||
}
|
||||
}
|
||||
|
||||
const result = await stream.result();
|
||||
|
||||
expect(text).toBe("Hello");
|
||||
expect(result.stopReason).toBe("stop");
|
||||
expect(startCount).toBe(1);
|
||||
expect(doneCount).toBe(1);
|
||||
expect(fetchMock).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
});
|
||||
@@ -1,53 +0,0 @@
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { extractRetryDelay } from "../src/providers/google-gemini-cli.js";
|
||||
|
||||
describe("extractRetryDelay header parsing", () => {
|
||||
afterEach(() => {
|
||||
vi.useRealTimers();
|
||||
});
|
||||
|
||||
it("prefers Retry-After seconds header", () => {
|
||||
vi.useFakeTimers();
|
||||
vi.setSystemTime(new Date("2025-01-01T00:00:00Z"));
|
||||
|
||||
const response = new Response("", { headers: { "Retry-After": "5" } });
|
||||
const delay = extractRetryDelay("Please retry in 1s", response);
|
||||
|
||||
expect(delay).toBe(6000);
|
||||
});
|
||||
|
||||
it("parses Retry-After HTTP date header", () => {
|
||||
vi.useFakeTimers();
|
||||
const now = new Date("2025-01-01T00:00:00Z");
|
||||
vi.setSystemTime(now);
|
||||
|
||||
const retryAt = new Date(now.getTime() + 12000).toUTCString();
|
||||
const response = new Response("", { headers: { "Retry-After": retryAt } });
|
||||
const delay = extractRetryDelay("", response);
|
||||
|
||||
expect(delay).toBe(13000);
|
||||
});
|
||||
|
||||
it("parses x-ratelimit-reset header", () => {
|
||||
vi.useFakeTimers();
|
||||
const now = new Date("2025-01-01T00:00:00Z");
|
||||
vi.setSystemTime(now);
|
||||
|
||||
const resetAtMs = now.getTime() + 20000;
|
||||
const resetSeconds = Math.floor(resetAtMs / 1000).toString();
|
||||
const response = new Response("", { headers: { "x-ratelimit-reset": resetSeconds } });
|
||||
const delay = extractRetryDelay("", response);
|
||||
|
||||
expect(delay).toBe(21000);
|
||||
});
|
||||
|
||||
it("parses x-ratelimit-reset-after header", () => {
|
||||
vi.useFakeTimers();
|
||||
vi.setSystemTime(new Date("2025-01-01T00:00:00Z"));
|
||||
|
||||
const response = new Response("", { headers: { "x-ratelimit-reset-after": "30" } });
|
||||
const delay = extractRetryDelay("", response);
|
||||
|
||||
expect(delay).toBe(31000);
|
||||
});
|
||||
});
|
||||
@@ -1,172 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { convertMessages } from "../src/providers/google-shared.js";
|
||||
import type { Context, Model } from "../src/types.js";
|
||||
|
||||
const SKIP_THOUGHT_SIGNATURE = "skip_thought_signature_validator";
|
||||
|
||||
function makeGemini3Model(id = "gemini-3-pro-preview"): Model<"google-generative-ai"> {
|
||||
return {
|
||||
id,
|
||||
name: "Gemini 3 Pro Preview",
|
||||
api: "google-generative-ai",
|
||||
provider: "google",
|
||||
baseUrl: "https://generativelanguage.googleapis.com",
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
};
|
||||
}
|
||||
|
||||
describe("google-shared convertMessages — Gemini 3 unsigned tool calls", () => {
|
||||
it("uses skip_thought_signature_validator for unsigned tool calls on Gemini 3", () => {
|
||||
const model = makeGemini3Model();
|
||||
const now = Date.now();
|
||||
const context: Context = {
|
||||
messages: [
|
||||
{ role: "user", content: "Hi", timestamp: now },
|
||||
{
|
||||
role: "assistant",
|
||||
content: [
|
||||
{
|
||||
type: "toolCall",
|
||||
id: "call_1",
|
||||
name: "bash",
|
||||
arguments: { command: "ls -la" },
|
||||
// No thoughtSignature: simulates Claude via Antigravity.
|
||||
},
|
||||
],
|
||||
api: "google-gemini-cli",
|
||||
provider: "google-antigravity",
|
||||
model: "claude-sonnet-4-20250514",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: now,
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const contents = convertMessages(model, context);
|
||||
|
||||
const modelTurn = contents.find((c) => c.role === "model");
|
||||
expect(modelTurn).toBeTruthy();
|
||||
|
||||
// Should be a structured functionCall, NOT text fallback
|
||||
const fcPart = modelTurn?.parts?.find((p) => p.functionCall !== undefined);
|
||||
expect(fcPart).toBeTruthy();
|
||||
expect(fcPart?.functionCall?.name).toBe("bash");
|
||||
expect(fcPart?.functionCall?.args).toEqual({ command: "ls -la" });
|
||||
expect(fcPart?.thoughtSignature).toBe(SKIP_THOUGHT_SIGNATURE);
|
||||
|
||||
// No text fallback should exist
|
||||
const textParts = modelTurn?.parts?.filter((p) => p.text !== undefined) ?? [];
|
||||
const historicalText = textParts.filter((p) => p.text?.includes("Historical context"));
|
||||
expect(historicalText).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("preserves valid thoughtSignature when present (same provider/model)", () => {
|
||||
const model = makeGemini3Model();
|
||||
const now = Date.now();
|
||||
// Valid base64 signature (16 bytes = 24 chars base64)
|
||||
const validSig = "AAAAAAAAAAAAAAAAAAAAAA==";
|
||||
const context: Context = {
|
||||
messages: [
|
||||
{ role: "user", content: "Hi", timestamp: now },
|
||||
{
|
||||
role: "assistant",
|
||||
content: [
|
||||
{
|
||||
type: "toolCall",
|
||||
id: "call_1",
|
||||
name: "bash",
|
||||
arguments: { command: "echo hi" },
|
||||
thoughtSignature: validSig,
|
||||
},
|
||||
],
|
||||
api: "google-generative-ai",
|
||||
provider: "google",
|
||||
model: "gemini-3-pro-preview",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: now,
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const contents = convertMessages(model, context);
|
||||
const modelTurn = contents.find((c) => c.role === "model");
|
||||
const fcPart = modelTurn?.parts?.find((p) => p.functionCall !== undefined);
|
||||
|
||||
expect(fcPart).toBeTruthy();
|
||||
expect(fcPart?.thoughtSignature).toBe(validSig);
|
||||
});
|
||||
|
||||
it("does not add sentinel for non-Gemini-3 models", () => {
|
||||
const model: Model<"google-generative-ai"> = {
|
||||
id: "gemini-2.5-flash",
|
||||
name: "Gemini 2.5 Flash",
|
||||
api: "google-generative-ai",
|
||||
provider: "google",
|
||||
baseUrl: "https://generativelanguage.googleapis.com",
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
};
|
||||
const now = Date.now();
|
||||
const context: Context = {
|
||||
messages: [
|
||||
{ role: "user", content: "Hi", timestamp: now },
|
||||
{
|
||||
role: "assistant",
|
||||
content: [
|
||||
{
|
||||
type: "toolCall",
|
||||
id: "call_1",
|
||||
name: "bash",
|
||||
arguments: { command: "ls" },
|
||||
// No thoughtSignature
|
||||
},
|
||||
],
|
||||
api: "google-gemini-cli",
|
||||
provider: "google-antigravity",
|
||||
model: "claude-sonnet-4-20250514",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: now,
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const contents = convertMessages(model, context);
|
||||
const modelTurn = contents.find((c) => c.role === "model");
|
||||
const fcPart = modelTurn?.parts?.find((p) => p.functionCall !== undefined);
|
||||
|
||||
expect(fcPart).toBeTruthy();
|
||||
// No sentinel, no thoughtSignature at all
|
||||
expect(fcPart?.thoughtSignature).toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -1,124 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { convertMessages } from "../src/providers/google-shared.js";
|
||||
import type { Context, Model } from "../src/types.js";
|
||||
|
||||
function makeModel<TApi extends "google-generative-ai" | "google-gemini-cli">(
|
||||
api: TApi,
|
||||
provider: Model<TApi>["provider"],
|
||||
id: string,
|
||||
): Model<TApi> {
|
||||
return {
|
||||
id,
|
||||
name: id,
|
||||
api,
|
||||
provider,
|
||||
baseUrl: "https://example.com",
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
};
|
||||
}
|
||||
|
||||
function makeContext(model: { api: string; provider: string; id: string }): Context {
|
||||
const now = Date.now();
|
||||
return {
|
||||
messages: [
|
||||
{ role: "user", content: "read the files", timestamp: now },
|
||||
{
|
||||
role: "assistant",
|
||||
content: [
|
||||
{ type: "toolCall", id: "call_a", name: "read", arguments: { path: "a.txt" } },
|
||||
{ type: "toolCall", id: "call_img", name: "read", arguments: { path: "image.png" } },
|
||||
{ type: "toolCall", id: "call_b", name: "read", arguments: { path: "b.txt" } },
|
||||
],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "toolUse",
|
||||
timestamp: now,
|
||||
},
|
||||
{
|
||||
role: "toolResult",
|
||||
toolCallId: "call_a",
|
||||
toolName: "read",
|
||||
content: [{ type: "text", text: "alpha text" }],
|
||||
isError: false,
|
||||
timestamp: now,
|
||||
},
|
||||
{
|
||||
role: "toolResult",
|
||||
toolCallId: "call_img",
|
||||
toolName: "read",
|
||||
content: [{ type: "image", data: "abc", mimeType: "image/png" }],
|
||||
isError: false,
|
||||
timestamp: now,
|
||||
},
|
||||
{
|
||||
role: "toolResult",
|
||||
toolCallId: "call_b",
|
||||
toolName: "read",
|
||||
content: [{ type: "text", text: "beta text" }],
|
||||
isError: false,
|
||||
timestamp: now,
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
describe("google-shared image tool result routing", () => {
|
||||
it("keeps separate synthetic image turn for Gemini 2.x Google API models", () => {
|
||||
const model = makeModel("google-generative-ai", "google", "gemini-2.5-flash");
|
||||
const contents = convertMessages(model, makeContext(model));
|
||||
|
||||
expect(contents).toHaveLength(5);
|
||||
expect(contents[2].parts?.every((part) => part.functionResponse)).toBe(true);
|
||||
expect(contents[3].parts?.[0]?.text).toBe("Tool result image:");
|
||||
expect(contents[3].parts?.[1]?.inlineData).toBeTruthy();
|
||||
expect(contents[4].parts?.[0]?.functionResponse).toBeTruthy();
|
||||
});
|
||||
|
||||
it("nests image tool results for Gemini 3 Google API models", () => {
|
||||
const model = makeModel("google-generative-ai", "google", "gemini-3-pro-preview");
|
||||
const contents = convertMessages(model, makeContext(model));
|
||||
|
||||
expect(contents).toHaveLength(3);
|
||||
const toolResultTurn = contents[2];
|
||||
expect(toolResultTurn.parts).toHaveLength(3);
|
||||
const imageResponse = toolResultTurn.parts?.[1]?.functionResponse;
|
||||
expect(imageResponse).toBeTruthy();
|
||||
expect(imageResponse?.parts).toHaveLength(1);
|
||||
expect(imageResponse?.parts?.[0]?.inlineData).toBeTruthy();
|
||||
});
|
||||
|
||||
it("nests image tool results for non-Gemini models on Antigravity / Cloud Code Assist", () => {
|
||||
const model = makeModel("google-gemini-cli", "google-antigravity", "claude-sonnet-4-6");
|
||||
const contents = convertMessages(model, makeContext(model));
|
||||
|
||||
expect(contents).toHaveLength(3);
|
||||
const toolResultTurn = contents[2];
|
||||
expect(toolResultTurn.parts).toHaveLength(3);
|
||||
const imageResponse = toolResultTurn.parts?.[1]?.functionResponse;
|
||||
expect(imageResponse).toBeTruthy();
|
||||
expect(imageResponse?.parts).toHaveLength(1);
|
||||
expect(imageResponse?.parts?.[0]?.inlineData).toBeTruthy();
|
||||
});
|
||||
|
||||
it("keeps separate synthetic image turn for Gemini 2.x Cloud Code Assist models", () => {
|
||||
const model = makeModel("google-gemini-cli", "google-gemini-cli", "gemini-2.5-flash");
|
||||
const contents = convertMessages(model, makeContext(model));
|
||||
|
||||
expect(contents).toHaveLength(5);
|
||||
expect(contents[3].parts?.[0]?.text).toBe("Tool result image:");
|
||||
expect(contents[3].parts?.[1]?.inlineData).toBeTruthy();
|
||||
});
|
||||
});
|
||||
@@ -1,183 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { streamSimple } from "../src/stream.js";
|
||||
import type { Api, Context, Model, SimpleStreamOptions } from "../src/types.js";
|
||||
import { resolveApiKey } from "./oauth.js";
|
||||
|
||||
type SimpleOptionsWithExtras = SimpleStreamOptions & Record<string, unknown>;
|
||||
|
||||
interface RunResult {
|
||||
thinkingEventCount: number;
|
||||
thinkingCharCount: number;
|
||||
text: string;
|
||||
outputTokens: number;
|
||||
contentTypes: string[];
|
||||
}
|
||||
|
||||
interface DisableExpectations {
|
||||
requestOptions?: SimpleOptionsWithExtras;
|
||||
minPongs?: number;
|
||||
maxOutputTokens?: number;
|
||||
}
|
||||
|
||||
const oauthTokens = await Promise.all([resolveApiKey("google-gemini-cli"), resolveApiKey("google-antigravity")]);
|
||||
const [geminiCliToken, antigravityToken] = oauthTokens;
|
||||
|
||||
function makeContext(): Context {
|
||||
return {
|
||||
systemPrompt: "You are a precise assistant. Follow the requested output format exactly.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content:
|
||||
"Before replying, carefully solve 36863 * 5279 internally. Then reply with the word pong repeated exactly 40 times, separated by single spaces. Do not add any other text.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
|
||||
function countPongs(text: string): number {
|
||||
return text.match(/\bpong\b/gi)?.length ?? 0;
|
||||
}
|
||||
|
||||
async function runWithoutReasoning<TApi extends Api>(
|
||||
model: Model<TApi>,
|
||||
options: SimpleOptionsWithExtras = {},
|
||||
): Promise<RunResult> {
|
||||
const s = streamSimple(model, makeContext(), {
|
||||
maxTokens: 160,
|
||||
temperature: 0,
|
||||
...options,
|
||||
});
|
||||
|
||||
let thinkingEventCount = 0;
|
||||
let thinkingCharCount = 0;
|
||||
|
||||
for await (const event of s) {
|
||||
if (event.type === "thinking_start" || event.type === "thinking_end") {
|
||||
thinkingEventCount += 1;
|
||||
}
|
||||
if (event.type === "thinking_delta") {
|
||||
thinkingEventCount += 1;
|
||||
thinkingCharCount += event.delta.length;
|
||||
}
|
||||
}
|
||||
|
||||
const response = await s.result();
|
||||
expect(response.stopReason, response.errorMessage).toBe("stop");
|
||||
|
||||
const text = response.content
|
||||
.filter((block) => block.type === "text")
|
||||
.map((block) => block.text)
|
||||
.join("")
|
||||
.trim();
|
||||
|
||||
return {
|
||||
thinkingEventCount,
|
||||
thinkingCharCount,
|
||||
text,
|
||||
outputTokens: response.usage.output,
|
||||
contentTypes: response.content.map((block) => block.type),
|
||||
};
|
||||
}
|
||||
|
||||
async function expectThinkingDisabledE2E<TApi extends Api>(model: Model<TApi>, expectations: DisableExpectations = {}) {
|
||||
const result = await runWithoutReasoning(model, expectations.requestOptions);
|
||||
|
||||
expect(result.thinkingEventCount).toBe(0);
|
||||
expect(result.thinkingCharCount).toBe(0);
|
||||
expect(result.contentTypes).not.toContain("thinking");
|
||||
expect(countPongs(result.text)).toBeGreaterThanOrEqual(expectations.minPongs ?? 35);
|
||||
if (expectations.maxOutputTokens !== undefined) {
|
||||
expect(result.outputTokens).toBeLessThan(expectations.maxOutputTokens);
|
||||
}
|
||||
}
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic thinking disable E2E", () => {
|
||||
it("disables thinking for budget-based reasoning models", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("anthropic", "claude-sonnet-4-5"), {
|
||||
requestOptions: { maxTokens: 320, temperature: 0 },
|
||||
});
|
||||
});
|
||||
|
||||
it("disables thinking for adaptive reasoning models", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("anthropic", "claude-sonnet-4-6"), {
|
||||
requestOptions: { maxTokens: 320, temperature: 0 },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.GEMINI_API_KEY)("Google thinking disable E2E", () => {
|
||||
it("disables thinking for Gemini 2.5", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("google", "gemini-2.5-flash"));
|
||||
});
|
||||
|
||||
it("disables thinking for Gemini 3.x", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("google", "gemini-3-flash-preview"));
|
||||
});
|
||||
|
||||
it("does not error when thinking is off for Gemini 3.1 Pro", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("google", "gemini-3.1-pro-preview"), {
|
||||
requestOptions: { maxTokens: 512 },
|
||||
minPongs: 20,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("Google Vertex thinking disable E2E", () => {
|
||||
const vertexProject = process.env.GOOGLE_CLOUD_PROJECT || process.env.GCLOUD_PROJECT;
|
||||
const vertexLocation = process.env.GOOGLE_CLOUD_LOCATION;
|
||||
const vertexApiKey = process.env.GOOGLE_CLOUD_API_KEY;
|
||||
const vertexOptions = vertexApiKey
|
||||
? ({ apiKey: vertexApiKey } satisfies SimpleOptionsWithExtras)
|
||||
: vertexProject && vertexLocation
|
||||
? ({ project: vertexProject, location: vertexLocation } satisfies SimpleOptionsWithExtras)
|
||||
: undefined;
|
||||
|
||||
it.skipIf(!vertexOptions)("disables thinking for Gemini 2.5", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("google-vertex", "gemini-2.5-flash"), {
|
||||
requestOptions: vertexOptions,
|
||||
});
|
||||
});
|
||||
|
||||
it.skipIf(!vertexOptions)("disables thinking for Gemini 3.x", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("google-vertex", "gemini-3-flash-preview"), {
|
||||
requestOptions: vertexOptions,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("Google Gemini CLI thinking disable E2E", () => {
|
||||
it.skipIf(!geminiCliToken)("disables thinking for Gemini 2.5", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("google-gemini-cli", "gemini-2.5-flash"), {
|
||||
requestOptions: { apiKey: geminiCliToken! },
|
||||
maxOutputTokens: 100,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("Google Antigravity thinking disable E2E", () => {
|
||||
it.skipIf(!antigravityToken)("disables thinking for Gemini 3.x", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("google-antigravity", "gemini-3-flash"), {
|
||||
requestOptions: { apiKey: antigravityToken! },
|
||||
maxOutputTokens: 100,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI thinking disable E2E", () => {
|
||||
it("disables thinking for Responses reasoning models", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("openai", "gpt-5.4-mini"), {
|
||||
requestOptions: { temperature: undefined },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENROUTER_API_KEY)("OpenRouter thinking disable E2E", () => {
|
||||
it("disables thinking for Qwen 3.5 reasoning models", { retry: 2, timeout: 30000 }, async () => {
|
||||
await expectThinkingDisabledE2E(getModel("openrouter", "qwen/qwen3.5-plus-02-15"), {
|
||||
maxOutputTokens: 100,
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,38 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { isThinkingPart, retainThoughtSignature } from "../src/providers/google-shared.js";
|
||||
|
||||
describe("Google thinking detection (thoughtSignature)", () => {
|
||||
it("treats part.thought === true as thinking", () => {
|
||||
expect(isThinkingPart({ thought: true, thoughtSignature: undefined })).toBe(true);
|
||||
expect(isThinkingPart({ thought: true, thoughtSignature: "opaque-signature" })).toBe(true);
|
||||
});
|
||||
|
||||
it("does not treat thoughtSignature alone as thinking", () => {
|
||||
// Per Google docs, thoughtSignature is for context replay and can appear on any part type.
|
||||
// Only thought === true indicates thinking content.
|
||||
// See: https://ai.google.dev/gemini-api/docs/thought-signatures
|
||||
expect(isThinkingPart({ thought: undefined, thoughtSignature: "opaque-signature" })).toBe(false);
|
||||
expect(isThinkingPart({ thought: false, thoughtSignature: "opaque-signature" })).toBe(false);
|
||||
});
|
||||
|
||||
it("does not treat empty/missing signatures as thinking if thought is not set", () => {
|
||||
expect(isThinkingPart({ thought: undefined, thoughtSignature: undefined })).toBe(false);
|
||||
expect(isThinkingPart({ thought: false, thoughtSignature: "" })).toBe(false);
|
||||
});
|
||||
|
||||
it("preserves the existing signature when subsequent deltas omit thoughtSignature", () => {
|
||||
const first = retainThoughtSignature(undefined, "sig-1");
|
||||
expect(first).toBe("sig-1");
|
||||
|
||||
const second = retainThoughtSignature(first, undefined);
|
||||
expect(second).toBe("sig-1");
|
||||
|
||||
const third = retainThoughtSignature(second, "");
|
||||
expect(third).toBe("sig-1");
|
||||
});
|
||||
|
||||
it("updates the signature when a new non-empty signature arrives", () => {
|
||||
const updated = retainThoughtSignature("sig-1", "sig-2");
|
||||
expect(updated).toBe("sig-2");
|
||||
});
|
||||
});
|
||||
@@ -1,105 +0,0 @@
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { streamGoogleGeminiCli } from "../src/providers/google-gemini-cli.js";
|
||||
import type { Context, Model, ToolCall } from "../src/types.js";
|
||||
|
||||
const emptySchema = Type.Object({});
|
||||
|
||||
const originalFetch = global.fetch;
|
||||
|
||||
afterEach(() => {
|
||||
global.fetch = originalFetch;
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe("google providers tool call missing args", () => {
|
||||
it("defaults arguments to empty object when provider omits args field", async () => {
|
||||
// Simulate a tool call response where args is missing (no-arg tool)
|
||||
const sse = `${[
|
||||
`data: ${JSON.stringify({
|
||||
response: {
|
||||
candidates: [
|
||||
{
|
||||
content: {
|
||||
role: "model",
|
||||
parts: [
|
||||
{
|
||||
functionCall: {
|
||||
name: "get_status",
|
||||
// args intentionally omitted
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
finishReason: "STOP",
|
||||
},
|
||||
],
|
||||
usageMetadata: {
|
||||
promptTokenCount: 10,
|
||||
candidatesTokenCount: 5,
|
||||
totalTokenCount: 15,
|
||||
},
|
||||
},
|
||||
})}`,
|
||||
].join("\n\n")}\n\n`;
|
||||
|
||||
const encoder = new TextEncoder();
|
||||
const dataStream = new ReadableStream<Uint8Array>({
|
||||
start(controller) {
|
||||
controller.enqueue(encoder.encode(sse));
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
|
||||
const fetchMock = vi.fn(async () => {
|
||||
return new Response(dataStream, {
|
||||
status: 200,
|
||||
headers: { "content-type": "text/event-stream" },
|
||||
});
|
||||
});
|
||||
|
||||
global.fetch = fetchMock as typeof fetch;
|
||||
|
||||
const model: Model<"google-gemini-cli"> = {
|
||||
id: "gemini-2.5-flash",
|
||||
name: "Gemini 2.5 Flash",
|
||||
api: "google-gemini-cli",
|
||||
provider: "google-gemini-cli",
|
||||
baseUrl: "https://cloudcode-pa.googleapis.com",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "Check status", timestamp: Date.now() }],
|
||||
tools: [
|
||||
{
|
||||
name: "get_status",
|
||||
description: "Get current status",
|
||||
parameters: emptySchema,
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const stream = streamGoogleGeminiCli(model, context, {
|
||||
apiKey: JSON.stringify({ token: "token", projectId: "project" }),
|
||||
});
|
||||
|
||||
for await (const _ of stream) {
|
||||
// consume stream
|
||||
}
|
||||
|
||||
const result = await stream.result();
|
||||
|
||||
expect(result.stopReason).toBe("toolUse");
|
||||
expect(result.content).toHaveLength(1);
|
||||
|
||||
const toolCall = result.content[0] as ToolCall;
|
||||
expect(toolCall.type).toBe("toolCall");
|
||||
expect(toolCall.name).toBe("get_status");
|
||||
expect(toolCall.arguments).toEqual({});
|
||||
});
|
||||
});
|
||||
@@ -1,144 +0,0 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
const googleGenAiMock = vi.hoisted(() => ({
|
||||
constructorCalls: [] as Array<Record<string, unknown>>,
|
||||
}));
|
||||
|
||||
vi.mock("@google/genai", () => {
|
||||
class GoogleGenAI {
|
||||
models = {
|
||||
generateContentStream: async function* () {
|
||||
yield {
|
||||
responseId: "vertex-response-id",
|
||||
candidates: [
|
||||
{
|
||||
content: { parts: [{ text: "ok" }] },
|
||||
finishReason: "STOP",
|
||||
},
|
||||
],
|
||||
usageMetadata: {
|
||||
promptTokenCount: 1,
|
||||
candidatesTokenCount: 1,
|
||||
totalTokenCount: 2,
|
||||
},
|
||||
};
|
||||
},
|
||||
};
|
||||
|
||||
constructor(config: Record<string, unknown>) {
|
||||
googleGenAiMock.constructorCalls.push(config);
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
GoogleGenAI,
|
||||
ThinkingLevel: {
|
||||
THINKING_LEVEL_UNSPECIFIED: "THINKING_LEVEL_UNSPECIFIED",
|
||||
MINIMAL: "MINIMAL",
|
||||
LOW: "LOW",
|
||||
MEDIUM: "MEDIUM",
|
||||
HIGH: "HIGH",
|
||||
},
|
||||
};
|
||||
});
|
||||
|
||||
import { getModel } from "../src/models.js";
|
||||
import { streamGoogleVertex } from "../src/providers/google-vertex.js";
|
||||
import type { Context } from "../src/types.js";
|
||||
|
||||
const model = getModel("google-vertex", "gemini-3-flash-preview");
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
|
||||
};
|
||||
|
||||
const originalGoogleCloudApiKey = process.env.GOOGLE_CLOUD_API_KEY;
|
||||
|
||||
beforeEach(() => {
|
||||
googleGenAiMock.constructorCalls.length = 0;
|
||||
delete process.env.GOOGLE_CLOUD_API_KEY;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
if (originalGoogleCloudApiKey === undefined) {
|
||||
delete process.env.GOOGLE_CLOUD_API_KEY;
|
||||
} else {
|
||||
process.env.GOOGLE_CLOUD_API_KEY = originalGoogleCloudApiKey;
|
||||
}
|
||||
});
|
||||
|
||||
describe("google-vertex api key resolution", () => {
|
||||
it("falls back to ADC when options.apiKey is a placeholder marker", async () => {
|
||||
const stream = streamGoogleVertex(model, context, {
|
||||
apiKey: "<authenticated>",
|
||||
project: "test-project",
|
||||
location: "us-central1",
|
||||
});
|
||||
|
||||
await stream.result();
|
||||
|
||||
expect(googleGenAiMock.constructorCalls).toHaveLength(1);
|
||||
expect(googleGenAiMock.constructorCalls[0]).toMatchObject({
|
||||
vertexai: true,
|
||||
project: "test-project",
|
||||
location: "us-central1",
|
||||
apiVersion: "v1",
|
||||
});
|
||||
expect(googleGenAiMock.constructorCalls[0]).not.toHaveProperty("apiKey");
|
||||
});
|
||||
|
||||
it("falls back to ADC when options.apiKey is the gcp-vertex-credentials marker", async () => {
|
||||
const stream = streamGoogleVertex(model, context, {
|
||||
apiKey: "gcp-vertex-credentials",
|
||||
project: "test-project",
|
||||
location: "us-central1",
|
||||
});
|
||||
|
||||
await stream.result();
|
||||
|
||||
expect(googleGenAiMock.constructorCalls).toHaveLength(1);
|
||||
expect(googleGenAiMock.constructorCalls[0]).toMatchObject({
|
||||
vertexai: true,
|
||||
project: "test-project",
|
||||
location: "us-central1",
|
||||
apiVersion: "v1",
|
||||
});
|
||||
expect(googleGenAiMock.constructorCalls[0]).not.toHaveProperty("apiKey");
|
||||
});
|
||||
|
||||
it("falls back to ADC when GOOGLE_CLOUD_API_KEY is a placeholder marker", async () => {
|
||||
process.env.GOOGLE_CLOUD_API_KEY = "<authenticated>";
|
||||
|
||||
const stream = streamGoogleVertex(model, context, {
|
||||
project: "test-project",
|
||||
location: "us-central1",
|
||||
});
|
||||
|
||||
await stream.result();
|
||||
|
||||
expect(googleGenAiMock.constructorCalls).toHaveLength(1);
|
||||
expect(googleGenAiMock.constructorCalls[0]).toMatchObject({
|
||||
vertexai: true,
|
||||
project: "test-project",
|
||||
location: "us-central1",
|
||||
apiVersion: "v1",
|
||||
});
|
||||
expect(googleGenAiMock.constructorCalls[0]).not.toHaveProperty("apiKey");
|
||||
});
|
||||
|
||||
it("still uses the API key client for real API keys", async () => {
|
||||
const stream = streamGoogleVertex(model, context, {
|
||||
apiKey: "AIzaSyExampleRealisticLookingApiKey123456",
|
||||
});
|
||||
|
||||
await stream.result();
|
||||
|
||||
expect(googleGenAiMock.constructorCalls).toHaveLength(1);
|
||||
expect(googleGenAiMock.constructorCalls[0]).toMatchObject({
|
||||
vertexai: true,
|
||||
apiKey: "AIzaSyExampleRealisticLookingApiKey123456",
|
||||
apiVersion: "v1",
|
||||
});
|
||||
expect(googleGenAiMock.constructorCalls[0]).not.toHaveProperty("project");
|
||||
expect(googleGenAiMock.constructorCalls[0]).not.toHaveProperty("location");
|
||||
});
|
||||
});
|
||||
@@ -1,481 +0,0 @@
|
||||
import { readFileSync } from "node:fs";
|
||||
import { join } from "node:path";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import type { Api, Context, Model, Tool, ToolResultMessage } from "../src/index.js";
|
||||
import { complete, getModel } from "../src/index.js";
|
||||
import type { StreamOptions } from "../src/types.js";
|
||||
|
||||
type StreamOptionsWithExtras = StreamOptions & Record<string, unknown>;
|
||||
|
||||
import { hasAzureOpenAICredentials, resolveAzureDeploymentName } from "./azure-utils.js";
|
||||
import { hasBedrockCredentials } from "./bedrock-utils.js";
|
||||
import { resolveApiKey } from "./oauth.js";
|
||||
|
||||
// Resolve OAuth tokens at module level (async, runs before tests)
|
||||
const oauthTokens = await Promise.all([
|
||||
resolveApiKey("anthropic"),
|
||||
resolveApiKey("github-copilot"),
|
||||
resolveApiKey("google-gemini-cli"),
|
||||
resolveApiKey("google-antigravity"),
|
||||
resolveApiKey("openai-codex"),
|
||||
]);
|
||||
const [anthropicOAuthToken, githubCopilotToken, geminiCliToken, antigravityToken, openaiCodexToken] = oauthTokens;
|
||||
|
||||
/**
|
||||
* Test that tool results containing only images work correctly across all providers.
|
||||
* This verifies that:
|
||||
* 1. Tool results can contain image content blocks
|
||||
* 2. Providers correctly pass images from tool results to the LLM
|
||||
* 3. The LLM can see and describe images returned by tools
|
||||
*/
|
||||
async function handleToolWithImageResult<TApi extends Api>(model: Model<TApi>, options?: StreamOptionsWithExtras) {
|
||||
// Check if the model supports images
|
||||
if (!model.input.includes("image")) {
|
||||
console.log(`Skipping tool image result test - model ${model.id} doesn't support images`);
|
||||
return;
|
||||
}
|
||||
|
||||
// Read the test image
|
||||
const imagePath = join(__dirname, "data", "red-circle.png");
|
||||
const imageBuffer = readFileSync(imagePath);
|
||||
const base64Image = imageBuffer.toString("base64");
|
||||
|
||||
// Define a tool that returns only an image (no text)
|
||||
const getImageSchema = Type.Object({});
|
||||
const getImageTool: Tool<typeof getImageSchema> = {
|
||||
name: "get_circle",
|
||||
description: "Returns a circle image for visualization",
|
||||
parameters: getImageSchema,
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant that uses tools when asked.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: "Call the get_circle tool to get an image, and describe what you see, shapes, colors, etc.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
tools: [getImageTool],
|
||||
};
|
||||
|
||||
// First request - LLM should call the tool
|
||||
const firstResponse = await complete(model, context, options);
|
||||
expect(firstResponse.stopReason).toBe("toolUse");
|
||||
|
||||
// Find the tool call
|
||||
const toolCall = firstResponse.content.find((b) => b.type === "toolCall");
|
||||
expect(toolCall).toBeTruthy();
|
||||
if (!toolCall || toolCall.type !== "toolCall") {
|
||||
throw new Error("Expected tool call");
|
||||
}
|
||||
expect(toolCall.name).toBe("get_circle");
|
||||
|
||||
// Add the tool call to context
|
||||
context.messages.push(firstResponse);
|
||||
|
||||
// Create tool result with ONLY an image (no text)
|
||||
const toolResult: ToolResultMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
content: [
|
||||
{
|
||||
type: "image",
|
||||
data: base64Image,
|
||||
mimeType: "image/png",
|
||||
},
|
||||
],
|
||||
isError: false,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
context.messages.push(toolResult);
|
||||
|
||||
// Second request - LLM should describe the image from the tool result
|
||||
const secondResponse = await complete(model, context, options);
|
||||
expect(secondResponse.stopReason).toBe("stop");
|
||||
expect(secondResponse.errorMessage).toBeFalsy();
|
||||
|
||||
// Verify the LLM can see and describe the image
|
||||
const textContent = secondResponse.content.find((b) => b.type === "text");
|
||||
expect(textContent).toBeTruthy();
|
||||
if (textContent && textContent.type === "text") {
|
||||
const lowerContent = textContent.text.toLowerCase();
|
||||
// Should mention red and circle since that's what the image shows
|
||||
expect(lowerContent).toContain("red");
|
||||
expect(lowerContent).toContain("circle");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Test that tool results containing both text and images work correctly across all providers.
|
||||
* This verifies that:
|
||||
* 1. Tool results can contain mixed content blocks (text + images)
|
||||
* 2. Providers correctly pass both text and images from tool results to the LLM
|
||||
* 3. The LLM can see both the text and images in tool results
|
||||
*/
|
||||
async function handleToolWithTextAndImageResult<TApi extends Api>(
|
||||
model: Model<TApi>,
|
||||
options?: StreamOptionsWithExtras,
|
||||
) {
|
||||
// Check if the model supports images
|
||||
if (!model.input.includes("image")) {
|
||||
console.log(`Skipping tool text+image result test - model ${model.id} doesn't support images`);
|
||||
return;
|
||||
}
|
||||
|
||||
// Read the test image
|
||||
const imagePath = join(__dirname, "data", "red-circle.png");
|
||||
const imageBuffer = readFileSync(imagePath);
|
||||
const base64Image = imageBuffer.toString("base64");
|
||||
|
||||
// Define a tool that returns both text and an image
|
||||
const getImageSchema = Type.Object({});
|
||||
const getImageTool: Tool<typeof getImageSchema> = {
|
||||
name: "get_circle_with_description",
|
||||
description: "Returns a circle image with a text description",
|
||||
parameters: getImageSchema,
|
||||
};
|
||||
|
||||
const context: Context = {
|
||||
systemPrompt: "You are a helpful assistant that uses tools when asked.",
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content:
|
||||
"Use the get_circle_with_description tool and tell me what you learned. Also say what color the shape is.",
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
tools: [getImageTool],
|
||||
};
|
||||
|
||||
// First request - LLM should call the tool
|
||||
const firstResponse = await complete(model, context, options);
|
||||
expect(firstResponse.stopReason).toBe("toolUse");
|
||||
|
||||
// Find the tool call
|
||||
const toolCall = firstResponse.content.find((b) => b.type === "toolCall");
|
||||
expect(toolCall).toBeTruthy();
|
||||
if (!toolCall || toolCall.type !== "toolCall") {
|
||||
throw new Error("Expected tool call");
|
||||
}
|
||||
expect(toolCall.name).toBe("get_circle_with_description");
|
||||
|
||||
// Add the tool call to context
|
||||
context.messages.push(firstResponse);
|
||||
|
||||
// Create tool result with BOTH text and image
|
||||
const toolResult: ToolResultMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: toolCall.id,
|
||||
toolName: toolCall.name,
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: "This is a geometric shape with specific properties: it has a diameter of 100 pixels.",
|
||||
},
|
||||
{
|
||||
type: "image",
|
||||
data: base64Image,
|
||||
mimeType: "image/png",
|
||||
},
|
||||
],
|
||||
isError: false,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
context.messages.push(toolResult);
|
||||
|
||||
// Second request - LLM should describe both the text and image from the tool result
|
||||
const secondResponse = await complete(model, context, options);
|
||||
expect(secondResponse.stopReason).toBe("stop");
|
||||
expect(secondResponse.errorMessage).toBeFalsy();
|
||||
|
||||
// Verify the LLM can see both text and image
|
||||
const textContent = secondResponse.content.find((b) => b.type === "text");
|
||||
expect(textContent).toBeTruthy();
|
||||
if (textContent && textContent.type === "text") {
|
||||
const lowerContent = textContent.text.toLowerCase();
|
||||
// Should mention details from the text (diameter/pixels)
|
||||
expect(lowerContent.match(/diameter|100|pixel/)).toBeTruthy();
|
||||
// Should also mention the visual properties (red and circle)
|
||||
expect(lowerContent).toContain("red");
|
||||
expect(lowerContent).toContain("circle");
|
||||
}
|
||||
}
|
||||
|
||||
describe("Tool Results with Images", () => {
|
||||
describe.skipIf(!process.env.GEMINI_API_KEY)("Google Provider (gemini-2.5-flash)", () => {
|
||||
const llm = getModel("google", "gemini-2.5-flash");
|
||||
|
||||
it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(llm);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Completions Provider (gpt-4o-mini)", () => {
|
||||
const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini");
|
||||
void _compat;
|
||||
const llm: Model<"openai-completions"> = {
|
||||
...baseModel,
|
||||
api: "openai-completions",
|
||||
};
|
||||
|
||||
it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(llm);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Responses Provider (gpt-5-mini)", () => {
|
||||
const llm = getModel("openai", "gpt-5-mini");
|
||||
|
||||
it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(llm);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!hasAzureOpenAICredentials())("Azure OpenAI Responses Provider (gpt-4o-mini)", () => {
|
||||
const llm = getModel("azure-openai-responses", "gpt-4o-mini");
|
||||
const azureDeploymentName = resolveAzureDeploymentName(llm.id);
|
||||
const azureOptions = azureDeploymentName ? { azureDeploymentName } : {};
|
||||
|
||||
it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(llm, azureOptions);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(llm, azureOptions);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic Provider (claude-haiku-4-5)", () => {
|
||||
const model = getModel("anthropic", "claude-haiku-4-5");
|
||||
|
||||
it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(model);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(model);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENROUTER_API_KEY)("OpenRouter Provider (glm-4.5v)", () => {
|
||||
const llm = getModel("openrouter", "z-ai/glm-4.5v");
|
||||
|
||||
it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(llm);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.MISTRAL_API_KEY)("Mistral Provider (pixtral-12b)", () => {
|
||||
const llm = getModel("mistral", "pixtral-12b");
|
||||
|
||||
it("should handle tool result with only image", { retry: 5, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(llm);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 5, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.KIMI_API_KEY)("Kimi For Coding Provider (kimi-for-coding)", () => {
|
||||
const llm = getModel("kimi-coding", "kimi-for-coding");
|
||||
|
||||
it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(llm);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.AI_GATEWAY_API_KEY)("Vercel AI Gateway Provider (google/gemini-2.5-flash)", () => {
|
||||
const llm = getModel("vercel-ai-gateway", "google/gemini-2.5-flash");
|
||||
|
||||
it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(llm);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(llm);
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!hasBedrockCredentials())("Amazon Bedrock Provider (claude-sonnet-4-5)", () => {
|
||||
const llm = getModel("amazon-bedrock", "global.anthropic.claude-sonnet-4-5-20250929-v1:0");
|
||||
|
||||
it("should handle tool result with only image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithImageResult(llm);
|
||||
});
|
||||
|
||||
it("should handle tool result with text and image", { retry: 3, timeout: 30000 }, async () => {
|
||||
await handleToolWithTextAndImageResult(llm);
|
||||
});
|
||||
});
|
||||
|
||||
// =========================================================================
|
||||
// OAuth-based providers (credentials from ~/.pi/agent/oauth.json)
|
||||
// =========================================================================
|
||||
|
||||
describe("Anthropic OAuth Provider (claude-sonnet-4-5)", () => {
|
||||
const model = getModel("anthropic", "claude-sonnet-4-5");
|
||||
|
||||
it.skipIf(!anthropicOAuthToken)(
|
||||
"should handle tool result with only image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
await handleToolWithImageResult(model, { apiKey: anthropicOAuthToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!anthropicOAuthToken)(
|
||||
"should handle tool result with text and image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
await handleToolWithTextAndImageResult(model, { apiKey: anthropicOAuthToken });
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
describe("GitHub Copilot Provider", () => {
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"gpt-4o - should handle tool result with only image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "gpt-4o");
|
||||
await handleToolWithImageResult(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"gpt-4o - should handle tool result with text and image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "gpt-4o");
|
||||
await handleToolWithTextAndImageResult(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"claude-sonnet-4 - should handle tool result with only image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "claude-sonnet-4");
|
||||
await handleToolWithImageResult(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!githubCopilotToken)(
|
||||
"claude-sonnet-4 - should handle tool result with text and image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("github-copilot", "claude-sonnet-4");
|
||||
await handleToolWithTextAndImageResult(llm, { apiKey: githubCopilotToken });
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
describe("Google Gemini CLI Provider", () => {
|
||||
it.skipIf(!geminiCliToken)(
|
||||
"gemini-2.5-flash - should handle tool result with only image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-gemini-cli", "gemini-2.5-flash");
|
||||
await handleToolWithImageResult(llm, { apiKey: geminiCliToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!geminiCliToken)(
|
||||
"gemini-2.5-flash - should handle tool result with text and image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-gemini-cli", "gemini-2.5-flash");
|
||||
await handleToolWithTextAndImageResult(llm, { apiKey: geminiCliToken });
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
describe("Google Antigravity Provider", () => {
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gemini-3-flash - should handle tool result with only image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gemini-3-flash");
|
||||
await handleToolWithImageResult(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"gemini-3-flash - should handle tool result with text and image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "gemini-3-flash");
|
||||
await handleToolWithTextAndImageResult(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
/** These two don't work, the model simply won't call the tool, works in pi
|
||||
it.skipIf(!antigravityToken)(
|
||||
"claude-sonnet-4-5 - should handle tool result with only image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "claude-sonnet-4-5");
|
||||
await handleToolWithImageResult(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!antigravityToken)(
|
||||
"claude-sonnet-4-5 - should handle tool result with text and image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("google-antigravity", "claude-sonnet-4-5");
|
||||
await handleToolWithTextAndImageResult(llm, { apiKey: antigravityToken });
|
||||
},
|
||||
);**/
|
||||
|
||||
// Note: gpt-oss-120b-medium does not support images, so not tested here
|
||||
});
|
||||
|
||||
describe("OpenAI Codex Provider", () => {
|
||||
it.skipIf(!openaiCodexToken)(
|
||||
"gpt-5.2-codex - should handle tool result with only image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("openai-codex", "gpt-5.2-codex");
|
||||
await handleToolWithImageResult(llm, { apiKey: openaiCodexToken });
|
||||
},
|
||||
);
|
||||
|
||||
it.skipIf(!openaiCodexToken)(
|
||||
"gpt-5.2-codex - should handle tool result with text and image",
|
||||
{ retry: 3, timeout: 30000 },
|
||||
async () => {
|
||||
const llm = getModel("openai-codex", "gpt-5.2-codex");
|
||||
await handleToolWithTextAndImageResult(llm, { apiKey: openaiCodexToken });
|
||||
},
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -1,145 +0,0 @@
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getEnvApiKey } from "../src/env-api-keys.js";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { completeSimple } from "../src/stream.js";
|
||||
import type { Api, Context, Model, StopReason, Tool, ToolCall, ToolResultMessage } from "../src/types.js";
|
||||
import { StringEnum } from "../src/utils/typebox-helpers.js";
|
||||
import { hasBedrockCredentials } from "./bedrock-utils.js";
|
||||
|
||||
const calculatorSchema = Type.Object({
|
||||
a: Type.Number({ description: "First number" }),
|
||||
b: Type.Number({ description: "Second number" }),
|
||||
operation: StringEnum(["add", "subtract", "multiply", "divide"], {
|
||||
description: "The operation to perform.",
|
||||
}),
|
||||
});
|
||||
|
||||
const calculatorTool: Tool<typeof calculatorSchema> = {
|
||||
name: "calculator",
|
||||
description: "Perform basic arithmetic operations",
|
||||
parameters: calculatorSchema,
|
||||
};
|
||||
|
||||
type CalculatorOperation = "add" | "subtract" | "multiply" | "divide";
|
||||
|
||||
type CalculatorArguments = {
|
||||
a: number;
|
||||
b: number;
|
||||
operation: CalculatorOperation;
|
||||
};
|
||||
|
||||
function asCalculatorArguments(args: ToolCall["arguments"]): CalculatorArguments {
|
||||
if (typeof args !== "object" || args === null) {
|
||||
throw new Error("Tool arguments must be an object");
|
||||
}
|
||||
|
||||
const value = args as Record<string, unknown>;
|
||||
const operation = value.operation;
|
||||
if (
|
||||
typeof value.a !== "number" ||
|
||||
typeof value.b !== "number" ||
|
||||
(operation !== "add" && operation !== "subtract" && operation !== "multiply" && operation !== "divide")
|
||||
) {
|
||||
throw new Error("Invalid calculator arguments");
|
||||
}
|
||||
|
||||
return { a: value.a, b: value.b, operation };
|
||||
}
|
||||
|
||||
function evaluateCalculatorCall(toolCall: ToolCall): number {
|
||||
const { a, b, operation } = asCalculatorArguments(toolCall.arguments);
|
||||
switch (operation) {
|
||||
case "add":
|
||||
return a + b;
|
||||
case "subtract":
|
||||
return a - b;
|
||||
case "multiply":
|
||||
return a * b;
|
||||
case "divide":
|
||||
return a / b;
|
||||
}
|
||||
}
|
||||
|
||||
async function assertSecondToolCallWithInterleavedThinking<TApi extends Api>(
|
||||
llm: Model<TApi>,
|
||||
reasoning: "high" | "xhigh",
|
||||
) {
|
||||
const context: Context = {
|
||||
systemPrompt: [
|
||||
"You are a helpful assistant that must use tools for arithmetic.",
|
||||
"Always think before every tool call, not just the first one.",
|
||||
"Do not answer with plain text when a tool call is required.",
|
||||
].join(" "),
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: [
|
||||
"Use calculator to calculate 328 * 29.",
|
||||
"You must call the calculator tool exactly once.",
|
||||
"Provide the final answer based on the best guess given the tool result, even if it seems unreliable.",
|
||||
"Start by thinking about the steps you will take to solve the problem.",
|
||||
].join(" "),
|
||||
timestamp: Date.now(),
|
||||
},
|
||||
],
|
||||
tools: [calculatorTool],
|
||||
};
|
||||
|
||||
const firstResponse = await completeSimple(llm, context, { reasoning });
|
||||
|
||||
expect(firstResponse.stopReason, `Error: ${firstResponse.errorMessage}`).toBe("toolUse" satisfies StopReason);
|
||||
expect(firstResponse.content.some((block) => block.type === "thinking")).toBe(true);
|
||||
expect(firstResponse.content.some((block) => block.type === "toolCall")).toBe(true);
|
||||
|
||||
const firstToolCall = firstResponse.content.find((block) => block.type === "toolCall");
|
||||
expect(firstToolCall?.type).toBe("toolCall");
|
||||
if (!firstToolCall || firstToolCall.type !== "toolCall") {
|
||||
throw new Error("Expected first response to include a tool call");
|
||||
}
|
||||
|
||||
context.messages.push(firstResponse);
|
||||
|
||||
const correctAnswer = evaluateCalculatorCall(firstToolCall);
|
||||
const firstToolResult: ToolResultMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: firstToolCall.id,
|
||||
toolName: firstToolCall.name,
|
||||
content: [{ type: "text", text: `The answer is ${correctAnswer} or ${correctAnswer * 2}.` }],
|
||||
isError: false,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
context.messages.push(firstToolResult);
|
||||
|
||||
const secondResponse = await completeSimple(llm, context, { reasoning });
|
||||
|
||||
expect(secondResponse.stopReason, `Error: ${secondResponse.errorMessage}`).toBe("stop" satisfies StopReason);
|
||||
expect(secondResponse.content.some((block) => block.type === "thinking")).toBe(true);
|
||||
expect(secondResponse.content.some((block) => block.type === "text")).toBe(true);
|
||||
}
|
||||
|
||||
const hasAnthropicCredentials = !!getEnvApiKey("anthropic");
|
||||
|
||||
describe.skipIf(!hasBedrockCredentials())("Amazon Bedrock interleaved thinking", () => {
|
||||
it("should do interleaved thinking on Claude Opus 4.5", { retry: 3 }, async () => {
|
||||
const llm = getModel("amazon-bedrock", "global.anthropic.claude-opus-4-5-20251101-v1:0");
|
||||
await assertSecondToolCallWithInterleavedThinking(llm, "high");
|
||||
});
|
||||
|
||||
it("should do interleaved thinking on Claude Opus 4.6", { retry: 3 }, async () => {
|
||||
const llm = getModel("amazon-bedrock", "global.anthropic.claude-opus-4-6-v1");
|
||||
await assertSecondToolCallWithInterleavedThinking(llm, "high");
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!hasAnthropicCredentials)("Anthropic interleaved thinking", () => {
|
||||
it("should do interleaved thinking on Claude Opus 4.5", { retry: 3 }, async () => {
|
||||
const llm = getModel("anthropic", "claude-opus-4-5");
|
||||
await assertSecondToolCallWithInterleavedThinking(llm, "high");
|
||||
});
|
||||
|
||||
it("should do interleaved thinking on Claude Opus 4.6", { retry: 3 }, async () => {
|
||||
const llm = getModel("anthropic", "claude-opus-4-6");
|
||||
await assertSecondToolCallWithInterleavedThinking(llm, "high");
|
||||
});
|
||||
});
|
||||
@@ -1,102 +0,0 @@
|
||||
import { spawnSync } from "node:child_process";
|
||||
import { createRequire } from "node:module";
|
||||
import { dirname, resolve } from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
import { describe, expect, it } from "vitest";
|
||||
|
||||
const require = createRequire(import.meta.url);
|
||||
const tsxLoader = require.resolve("tsx/esm");
|
||||
const packageRoot = resolve(dirname(fileURLToPath(import.meta.url)), "..");
|
||||
const aiEntryUrl = new URL("../src/index.ts", import.meta.url).href;
|
||||
|
||||
const SDK_SPECIFIERS = [
|
||||
"@anthropic-ai/sdk",
|
||||
"openai",
|
||||
"@google/genai",
|
||||
"@mistralai/mistralai",
|
||||
"@aws-sdk/client-bedrock-runtime",
|
||||
] as const;
|
||||
|
||||
type ProbeResult = {
|
||||
loadedSpecifiers: string[];
|
||||
};
|
||||
|
||||
function runProbe(action: string): ProbeResult {
|
||||
const script = `
|
||||
import { registerHooks } from "node:module";
|
||||
|
||||
const targets = new Set(${JSON.stringify(SDK_SPECIFIERS)});
|
||||
const loaded = [];
|
||||
|
||||
registerHooks({
|
||||
resolve(specifier, context, nextResolve) {
|
||||
if (targets.has(specifier)) {
|
||||
loaded.push(specifier);
|
||||
}
|
||||
return nextResolve(specifier, context);
|
||||
},
|
||||
});
|
||||
|
||||
const mod = await import(${JSON.stringify(aiEntryUrl)});
|
||||
${action}
|
||||
console.log(JSON.stringify({ loadedSpecifiers: [...new Set(loaded)] }));
|
||||
`;
|
||||
|
||||
const result = spawnSync(process.execPath, ["--import", tsxLoader, "--input-type=module", "--eval", script], {
|
||||
cwd: packageRoot,
|
||||
encoding: "utf8",
|
||||
});
|
||||
|
||||
if (result.status !== 0) {
|
||||
throw new Error(`Probe failed (exit ${result.status})\nSTDOUT:\n${result.stdout}\nSTDERR:\n${result.stderr}`);
|
||||
}
|
||||
|
||||
const stdoutLines = result.stdout
|
||||
.split(/\r?\n/)
|
||||
.map((line) => line.trim())
|
||||
.filter((line) => line.length > 0);
|
||||
const lastLine = stdoutLines.at(-1);
|
||||
if (!lastLine) {
|
||||
throw new Error(`Probe produced no output\nSTDERR:\n${result.stderr}`);
|
||||
}
|
||||
|
||||
return JSON.parse(lastLine) as ProbeResult;
|
||||
}
|
||||
|
||||
describe("lazy provider module loading", () => {
|
||||
it("does not load provider SDKs when importing the root barrel", () => {
|
||||
const result = runProbe("");
|
||||
expect(result.loadedSpecifiers).toEqual([]);
|
||||
});
|
||||
|
||||
it("loads only the Anthropic SDK when calling the root lazy wrapper", () => {
|
||||
const result = runProbe(`
|
||||
const model = {
|
||||
id: "claude-sonnet-4-20250514",
|
||||
name: "Claude Sonnet 4",
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
baseUrl: "https://api.anthropic.com",
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 200000,
|
||||
maxTokens: 8192,
|
||||
};
|
||||
const context = { messages: [{ role: "user", content: "hi" }] };
|
||||
await mod.streamSimpleAnthropic(model, context).result();
|
||||
`);
|
||||
|
||||
expect(result.loadedSpecifiers).toEqual(["@anthropic-ai/sdk"]);
|
||||
});
|
||||
|
||||
it("loads only the Anthropic SDK when dispatching through streamSimple", () => {
|
||||
const result = runProbe(`
|
||||
const model = mod.getModel("anthropic", "claude-sonnet-4-20250514");
|
||||
const context = { messages: [{ role: "user", content: "hi" }] };
|
||||
await mod.streamSimple(model, context).result();
|
||||
`);
|
||||
|
||||
expect(result.loadedSpecifiers).toEqual(["@anthropic-ai/sdk"]);
|
||||
});
|
||||
});
|
||||
@@ -1,66 +0,0 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { streamSimple } from "../src/stream.js";
|
||||
import type { Context, Model, SimpleStreamOptions } from "../src/types.js";
|
||||
|
||||
interface MistralPayload {
|
||||
promptMode?: "reasoning";
|
||||
reasoningEffort?: "none" | "high";
|
||||
}
|
||||
|
||||
function makeContext(): Context {
|
||||
return {
|
||||
messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
|
||||
};
|
||||
}
|
||||
|
||||
async function capturePayload(
|
||||
model: Model<"mistral-conversations">,
|
||||
options?: SimpleStreamOptions,
|
||||
): Promise<MistralPayload> {
|
||||
let capturedPayload: MistralPayload | undefined;
|
||||
const payloadCaptureModel: Model<"mistral-conversations"> = {
|
||||
...model,
|
||||
baseUrl: "http://127.0.0.1:9",
|
||||
};
|
||||
|
||||
const stream = streamSimple(payloadCaptureModel, makeContext(), {
|
||||
...options,
|
||||
apiKey: "fake-key",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload as MistralPayload;
|
||||
return payload;
|
||||
},
|
||||
});
|
||||
|
||||
await stream.result();
|
||||
|
||||
if (!capturedPayload) {
|
||||
throw new Error("Expected payload to be captured before request failure");
|
||||
}
|
||||
|
||||
return capturedPayload;
|
||||
}
|
||||
|
||||
describe("Mistral reasoning mode selection", () => {
|
||||
it("uses reasoning_effort for Mistral Small 4", async () => {
|
||||
const payload = await capturePayload(getModel("mistral", "mistral-small-2603"), { reasoning: "medium" });
|
||||
|
||||
expect(payload.reasoningEffort).toBe("high");
|
||||
expect(payload.promptMode).toBeUndefined();
|
||||
});
|
||||
|
||||
it("omits reasoning controls for Mistral Small 4 when thinking is off", async () => {
|
||||
const payload = await capturePayload(getModel("mistral", "mistral-small-2603"));
|
||||
|
||||
expect(payload.reasoningEffort).toBeUndefined();
|
||||
expect(payload.promptMode).toBeUndefined();
|
||||
});
|
||||
|
||||
it("uses prompt_mode for Magistral reasoning models", async () => {
|
||||
const payload = await capturePayload(getModel("mistral", "magistral-medium-latest"), { reasoning: "medium" });
|
||||
|
||||
expect(payload.promptMode).toBe("reasoning");
|
||||
expect(payload.reasoningEffort).toBeUndefined();
|
||||
});
|
||||
});
|
||||
@@ -1,61 +0,0 @@
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { getModel } from "../src/models.js";
|
||||
import { complete } from "../src/stream.js";
|
||||
import type { Context, Model } from "../src/types.js";
|
||||
|
||||
interface MistralToolPayload {
|
||||
tools?: Array<{
|
||||
type: "function";
|
||||
function: {
|
||||
name: string;
|
||||
parameters: Record<string, unknown>;
|
||||
};
|
||||
}>;
|
||||
}
|
||||
|
||||
describe("Mistral tool schema serialization", () => {
|
||||
it("strips TypeBox symbol keys before the SDK validates tool schemas", async () => {
|
||||
const model: Model<"mistral-conversations"> = {
|
||||
...getModel("mistral", "devstral-medium-latest"),
|
||||
baseUrl: "http://127.0.0.1:9",
|
||||
};
|
||||
const parameters = Type.Object({
|
||||
nested: Type.Object({
|
||||
value: Type.String(),
|
||||
}),
|
||||
});
|
||||
const context: Context = {
|
||||
messages: [{ role: "user", content: "Hi", timestamp: Date.now() }],
|
||||
tools: [
|
||||
{
|
||||
name: "inspect_schema",
|
||||
description: "Inspect the schema",
|
||||
parameters,
|
||||
},
|
||||
],
|
||||
};
|
||||
let capturedPayload: MistralToolPayload | undefined;
|
||||
|
||||
const response = await complete(model, context, {
|
||||
apiKey: "fake-key",
|
||||
onPayload: (payload) => {
|
||||
capturedPayload = payload as MistralToolPayload;
|
||||
return payload;
|
||||
},
|
||||
});
|
||||
|
||||
expect(capturedPayload?.tools).toHaveLength(1);
|
||||
const payloadParameters = capturedPayload?.tools?.[0]?.function.parameters;
|
||||
expect(payloadParameters).toBeDefined();
|
||||
expect(Object.getOwnPropertySymbols(payloadParameters ?? {})).toHaveLength(0);
|
||||
const properties = payloadParameters?.properties;
|
||||
expect(properties).toBeTruthy();
|
||||
expect(Object.getOwnPropertySymbols((properties as Record<string, unknown>) ?? {})).toHaveLength(0);
|
||||
const nested = (properties as Record<string, unknown> | undefined)?.nested;
|
||||
expect(nested).toBeTruthy();
|
||||
expect(Object.getOwnPropertySymbols((nested as Record<string, unknown>) ?? {})).toHaveLength(0);
|
||||
expect(response.stopReason).toBe("error");
|
||||
expect(response.errorMessage).not.toContain("Input validation failed");
|
||||
});
|
||||
});
|
||||
Loaded 100 of 684 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user