@bastani/pi-ai 0.9.14 → 0.9.15

package/dist/api/anthropic-messages.js

@@ -8,8 +8,9 @@ import { getPiUserAgent } from "../utils/pi-user-agent.js";
 import { getProviderEnvValue } from "../utils/provider-env.js";
 import { retryProviderRequest } from "../utils/provider-retry.js";
 import { sanitizeSurrogates } from "../utils/sanitize-unicode.js";
+import { createStreamDeadline, withStreamDeadline } from "../utils/stream-deadline.js";
 import { getJsonSchemaToolParameters, resolveJsonSchemaStrictSampling } from "./constrained-sampling.js";
-import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.js";
+import { buildCopilotDynamicHeaders, hasCopilotVisionInput, preserveCopilotIntegrationHeader, } from "./github-copilot-headers.js";
 import { adjustMaxTokensForThinking, buildBaseOptions, clampMaxTokensToContext } from "./simple-options.js";
 import { transformMessages } from "./transform-messages.js";
 /**
@@ -110,6 +111,10 @@ function convertContentBlocks(content) {
 }
 const FINE_GRAINED_TOOL_STREAMING_BETA = "fine-grained-tool-streaming-2025-05-14";
 const INTERLEAVED_THINKING_BETA = "interleaved-thinking-2025-05-14";
+const SERVER_SIDE_FALLBACK_BETA = "server-side-fallback-2026-07-01";
+function shouldUseServerSideFallbackBeta(model) {
+    return (model.compat?.allowedFallbackModels?.length ?? 0) > 0;
+}
 function getAnthropicCompat(model) {
     return {
         supportsEagerToolInputStreaming: model.compat?.supportsEagerToolInputStreaming ?? true,
@@ -146,16 +151,8 @@ function mergeHeaders(...headerSources) {
     }
     return merged;
 }
-function mergeClientHeaders(model, ...headerSources) {
-    const merged = mergeHeaders(...headerSources);
-    if (model.provider === "kimi-coding") {
-        for (const name of Object.keys(merged)) {
-            if (name.toLowerCase() === "user-agent")
-                delete merged[name];
-        }
-        merged["User-Agent"] = getPiUserAgent();
-    }
-    return merged;
+function mergeClientHeaders(...headerSources) {
+    return mergeHeaders({ "User-Agent": getPiUserAgent() }, ...headerSources);
 }
 function hasHeader(headers, name) {
     if (!headers)
@@ -251,6 +248,10 @@ async function* iterateSseMessages(body, signal) {
     const decoder = new TextDecoder();
     const state = { event: null, data: [], raw: [] };
     let buffer = "";
+    const onAbort = () => {
+        void reader.cancel().catch(() => { });
+    };
+    signal?.addEventListener("abort", onAbort, { once: true });
     try {
         while (true) {
             if (signal?.aborted) {
@@ -293,6 +294,11 @@ async function* iterateSseMessages(body, signal) {
         }
     }
     finally {
+        signal?.removeEventListener("abort", onAbort);
+        try {
+            await reader.cancel();
+        }
+        catch { }
         reader.releaseLock();
     }
 }
@@ -348,9 +354,11 @@ export const stream = (model, context, options) => {
             stopReason: "pending",
             timestamp: Date.now(),
         };
+        const streamDeadline = createStreamDeadline(options?.streamDeadlineMs, options?.signal);
         try {
             let client;
             let isOAuth;
+            let usageModel = model;
             if (options?.client) {
                 client = options.client;
                 isOAuth = false;
@@ -361,14 +369,15 @@ export const stream = (model, context, options) => {
                 let copilotDynamicHeaders;
                 if (model.provider === "github-copilot") {
                     const hasImages = hasCopilotVisionInput(context.messages);
-                    copilotDynamicHeaders = buildCopilotDynamicHeaders({
+                    copilotDynamicHeaders = preserveCopilotIntegrationHeader(model.headers, buildCopilotDynamicHeaders({
                         messages: context.messages,
                         hasImages,
-                    });
+                        apiKey,
+                    }));
                 }
                 const cacheRetention = resolveCacheRetention(options?.cacheRetention, options?.env);
                 const cacheSessionId = cacheRetention === "none" ? undefined : options?.sessionId;
-                const created = createClient(model, apiKey, options?.interleavedThinking ?? true, shouldUseFineGrainedToolStreamingBeta(model, context), options?.headers, options?.fetch, copilotDynamicHeaders, cacheSessionId);
+                const created = createClient(model, apiKey, options?.interleavedThinking ?? true, shouldUseFineGrainedToolStreamingBeta(model, context), shouldUseServerSideFallbackBeta(model), options?.headers, options?.fetch, copilotDynamicHeaders, cacheSessionId);
                 client = created.client;
                 isOAuth = created.isOAuthToken;
             }
@@ -378,21 +387,26 @@ export const stream = (model, context, options) => {
                 params = nextParams;
             }
             const requestOptions = {
-                ...(options?.signal ? { signal: options.signal } : {}),
+                ...(streamDeadline.signal ? { signal: streamDeadline.signal } : {}),
                 ...(options?.timeoutMs !== undefined ? { timeout: options.timeoutMs } : {}),
                 maxRetries: 0,
             };
             const response = await retryProviderRequest(() => client.messages.create({ ...params, stream: true }, requestOptions).asResponse(), {
                 maxRetries: options?.maxRetries,
                 maxRetryDelayMs: options?.maxRetryDelayMs,
-                signal: options?.signal,
+                signal: streamDeadline.signal,
             });
             await options?.onResponse?.({ status: response.status, headers: headersToRecord(response.headers) }, model);
             stream.push({ type: "start", partial: output });
             const blocks = output.content;
-            for await (const event of iterateAnthropicEvents(response, options?.signal)) {
+            for await (const event of withStreamDeadline(iterateAnthropicEvents(response, streamDeadline.signal), streamDeadline.deadlineMs, streamDeadline.abort)) {
                 if (event.type === "message_start") {
                     output.responseId = event.message.id;
+                    output.model = event.message.model;
+                    const fallbackCost = output.model === model.id
+                        ? undefined
+                        : model.compat?.allowedFallbackModels?.find((fallback) => fallback.provider === model.provider && fallback.model === output.model)?.cost;
+                    usageModel = fallbackCost ? { ...model, id: output.model, cost: fallbackCost } : model;
                     // 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;
@@ -403,7 +417,7 @@ export const stream = (model, context, options) => {
                     // 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);
+                    calculateCost(usageModel, output.usage);
                 }
                 else if (event.type === "content_block_start") {
                     if (event.content_block.type === "text") {
@@ -572,7 +586,7 @@ export const stream = (model, context, options) => {
                     // 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);
+                    calculateCost(usageModel, output.usage);
                 }
             }
             if (options?.signal?.aborted) {
@@ -598,6 +612,9 @@ export const stream = (model, context, options) => {
             stream.push({ type: "error", reason: output.stopReason, error: output });
             stream.end();
         }
+        finally {
+            streamDeadline.cleanup();
+        }
     })();
     return stream;
 };
@@ -624,9 +641,15 @@ function mapThinkingLevelToEffort(model, level) {
 }
 export const streamSimple = (model, context, options) => {
     assertRequestAuth(model.provider, options?.apiKey, options?.headers);
-    const base = buildBaseOptions(model, context, options, options?.apiKey);
+    const base = {
+        ...buildBaseOptions(model, context, options, options?.apiKey),
+        toolChoice: options?.toolChoice,
+    };
     if (!options?.reasoning) {
-        return stream(model, context, { ...base, thinkingEnabled: false });
+        return stream(model, context, {
+            ...base,
+            thinkingEnabled: false,
+        });
     }
     // For models with adaptive thinking: use an effort level.
     // For older models: use budget-based thinking.
@@ -652,7 +675,7 @@ export const streamSimple = (model, context, options) => {
 function isOAuthToken(apiKey) {
     return apiKey.includes("sk-ant-oat");
 }
-function createClient(model, apiKey, interleavedThinking, useFineGrainedToolStreamingBeta, optionsHeaders, fetch, dynamicHeaders, sessionId) {
+function createClient(model, apiKey, interleavedThinking, useFineGrainedToolStreamingBeta, useServerSideFallbackBeta, optionsHeaders, fetch, dynamicHeaders, sessionId) {
     // Adaptive thinking models have interleaved thinking built in, so skip the beta header.
     const needsInterleavedBeta = interleavedThinking && model.compat?.forceAdaptiveThinking !== true;
     const betaFeatures = [];
@@ -662,6 +685,9 @@ function createClient(model, apiKey, interleavedThinking, useFineGrainedToolStre
     if (needsInterleavedBeta) {
         betaFeatures.push(INTERLEAVED_THINKING_BETA);
     }
+    if (useServerSideFallbackBeta) {
+        betaFeatures.push(SERVER_SIDE_FALLBACK_BETA);
+    }
     // Copilot: Bearer auth, selective betas.
     if (model.provider === "github-copilot") {
         const client = new Anthropic({
@@ -670,7 +696,7 @@ function createClient(model, apiKey, interleavedThinking, useFineGrainedToolStre
             baseURL: model.baseUrl,
             dangerouslyAllowBrowser: true,
             fetch,
-            defaultHeaders: mergeClientHeaders(model, {
+            defaultHeaders: mergeClientHeaders({
                 accept: "application/json",
                 "anthropic-dangerous-direct-browser-access": "true",
                 ...(betaFeatures.length > 0 ? { "anthropic-beta": betaFeatures.join(",") } : {}),
@@ -686,7 +712,7 @@ function createClient(model, apiKey, interleavedThinking, useFineGrainedToolStre
             baseURL: model.baseUrl,
             dangerouslyAllowBrowser: true,
             fetch,
-            defaultHeaders: mergeClientHeaders(model, {
+            defaultHeaders: mergeClientHeaders({
                 accept: "application/json",
                 "anthropic-dangerous-direct-browser-access": "true",
                 "anthropic-beta": ["claude-code-20250219", "oauth-2025-04-20", ...betaFeatures].join(","),
@@ -698,7 +724,7 @@ function createClient(model, apiKey, interleavedThinking, useFineGrainedToolStre
     }
     // API key or header-owned auth.
     const sessionAffinityHeaders = sessionId && getAnthropicCompat(model).sendSessionAffinityHeaders ? { "x-session-affinity": sessionId } : {};
-    const defaultHeaders = mergeClientHeaders(model, {
+    const defaultHeaders = mergeClientHeaders({
         accept: "application/json",
         "anthropic-dangerous-direct-browser-access": "true",
         ...(betaFeatures.length > 0 ? { "anthropic-beta": betaFeatures.join(",") } : {}),
@@ -813,6 +839,10 @@ function buildParams(model, context, isOAuthToken, options) {
             params.tool_choice = options.toolChoice;
         }
     }
+    const allowedFallbackModels = model.compat?.allowedFallbackModels;
+    if (allowedFallbackModels && allowedFallbackModels.length > 0) {
+        params.fallbacks = allowedFallbackModels.map((fallback) => ({ model: fallback.model }));
+    }
     return params;
 }
 // Normalize tool call IDs to match Anthropic's required pattern and length