diff --git a/src/agents/pi-bundle-mcp-runtime.test.ts b/src/agents/pi-bundle-mcp-runtime.test.ts index 81576f3c6b5..2645cd667df 100644 --- a/src/agents/pi-bundle-mcp-runtime.test.ts +++ b/src/agents/pi-bundle-mcp-runtime.test.ts @@ -266,7 +266,7 @@ describe("session MCP runtime", () => { const serverScriptPath = path.join(pluginRoot, "servers", "bundle-probe.mjs"); await writeBundleProbeMcpServer(serverScriptPath, { startupCounterPath, - startupDelayMs: 1_000, + startupDelayMs: 100, pidPath, exitMarkerPath, }); diff --git a/src/agents/pi-embedded-runner-extraparams.test.ts b/src/agents/pi-embedded-runner-extraparams.test.ts index 28143c9bf02..cacd71b2ec5 100644 --- a/src/agents/pi-embedded-runner-extraparams.test.ts +++ b/src/agents/pi-embedded-runner-extraparams.test.ts @@ -1,16 +1,15 @@ import type { StreamFn } from "@mariozechner/pi-agent-core"; import type { Context, Model, SimpleStreamOptions } from "@mariozechner/pi-ai"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { - createAnthropicBetaHeadersWrapper, - createAnthropicFastModeWrapper, - createAnthropicServiceTierWrapper, - resolveAnthropicBetas, - resolveAnthropicFastMode, - resolveAnthropicServiceTier, -} from "../../test/helpers/providers/anthropic-contract.js"; import { __testing as extraParamsTesting } from "./pi-embedded-runner/extra-params.js"; +const ANTHROPIC_DEFAULT_BETAS = [ + "fine-grained-tool-streaming-2025-05-14", + "interleaved-thinking-2025-05-14", +]; +const ANTHROPIC_CONTEXT_1M_BETA = "context-1m-2025-08-07"; +const ANTHROPIC_OAUTH_BETAS = ["oauth-2025-04-20", "claude-code-20250219"]; + const XAI_FAST_MODEL_IDS = new Map([ ["grok-3", "grok-3-fast"], ["grok-3-mini", "grok-3-mini-fast"], @@ -111,6 +110,87 @@ function createTestToolStreamWrapper( }; } +function resolveAnthropicBetas( + extraParams: Record | undefined, + modelId: string, +): string[] { + const configuredBetas = Array.isArray(extraParams?.anthropicBeta) + ? extraParams.anthropicBeta.filter((value): value is string => typeof value === "string") + : []; + if (!extraParams?.context1m || !/(opus|sonnet)/i.test(modelId)) { + return configuredBetas; + } + return [...ANTHROPIC_DEFAULT_BETAS, ...configuredBetas, ANTHROPIC_CONTEXT_1M_BETA]; +} + +function resolveAnthropicServiceTier(extraParams: Record | undefined) { + const serviceTier = extraParams?.service_tier ?? extraParams?.serviceTier; + return serviceTier === "auto" || serviceTier === "standard_only" ? serviceTier : undefined; +} + +function resolveAnthropicFastMode(extraParams: Record | undefined) { + return typeof extraParams?.fastMode === "boolean" ? extraParams.fastMode : undefined; +} + +function isAnthropicOauthApiKey(apiKey: unknown): boolean { + return typeof apiKey === "string" && apiKey.startsWith("sk-ant-oat"); +} + +function isDirectAnthropicModel(model: { provider?: string; baseUrl?: string }): boolean { + const baseUrl = typeof model.baseUrl === "string" ? model.baseUrl : ""; + return model.provider === "anthropic" && (!baseUrl || baseUrl.includes("api.anthropic.com")); +} + +function createAnthropicBetaHeadersWrapper(baseStreamFn: StreamFn | undefined, betas: string[]) { + const underlying = baseStreamFn ?? (() => ({}) as ReturnType); + return ((model, context, options) => { + const nextBetas = isAnthropicOauthApiKey(options?.apiKey) + ? [...ANTHROPIC_OAUTH_BETAS, ...betas.filter((beta) => beta !== ANTHROPIC_CONTEXT_1M_BETA)] + : betas; + const existingBeta = + typeof options?.headers?.["anthropic-beta"] === "string" + ? options.headers["anthropic-beta"] + : ""; + const betaHeader = [...(existingBeta ? [existingBeta] : []), ...nextBetas].join(","); + return underlying(model, context, { + ...options, + headers: { + ...options?.headers, + ...(betaHeader ? { "anthropic-beta": betaHeader } : {}), + }, + }); + }) as StreamFn; +} + +function createAnthropicServiceTierWrapper( + baseStreamFn: StreamFn | undefined, + serviceTier: string, +) { + const underlying = baseStreamFn ?? (() => ({}) as ReturnType); + return ((model, context, options) => { + const originalOnPayload = options?.onPayload; + return underlying(model, context, { + ...options, + onPayload: (payload) => { + if ( + payload && + typeof payload === "object" && + isDirectAnthropicModel(model) && + !isAnthropicOauthApiKey(options?.apiKey) + ) { + const payloadObj = payload as Record; + payloadObj.service_tier ??= serviceTier; + } + return originalOnPayload?.(payload, model); + }, + }); + }) as StreamFn; +} + +function createAnthropicFastModeWrapper(baseStreamFn: StreamFn | undefined, fastMode: boolean) { + return createAnthropicServiceTierWrapper(baseStreamFn, fastMode ? "auto" : "standard_only"); +} + import { createAnthropicToolPayloadCompatibilityWrapper } from "./pi-embedded-runner/anthropic-family-tool-payload-compat.js"; import { createBedrockNoCacheWrapper,