diff --git a/changelog.d/fixes/12696-server-owned-tool-loop.md b/changelog.d/fixes/12696-server-owned-tool-loop.md new file mode 100644 index 0000000000..02d508ceee --- /dev/null +++ b/changelog.d/fixes/12696-server-owned-tool-loop.md @@ -0,0 +1 @@ +- **fix(chat):** Chat Completions no longer return empty `content` after a server-owned memory or skills tool; the first provider send and account/model recovery now share one pipeline so a follow-up round-trip can fill the reply ([#12696](https://github.com/diegosouzapw/OmniRoute/issues/12696)) — thanks @HouMinXi diff --git a/config/quality/file-size-baseline.json b/config/quality/file-size-baseline.json index de4d73ff11..769a6552f0 100644 --- a/config/quality/file-size-baseline.json +++ b/config/quality/file-size-baseline.json @@ -4,6 +4,7 @@ "_rebaseline_2026_09_04_12697_combo_pin_allowlist": "PR #12697 own growth: src/sse/handlers/chat.ts 2454->2458 (+4). checkModelAvailable preflight and handleSingleModelChat now call comboPinAllowlist so a pin-only combo step cannot scan the provider pool after 502/429. Helper lives in src/lib/combos/steps.ts under cap. Covered by tests/unit/combo-pin-implicit-allowlist.test.ts (11/11).", "_rebaseline_2026_09_05_quota_weighted": "feat/quota-weighted-routing own growth: src/app/(dashboard)/dashboard/combos/page.tsx 5066->5080 (+14 = STRATEGY_GUIDANCE_FALLBACK + STRATEGY_RECOMMENDATIONS_FALLBACK entries for quota-weighted; copy is the spec-mandated when/avoid/example and tips, irreducible at the existing fallback maps). Rebased onto 9d1a896c6 where #12671 already grew the same file 5018->5066. Covered by tests/unit/combo/quota-weighted-strategy.test.ts + autocombo-unification.test.ts.", "_rebaseline_2026_09_03_combo_execute_target_attempt": "Task 3 of handleComboChat split: new leaf open-sse/services/combo/executeTargetAttempt.ts lands at 1205 (check-file-size split-newline; wc -l 1204) above new-file cap 1200. Lift-as-is from combo.ts:1533-2616 retry loop. Pure classify predicates already extracted to executeTargetClassify.ts (54 LOC). Remaining growth is I/O + side effects (handleSingleModel, quality, pin/LKGP, cooldown, lockout) that cannot leave this file without splitting the retry loop mid-request. Frozen at exact LOC so it can only shrink. Covered by tests/unit/combo/execute-target-attempt.test.ts (7/7).", + "_rebaseline_2026_09_06_12696_chat_pipeline_retry_off": "PR #12696 own test growth: tests/integration/chat-pipeline.test.ts 1644->1648 (+4). The upstream-500 structured-error case now sets requestRetry/maxRetryIntervalSec to 0 so the new provider-execution pipeline cannot retry the mock 500 and double-count fetch. Irreducible at the existing seed+fetch mock; covered by the same test.", "_rebaseline_2026_09_03_reset_aware_model_family": "Own growth: open-sse/services/combo.ts 4036->4041 (+5). buildAutoCandidates now keys the reset-aware quota cache by getQuotaFetchScope and spreads requestedModel onto the connection so Gemini windows stay off a Claude-empty Antigravity account. Irreducible wiring at the existing fetchResetAwareQuotaWithCache call site; the family helper itself lives in antigravityQuotaFamily.ts. Covered by tests/unit/reset-aware-request-scope-12600.test.ts.", "_rebaseline_2026_09_03_overloaded_not_provider_breaker": "fix/overloaded-not-provider-breaker own growth: open-sse/services/combo.ts 4036->4075 (check-file-size split-newline, +39). Circuit-open pre-skip now records the breaker retryAfter and, when every target was skipped that way, waits the short reset via resolveCircuitOpenWaitDecision (new leaf in comboCooldownRetry.ts) instead of crystallizing ALL_TARGETS_SKIPPED in ~43ms. skippedForCircuitOpen / earliestCircuitOpenRetryMs reset each setTry so a later iteration cannot inherit a stale retryAfter. Irreducible at the existing ALL_TARGETS_SKIPPED chokepoint (same pattern as #7301/#8213 cooldown-wait). Predicate itself lives in circuitBreaker.ts / comboPredicates.ts / chatPredicates.ts, all under cap. Covered by tests/unit/overloaded-not-provider-breaker.test.ts + combo-cooldown-retry.test.ts.", "_rebaseline_2026_09_03_12649_free_tier_reaudit_gateways": "PR #12649 (fix/free-tier-quota-reaudit) own growth: src/shared/constants/providers/apikey/gateways.ts 1459->1462 (+3 = the nara authHint rewritten for the re-audited 7M/day plan now wraps to two lines, plus the Prettier reflow of two pre-existing >100-col authHint lines (oneminai, freebuff) that lint-staged enforces on any touch of the file; additive text at the existing registry chokepoint, same god-file no-split rationale as prior gateways.ts rebaselines: #11786 seekai, #10987 logfare, #10531 freebuff). Covered by tests/unit/free-tier-reaudit-2026-09.test.ts and tests/unit/free-providers-batch-2026-07.test.ts.", @@ -214,7 +215,7 @@ "_rebaseline_2026_08_24_video_bridge_fu01_fu03_fu04_result_cache_tests": "PRs #11362 (FU-01 cache hardening) + #11382 (FU-03 visual dedup policy identity) + #11383 (FU-04 focused analysis mode) own test growth: videoBridgeResultCache.test.ts <1000->1040, +40 (sum of three stacked PRs boarded together in the same merge-batch, each adding its own cache-identity assertions on the shared result-cache seam). Owner pre-authorized rebaseline for legitimate PR growth (2026-08-19 directive).", "_rebaseline_basered_codebuddy_cn": "Base-red fix (#4664 CodeBuddy CN): oauth-providers-config.test.ts 867->870 (+3) to align the EXPECTED provider list/config with the codebuddy-cn provider that #4664 added to the registry without updating this test (it asserts 'exactly once').", "_rebaseline_pr4613_compatible_provider_groups": "Reconcile #4613 already-merged growth: providers-page-utils.test.ts 1004->1052 (+48, buildCompatibleProviderGroups partition unit test). Fast-gate PR->release does not run check:file-size, so this surfaced post-merge.", - "tests/integration/chat-pipeline.test.ts": 1644, + "tests/integration/chat-pipeline.test.ts": 1648, "tests/unit/account-fallback-service.test.ts": 2056, "tests/unit/batch_api.test.ts": 1345, "tests/unit/cc-compatible-provider.test.ts": 1225, diff --git a/open-sse/handlers/chatCore.ts b/open-sse/handlers/chatCore.ts index 2c4a8babae..63395b47ad 100644 --- a/open-sse/handlers/chatCore.ts +++ b/open-sse/handlers/chatCore.ts @@ -2,10 +2,17 @@ import { extractRequestToolIdentityMap, resolveResponseToolNameMap, } from "./chatCore/requestToolIdentity.ts"; -import { injectMemoryAndSkills } from "./chatCore/memorySkillsInjection.ts"; +import { + injectMemoryAndSkills, + mergeInjectedFallbackOwnerNames, +} from "./chatCore/memorySkillsInjection.ts"; import { resolveChatCoreRequestSetup } from "./chatCore/requestSetup.ts"; import { normalizeOpenAICompatibleTools } from "./chatCore/openAICompatibleTools.ts"; -import { buildFailureUsageRecord, projectFailureUsageErrorCode } from "./chatCore/failureUsage.ts"; +import { + buildFailureUsageRecord, + projectFailureUsageErrorCode, + type FailureUsageAggregate, +} from "./chatCore/failureUsage.ts"; import { createTranslationFailureResult } from "./chatCore/translationFailure.ts"; import { estimateFinalInputTokens } from "./chatCore/contextEstimation.ts"; import { @@ -24,7 +31,6 @@ import { detectClassifierFormat, buildDefaultAllowClaudeMessage, } from "./chatCore/claudeClassifierCompat.ts"; -import { applyClientUsageBuffer } from "./chatCore/clientUsageBuffer.ts"; import { buildPostCallGuardrailContext } from "./chatCore/postCallGuardrailContext.ts"; import { storeSemanticCacheResponse } from "./chatCore/semanticCacheStore.ts"; import { buildNonStreamingResponseHeaders } from "./chatCore/nonStreamingResponseHeaders.ts"; @@ -75,9 +81,8 @@ import { getHeaderValueCaseInsensitive, isNoMemoryRequested, resolveCompressionHeader, - isStripReasoningRequested, } from "./chatCore/headers.ts"; -import { markCodexScopeRateLimited } from "./chatCore/codexFailover.ts"; + import { getCodexClientSessionId, isCodexOriginatedHeaders, @@ -104,6 +109,17 @@ import { isClaudeCodeSemanticPassthroughRequest, } from "./chatCore/passthroughHelpers.ts"; import { recoverAnthropicThinkingSignature } from "./chatCore/thinkingSignatureRecovery.ts"; +import { runProviderExecutionPipeline } from "./chatCore/providerExecutionPipeline.ts"; +import { runNonStreamingProviderLeg } from "./chatCore/nonStreamingProviderLeg.ts"; +import type { ChatCoreErrorResult } from "@/lib/skills/toolLoopTypes.ts"; +import { + applyServerOwnedToolLoopIfNeeded, + derivePostInjectionRequestIdentity, + followUpLegInput, +} from "./chatCore/serverOwnedToolLoopWire.ts"; +import { finalizeToolLoopError } from "./chatCore/nonStreamingFinalization.ts"; +import { markCodexScopeRateLimited } from "./chatCore/codexFailover.ts"; +import { deleteSessionAccountAffinity } from "@/lib/db/sessionAccountAffinity"; import { buildStreamingResponseHeaders, materializeDeduplicatedExecutionResult, @@ -125,7 +141,6 @@ export { stripStaleForwardingHeaders, }; import { resolveMemoryOwnerId, runMemoryExtractionGate } from "./chatCore/memoryExtraction.ts"; -import { CORS_HEADERS } from "../utils/cors.ts"; import { checkResourcePressureGuard } from "../utils/resourcePressure.ts"; import { normalizeHeaders } from "../utils/headers.ts"; import { resolveChatCoreRequestFormat } from "./chatCore/requestFormat.ts"; @@ -144,7 +159,6 @@ import { createSSETransformStreamWithLogger, createPassthroughStreamWithLogger, COLORS, - withBodyTimeout, } from "../utils/stream.ts"; import { ensureStreamReadiness } from "../utils/streamReadiness.ts"; import { resolveSuppressThinkClose, THINKING_MARKER_HEADER } from "../utils/thinkCloseMarker.ts"; @@ -152,14 +166,7 @@ import { resolveStreamReadinessTimeout } from "../utils/streamReadinessPolicy.ts import { resolveAgentGoalPolicy } from "../utils/agentGoalPolicy.ts"; import { createStreamController } from "../utils/streamHandler.ts"; import * as streamFailure from "../utils/streamFailureFinalization.ts"; -import { createSseHeartbeatTransform, shapeForClientFormat } from "../utils/sseHeartbeat.ts"; -import { - addBufferToUsage, - filterUsageForFormat, - estimateUsage, - normalizeUsage, - sanitizeUsagePayloadForRequest, -} from "../utils/usageTracking.ts"; +import { normalizeUsage } from "../utils/usageTracking.ts"; import { refreshWithRetry, isUnrecoverableRefreshError, @@ -210,6 +217,7 @@ import { import { areContextWindowChecksDisabled, isFeatureFlagEnabled, + isServerOwnedToolLoopEnabled, } from "@/shared/utils/featureFlags.ts"; import { resolveNoAuthEchoModel } from "./chatCore/noAuthEchoModel.ts"; import { @@ -232,16 +240,11 @@ import { detectMalformedNonStream, describeMalformedNonStream, } from "../utils/diagnostics.ts"; -import { - checkTokenLimits, - recordTokenUsage, -} from "@omniroute/open-sse/services/tokenLimitCounter.ts"; +import { checkTokenLimits } from "@omniroute/open-sse/services/tokenLimitCounter.ts"; import { COOLDOWN_MS, HTTP_STATUS, - FETCH_BODY_TIMEOUT_MS, PROVIDER_MAX_TOKENS, - STREAM_IDLE_TIMEOUT_MS, STREAM_READINESS_MAX_TIMEOUT_MS, STREAM_READINESS_TIMEOUT_MS, ANTIGRAVITY_PRE_RESPONSE_TIMEOUT_CODE, @@ -258,7 +261,6 @@ import { import { classifyProviderError, PROVIDER_ERROR_TYPES, - isEmptyContentResponse, } from "../services/errorClassifier.ts"; import { updateProviderConnection, getProviderConnectionById } from "@/lib/db/providers"; import { wasRefreshTokenRotated } from "@omniroute/open-sse/services/refreshSerializer.ts"; @@ -304,8 +306,6 @@ import { calculateCost } from "@/lib/usage/costCalculator"; import { buildClaudePassthroughToolNameMap, mergeResponseToolNameMap, - normalizeOpenAIToolFinishReasons, - restoreNonStreamingToolNames, } from "./chatCore/passthroughToolNames.ts"; import { createDisabledCompressionConfig, @@ -336,20 +336,9 @@ import { import { scheduleStreamingQuotaShareConsumption } from "./chatCore/streamingQuotaShare.ts"; import { recordStreamingUsageStats } from "./chatCore/streamingUsageStats.ts"; import { recordStreamingCost } from "./chatCore/streamingCost.ts"; -import { - appendNonStreamingSseTerminalSignal, - type NonStreamingSseTerminalState, -} from "./chatCore/nonStreamingSse.ts"; -import { - isJsonRecord, - parseNonStreamingResponseBody, -} from "./chatCore/nonStreamingResponseParse.ts"; -import { unwrapClinepassEnvelope } from "../utils/clinepassEnvelope.ts"; +import { isJsonRecord } from "./chatCore/nonStreamingResponseParse.ts"; import { recordNonStreamingUsageStats } from "./chatCore/nonStreamingUsageStats.ts"; import { - createBodyTimeoutError, - readStreamChunkWithTimeout, - computeBillableTokens, normalizeExecutorResult, executeWithUpstreamStartTimeout, resolveConnectionTimeoutMs, @@ -357,7 +346,7 @@ import { import { getModelNormalizeToolCallId, getModelPreserveOpenAIDeveloperRole } from "@/lib/db/models"; import { getProviderCredentials, extractSessionAffinityKey } from "@/sse/services/auth"; import { assertExclusiveConnectionLeaseFence } from "@/lib/db/exclusiveConnectionLeases"; -import { deleteSessionAccountAffinity } from "@/lib/db/sessionAccountAffinity"; + import { getCacheControlSettings } from "@/lib/cacheControlSettings"; import { guardrailRegistry } from "@/lib/guardrails"; import type { VideoBridgeLogRedactionEntry } from "@/lib/guardrails/videoBridge"; @@ -381,20 +370,13 @@ import { cacheReasoningFromAssistantMessage, requiresReasoningReplay, } from "../services/reasoningCache.ts"; -import { sanitizeOpenAITool } from "../services/toolSchemaSanitizer.ts"; import { isCompactResponsesEndpoint } from "../executors/codex.ts"; import { persistCodexChildQuotaResponse } from "../services/codexAccount/index.ts"; import { invalidateCodexQuotaCache } from "../services/codexQuotaFetcher.ts"; import { invalidateGenericQuotaCacheOnStatus } from "../services/genericQuotaFetcher.ts"; import { translateNonStreamingResponse } from "./responseTranslator.ts"; import { extractToolSchemaMap } from "../translator/response/openai-responses/toolSchemas.ts"; -import { unwrapClineNonStreamingEnvelope } from "./chatCore/clineResponseEnvelope.ts"; import { extractUsageFromResponse } from "./usageExtractor.ts"; -import { - sanitizeOpenAIResponse, - sanitizeResponsesApiResponse, - shouldParseTextualReasoningTags, -} from "./responseSanitizer.ts"; import { withRateLimit, updateFromHeaders, @@ -412,13 +394,6 @@ import { recordCoreOwnedAntigravityQuotaState, shouldDeferAntigravityQuotaStateToCaller, } from "../services/accountFallback.ts"; -import { - generateSignature, - getCachedResponse, - setCachedResponse, - isCacheableForRead, - isCacheableForWrite, -} from "@/lib/semanticCache"; import { saveIdempotency } from "@/lib/idempotencyLayer"; import { isModelUnavailableError, @@ -445,11 +420,7 @@ import { generateSessionId } from "../services/sessionManager.ts"; import { prepareWebSearchFallbackBody } from "../services/webSearchFallback.ts"; import { prepareWebFetchFallbackBody } from "../services/webFetchInterception.ts"; import { resolveInterceptSearch, resolveInterceptFetch } from "@/lib/db/interceptionRules"; -import { - resolveExplicitStreamAlias, - resolveStreamFlag, - stripMarkdownCodeFence, -} from "../utils/aiSdkCompat.ts"; +import { resolveExplicitStreamAlias, resolveStreamFlag } from "../utils/aiSdkCompat.ts"; import { generateRequestId } from "@/shared/utils/requestId"; import { isLocalStreamLifecycleError } from "@/shared/utils/circuitBreaker"; import { shouldIsolateProbeFailures } from "@/shared/utils/probeOrigin"; @@ -457,7 +428,6 @@ import { writeTerminalStatus } from "@/shared/utils/terminalStatus"; import { extractFacts } from "@/lib/memory/extraction"; import { handleToolCallExecution } from "@/lib/skills/interception"; import { MEMORY_BUILTIN_TOOL_NAMES } from "@/lib/skills/memoryBuiltins"; -import { OMNIROUTE_RESPONSE_HEADERS } from "@/shared/constants/headers"; import { resolveProviderId } from "@/shared/constants/providers"; import { getClaudeCodeCompatibleRequestDefaults } from "@/lib/providers/requestDefaults"; import { @@ -465,8 +435,6 @@ import { resolveClaudeCodeCompatibleSessionId, } from "../services/claudeCodeCompatible.ts"; import { setGeminiThoughtSignatureMode } from "../services/geminiThoughtSignatureStore.ts"; -import { fetchLiveProviderLimits } from "@/lib/usage/providerLimits"; -import { isClaudeExtraUsageBlockEnabled } from "@/lib/providers/claudeExtraUsage"; import { classifyModelScope429, getModelScopeRetryDelayMs, @@ -476,9 +444,7 @@ import { incrementRequestCount, incrementTokenUsage, isTpmExhausted, - isRpmExhausted, } from "../services/geminiRateLimitTracker.ts"; -import { isSmallEnoughForSemanticCache } from "../utils/estimateSize.ts"; import { getProactiveCompressionRatio } from "@/lib/db/compression"; type ChatCoreExecutorResult = ReturnType & { @@ -718,7 +684,11 @@ export async function handleChatCore({ ): EffectiveServiceTier | null => resolveReportedServiceTierFor(provider, payload, maxDepth); // Failure usage record building extracted to chatCore/failureUsage.ts (#3501); the handler keeps // the fire-and-forget save + computes latencyMs, so the call sites stay byte-identical. - const persistFailureUsage = (statusCode: number, errorCode?: string | null) => { + const persistFailureUsage = ( + statusCode: number, + errorCode?: string | null, + aggregate?: FailureUsageAggregate | null + ) => { saveRequestUsage( buildFailureUsageRecord({ provider, @@ -732,6 +702,7 @@ export async function handleChatCore({ errorCode, latencyMs: Date.now() - startTime, endpoint: endpointPath, + aggregate: aggregate ?? undefined, }) ).catch(() => {}); }; @@ -947,6 +918,33 @@ export async function handleChatCore({ // native-bypass defaults below when the operator explicitly configured it for this // provider/model pair; undefined falls through to the existing bypass logic. const interceptSearchOverride = resolveInterceptSearch(provider, effectiveModel); + + // Capture client tool names BEFORE fallback injection so the owner-provenance + // merge can distinguish tools the client already declared from synthetic tools + // added by the fallback preparer. Without this, a client function named + // `omniroute_web_search` (colliding with the fallback tool name) would be + // marked server-owned even though the client owns it. + const preConversionClientToolNames: string[] = ( + Array.isArray((body as Record).tools) + ? ((body as Record).tools as unknown[]) + : [] + ) + .map((tool) => { + if (!tool || typeof tool !== "object") return ""; + const record = tool as Record; + if (typeof record.name === "string") return record.name; + const fn = record.function; + if ( + fn && + typeof fn === "object" && + typeof (fn as Record).name === "string" + ) { + return (fn as Record).name as string; + } + return ""; + }) + .filter(Boolean); + const { body: bodyWithWebSearchFallback, fallback: webSearchFallbackPlan } = prepareWebSearchFallbackBody(body as Record, { provider, @@ -1320,6 +1318,17 @@ export async function handleChatCore({ body = injectionResult.body; const memorySettings = injectionResult.memorySettings; + // Merge web-search/web-fetch fallback tool names into the builtin owner set. + // injectMemoryAndSkills only tracks memory tools; the fallback names were + // injected into body.tools by prepareWebSearchFallbackBody/prepareWebFetchFallbackBody + // above, so they must be carried into the owner provenance chain here. + const mergedOwnerNames = mergeInjectedFallbackOwnerNames( + injectionResult, + [webSearchFallbackPlan, webFetchFallbackPlan], + preConversionClientToolNames + ); + injectionResult.builtinToolNames = mergedOwnerNames.builtinToolNames; + // Translate request (pass reqLogger for intermediate logging) // ── Proactive Context Compression (Phase 4) ── // Check if context exceeds 70% of limit and compress proactively before sending to provider. @@ -3081,27 +3090,7 @@ export async function handleChatCore({ const isModelScopeForRequest = isModelScope(); const maxAttempts = isModelScopeForRequest ? 3 : provider === "codex" ? 3 : 1; - // ── Codex 429 account-rotation state ───────────────────────────────── - // Track excluded connection IDs for codex failover across attempts. - const codexExcludedIds: string[] = []; - // Derive session affinity key once for codex failover (used to clear affinity on 429). - const codexSessionAffinityKey = - provider === "codex" - ? (extractSessionAffinityKey(body, clientRawRequest?.headers) ?? null) - : null; - - // ── Antigravity BYOP 422 account-rotation state ───────────────────── - // A GCP_PROJECT_REQUIRED 422 is account-specific (that Google - // account lacks a GCP Project ID). Rotate to a sibling antigravity - // account instead of surfacing the error, so multi-account setups - // keep working without user action. Tracked separately from - // maxAttempts so non-BYOP antigravity failures never get a second - // shot (no double upstream calls). - const antigravityByopExcludedIds: string[] = []; - let antigravityByopRotationPending = false; - - while (attempts < maxAttempts || antigravityByopRotationPending) { - antigravityByopRotationPending = false; // consumed per iteration + while (attempts < maxAttempts) { trace("pre_executor", { attempt: attempts }); updatePendingScope(pendingScope, { stage: "sending_to_provider", @@ -3286,169 +3275,25 @@ export async function handleChatCore({ } } - // Codex 429 account-rotation failover (disabled for context-relay so combo.ts can inject handoff) - if ( - provider === "codex" && - !managedLease && - comboStrategy !== "context-relay" && - res.response.status === 429 && - attempts < maxAttempts - 1 && - // Probe-origin (test-all) 429 must not rotate accounts or persist - // cooldowns — routing state untouched (#9817). - !(await shouldIsolateProbeFailures()) - ) { - const failedConnectionId = - executionConnectionId || credentials?.connectionId || connectionId; - const normalizedHeaders = normalizeHeaders(res.response.headers); - const retryAfterHeader = normalizedHeaders["retry-after"] ?? null; - const retryAfterMs = retryAfterHeader - ? Number.parseFloat(retryAfterHeader) * 1000 - : null; - - log?.warn?.( - "CODEX_FAILOVER", - `429 on connection ${String(failedConnectionId).slice(0, 8)} (attempt ${attempts + 1}/${maxAttempts}), rotating account` - ); - - // Mark only the current Codex model scope as rate-limited. A connection-wide - // cooldown here would let a Spark limit suppress independent Sol/Terra traffic. - if (failedConnectionId) { - await markCodexScopeRateLimited({ - failedConnectionId: String(failedConnectionId), - model: modelToCall || model || requestedModel || null, - rateLimitedUntil: new Date(Date.now() + (retryAfterMs || 60_000)).toISOString(), - credentials: execCreds || credentials, - }); - if (!codexExcludedIds.includes(String(failedConnectionId))) { - codexExcludedIds.push(String(failedConnectionId)); - } - } - - // Clear session affinity so next request won't be pinned to the failing account - if (codexSessionAffinityKey) { - try { - deleteSessionAccountAffinity(codexSessionAffinityKey, "codex"); - } catch { - // best-effort - } - } - - // Fetch next available codex connection (excluding all previously failed ones) - const nextCreds = await getProviderCredentials( - "codex", - null, - null, - modelToCall || model || requestedModel || null, - { - excludeConnectionIds: [...codexExcludedIds], - } - ).catch(() => null); - - if (!nextCreds || nextCreds.allRateLimited) { - log?.warn?.("CODEX_FAILOVER", "No more codex accounts available — returning 429"); - if (stream) { - releaseAccountSemaphore(); - return { - ...res, - _executionCredentials: execCreds, - }; - } - return { - ...res, - _accountSemaphoreRelease: releaseAccountSemaphore, - _executionCredentials: execCreds, - }; - } - - const newConnectionId = nextCreds.connectionId; - log?.info?.( - "CODEX_FAILOVER", - `Rotating codex account: ${String(failedConnectionId).slice(0, 8)} → ${newConnectionId.slice(0, 8)} (attempt ${attempts + 2}/${maxAttempts})` - ); - - logAuditEvent({ - action: "codex.account_rotation", - actor: apiKeyInfo?.name || "system", - target: newConnectionId, - details: { - failed_connection_id: failedConnectionId, - new_connection_id: newConnectionId, - attempt: attempts + 1, - retry_after_ms: retryAfterMs, - }, - }); - - // Update credentials in-place so getExecutionCredentials() picks up the new account - Object.assign(credentials, nextCreds); - - releaseAccountSemaphore(); - attempts++; - continue; - } - - // ── Antigravity BYOP 422 account rotation ─────────────────────── - // GCP_PROJECT_REQUIRED (422, code gcp_project_required) means - // THIS Google account must Bring Its Own GCP Project. Mark the - // connection excluded (rateLimitedUntil, best-effort) and rotate - // to a sibling antigravity account so the request succeeds - // without user action. When no sibling exists (or all are BYOP), - // fall through: the error-state block excludes the connection - // and the actionable 422 is surfaced. - if (provider === "antigravity" && res.response.status === 422) { - const byopBody = await res.response - .clone() - .text() - .catch(() => ""); - if (byopBody.includes("gcp_project_required")) { - const byopFailedId = - executionConnectionId || credentials?.connectionId || connectionId; - if (byopFailedId) { - if (!antigravityByopExcludedIds.includes(String(byopFailedId))) { - antigravityByopExcludedIds.push(String(byopFailedId)); - } - try { - const { setConnectionRateLimitUntil } = await import("@/lib/db/providers"); - setConnectionRateLimitUntil( - String(byopFailedId), - Date.now() + COOLDOWN_MS.gcpProjectRequired - ); - } catch { - // best-effort — never break the rotation path - } - } - const byopNextCreds = await getProviderCredentials( - "antigravity", - null, - null, - modelToCall || model || requestedModel || null, - { excludeConnectionIds: [...antigravityByopExcludedIds] } - ).catch(() => null); - if (byopNextCreds && !byopNextCreds.allRateLimited) { - log?.warn?.( - "ANTIGRAVITY_BYOP_ROTATION", - `BYOP 422 on connection ${String(byopFailedId).slice(0, 8)} → rotating to ${String(byopNextCreds.connectionId).slice(0, 8)}` - ); - releaseAccountSemaphore(); - Object.assign(credentials, byopNextCreds); - antigravityByopRotationPending = true; - continue; - } - } - } - // For streaming: release the semaphore when the client drains or cancels the stream. + // Non-2xx streams must drop the slot before returning so the pipeline can rotate + // accounts without holding the failed connection's concurrency gate. Do NOT + // cancel() the body here — the pipeline clones it (BYOP 422 / toOutcome). if (stream) { const originalBody = res.response.body; - if (!originalBody) { + const okStatus = res.response.status >= 200 && res.response.status < 300; + if (!originalBody || !okStatus) { releaseAccountSemaphore(); - return res; + return { + ...res, + _executionCredentials: execCreds, + }; } // Opt-in transparent stream recovery (free-claude-code port, default OFF). // Only engages for a successful (2xx) stream — an error body must never be // held or replayed. Setting is read once here from the cached resolved // resilience settings; the default path is byte-for-byte unchanged. - const okStatus = res.response.status >= 200 && res.response.status < 300; let streamRecoveryEnabled = false; let continueMidStreamEnabled = false; let throughputWatchdog = @@ -3779,13 +3624,124 @@ export async function handleChatCore({ let finalBody; let claudePromptCacheLogMeta = null; + let pipelineRecovered = false; + if (stream) { try { - const result = await executeProviderRequest(effectiveModel, true); + const pipelineOutcome = await runProviderExecutionPipeline({ + policy: { + allowAccountRotation: !managedLease && comboStrategy !== "context-relay", + allowModelFallback: true, + expectedConnectionId: managedLease + ? String(getCurrentConnectionId() || connectionId || "") || undefined + : undefined, + }, + target: { + provider, + requestedModel: effectiveModel, + sourceFormat, + targetFormat, + stream, + }, + connection: { + initialConnectionId: String(getCurrentConnectionId() || connectionId || ""), + getCurrentConnectionId: () => getCurrentConnectionId() || undefined, + getCredentials: () => (credentials || {}) as Record, + replaceCredentials: (next) => { + Object.assign(credentials, next); + }, + onCredentialsRefreshed: async () => {}, + assertManagedLeaseFence: (id) => { + assertManagedLeaseFence(id); + }, + getProviderCredentials, + }, + wire: { + body: translatedBody as Record, + currentModel, + triedModels, + setBodyAndModel: (body, model) => { + translatedBody = body as typeof translatedBody; + currentModel = model; + triedModels.add(model); + }, + }, + state: { + updatePendingStage: (stage, data) => { + updatePendingScope(pendingScope, { stage, ...(data || {}) }); + }, + recordRateLimitHeaders: updateFromHeaders, + recordRateLimitBody: updateFromResponseBody, + writeTerminalStatus, + persistConnectionPatch: updateProviderConnection, + setConnectionRateLimitedUntil: async (id, untilMs) => { + const { setConnectionRateLimitUntil } = await import("@/lib/db/providers"); + setConnectionRateLimitUntil(id, untilMs); + }, + lockModel, + recordAntigravityQuotaState: recordCoreOwnedAntigravityQuotaState, + markAccountSemaphoreBlocked: (key) => { + markAccountSemaphoreBlocked(key, Date.now() + 60_000); + }, + isolateProbeFailures: () => shouldIsolateProbeFailures(), + onCodexScopeRateLimited: async (params) => { + await markCodexScopeRateLimited({ + failedConnectionId: params.failedConnectionId, + model: params.model, + rateLimitedUntil: params.rateLimitedUntil, + credentials: (params.credentials || credentials) as { + connectionId?: string | null; + providerSpecificData?: unknown; + }, + }); + }, + onClearSessionAffinity: () => { + const key = + sessionAffinityKey || + extractSessionAffinityKey(body, clientRawRequest?.headers) || + null; + if (!key) return; + try { + deleteSessionAccountAffinity(key, "codex"); + } catch { + // best-effort + } + }, + onAuditAccountRotation: (params) => { + logAuditEvent({ + action: params.action, + actor: apiKeyInfo?.name || "system", + target: params.newConnectionId, + details: { + failed_connection_id: params.failedConnectionId, + new_connection_id: params.newConnectionId, + attempt: params.attempt, + retry_after_ms: params.retryAfterMs, + }, + }); + }, + }, + sendProviderAttempt: (modelToCall, allowDedup) => executeProviderRequest(modelToCall, allowDedup), + }); - providerResponse = result.response; - providerUrl = result.url; - providerHeaders = result.headers; - finalBody = providerRequestCapture.body(result.transformedBody); + pipelineRecovered = true; + currentModel = pipelineOutcome.model; + if (pipelineOutcome.kind === "error") { + providerResponse = pipelineOutcome.result.response; + providerUrl = ""; + providerHeaders = normalizeHeaders(pipelineOutcome.result.response.headers); + finalBody = translatedBody; + } else { + const result = { + response: pipelineOutcome.response, + url: pipelineOutcome.url, + headers: pipelineOutcome.headers, + transformedBody: pipelineOutcome.transformedBody, + }; + providerResponse = result.response; + providerUrl = result.url; + providerHeaders = result.headers; + finalBody = providerRequestCapture.body(result.transformedBody); + } const responseConnectionId = getCurrentConnectionId(); effectiveServiceTier = resolveEffectiveServiceTier(finalBody); claudePromptCacheLogMeta = buildClaudePromptCacheLogMeta( @@ -4168,7 +4124,9 @@ export async function handleChatCore({ ); } - const signatureRecovery = await recoverAnthropicThinkingSignature({ + const signatureRecovery = pipelineRecovered + ? { attempted: false, succeeded: false, execution: null, error: null, recoveryBody: null } + : await recoverAnthropicThinkingSignature({ provider, statusCode, message, @@ -4179,7 +4137,7 @@ export async function handleChatCore({ }, parseError: (response) => parseUpstreamError(response, provider), }); - if (signatureRecovery.attempted && signatureRecovery.execution) { + if (!pipelineRecovered && signatureRecovery.attempted && signatureRecovery.execution) { providerResponse = signatureRecovery.execution.response; if (signatureRecovery.succeeded) { providerUrl = signatureRecovery.execution.url; @@ -4591,7 +4549,7 @@ export async function handleChatCore({ // Before returning a model-unavailable error upstream, try sibling models // from the same family. This keeps the request alive on the same account // instead of failing the entire combo. - if (isModelUnavailableError(statusCode, message, provider)) { + if (!pipelineRecovered && isModelUnavailableError(statusCode, message, provider)) { const nextModel = getNextFamilyFallback(currentModel, triedModels, provider); if (nextModel) { triedModels.add(nextModel); @@ -4795,220 +4753,319 @@ export async function handleChatCore({ } // ── End T5 ─────────────────────────────────────────────────────────────── } + } // Non-streaming response if (!stream) { - const parsed = await parseNonStreamingResponseBody({ - providerResponse, - upstreamStream, - providerHeaders, - finalBody, + try { + const runNonStreamingPipeline = async ({ policy, model: pipelineModel, translatedBody: wireBody }) => { + translatedBody = wireBody as typeof translatedBody; + currentModel = pipelineModel; + triedModels.add(pipelineModel); + return runProviderExecutionPipeline({ + policy, + target: { + provider, + requestedModel: pipelineModel, + sourceFormat, + targetFormat, + stream: false, + }, + connection: { + initialConnectionId: String(getCurrentConnectionId() || connectionId || ""), + getCurrentConnectionId: () => getCurrentConnectionId() || undefined, + getCredentials: () => (credentials || {}) as Record, + replaceCredentials: (next) => { + Object.assign(credentials, next); + }, + onCredentialsRefreshed: async () => {}, + assertManagedLeaseFence: (id) => { + assertManagedLeaseFence(id); + }, + getProviderCredentials, + }, + wire: { + body: translatedBody as Record, + currentModel, + triedModels, + setBodyAndModel: (nextBody, nextModel) => { + translatedBody = nextBody as typeof translatedBody; + currentModel = nextModel; + triedModels.add(nextModel); + }, + }, + state: { + updatePendingStage: (stage, data) => { + updatePendingScope(pendingScope, { stage, ...(data || {}) }); + }, + recordRateLimitHeaders: updateFromHeaders, + recordRateLimitBody: updateFromResponseBody, + writeTerminalStatus, + persistConnectionPatch: updateProviderConnection, + setConnectionRateLimitedUntil: async (id, untilMs) => { + const { setConnectionRateLimitUntil } = await import("@/lib/db/providers"); + setConnectionRateLimitUntil(id, untilMs); + }, + lockModel, + recordAntigravityQuotaState: recordCoreOwnedAntigravityQuotaState, + markAccountSemaphoreBlocked: (key) => { + markAccountSemaphoreBlocked(key, Date.now() + 60_000); + }, + isolateProbeFailures: () => shouldIsolateProbeFailures(), + onCodexScopeRateLimited: async (params) => { + await markCodexScopeRateLimited({ + failedConnectionId: params.failedConnectionId, + model: params.model, + rateLimitedUntil: params.rateLimitedUntil, + credentials: (params.credentials || credentials) as { + connectionId?: string | null; + providerSpecificData?: unknown; + }, + }); + }, + onClearSessionAffinity: () => { + const key = + sessionAffinityKey || + extractSessionAffinityKey(body, clientRawRequest?.headers) || + null; + if (!key) return; + try { + deleteSessionAccountAffinity(key, "codex"); + } catch { + // best-effort + } + }, + onAuditAccountRotation: (params) => { + logAuditEvent({ + action: params.action, + actor: apiKeyInfo?.name || "system", + target: params.newConnectionId, + details: { + failed_connection_id: params.failedConnectionId, + new_connection_id: params.newConnectionId, + attempt: params.attempt, + retry_after_ms: params.retryAfterMs, + }, + }); + }, + }, + sendProviderAttempt: (modelToCall, allowDedup) => + executeProviderRequest(modelToCall, allowDedup), + }); + }; + + let toolLoopRan = false; + let toolLoopUsage = null; + let legResult = await runNonStreamingProviderLeg({ + phase: "initial", + sourceBody: (body || {}) as Record, + expectedConnectionId: managedLease + ? String(getCurrentConnectionId() || connectionId || "") || undefined + : undefined, + allowAccountRotation: !managedLease && comboStrategy !== "context-relay", + allowModelFallback: true, + executeProviderRequest: (modelToCall, allowDedup) => + executeProviderRequest(modelToCall, allowDedup), + runProviderExecution: runNonStreamingPipeline, + setRequestWireState: ({ translatedBody: nextBody, effectiveModel: nextModel }) => { + translatedBody = nextBody as typeof translatedBody; + currentModel = nextModel; + triedModels.add(nextModel); + }, + sourceFormat, targetFormat, - model, + clientResponseFormat, + provider, + model: effectiveModel, + connectionId: String(getCurrentConnectionId() || connectionId || ""), + getCurrentConnectionId: () => getCurrentConnectionId() || undefined, + effectiveModel: currentModel, + translatedBody: translatedBody as Record, + toolNameMap, + requestToolIdentityMap, + reasoningCacheScope, + clientHeaders: clientRawRequest?.headers ?? null, + isClaudeCodeCompatible, log, }); - const normalizedProviderPayload = parsed.normalizedProviderPayload; - const looksLikeSSE = parsed.looksLikeSSE; - if (parsed.kind === "invalid_sse") { - appendRequestLog({ - model, - provider, - connectionId, - status: `FAILED ${HTTP_STATUS.BAD_GATEWAY}`, - }).catch(() => {}); - const invalidSseMessage = parsed.message; - persistAttemptLogs({ - status: HTTP_STATUS.BAD_GATEWAY, - error: invalidSseMessage, - providerRequest: finalBody || translatedBody, - providerResponse: normalizedProviderPayload, - clientResponse: buildErrorBody(HTTP_STATUS.BAD_GATEWAY, invalidSseMessage), - cacheSource: "upstream", - }); - persistFailureUsage(HTTP_STATUS.BAD_GATEWAY, "invalid_sse_payload"); - trackPendingRequest(model, provider, pendingConnId, false); - return createErrorResult(HTTP_STATUS.BAD_GATEWAY, invalidSseMessage); - } - - if (parsed.kind === "invalid_json") { - appendRequestLog({ - model, - provider, - connectionId, - status: `FAILED ${HTTP_STATUS.BAD_GATEWAY}`, - }).catch(() => {}); - const detailedError = parsed.detailedError; - const invalidJsonMessage = parsed.message; - persistAttemptLogs({ - status: HTTP_STATUS.BAD_GATEWAY, - error: detailedError, - providerRequest: finalBody || translatedBody, - providerResponse: normalizedProviderPayload, - clientResponse: buildErrorBody(HTTP_STATUS.BAD_GATEWAY, invalidJsonMessage), - cacheSource: "upstream", - }); - persistFailureUsage(HTTP_STATUS.BAD_GATEWAY, "invalid_json_payload"); - trackPendingRequest(model, provider, connectionId, false); - return createErrorResult(HTTP_STATUS.BAD_GATEWAY, invalidJsonMessage); - } - - let responseBody = parsed.responseBody; - let responsePayloadFormat = parsed.responsePayloadFormat; - - // ── ClinePass {success,data} envelope unwrap (before translation) ────────── - // ClinePass wraps non-streaming JSON in a {success, data} envelope; errors - // use {success:false, error}. Transient {success:false, error:"empty..."} - // responses get one 2s retry before surfacing. CLINEPASS-GATED — untouched - // for every other provider. Envelope errors route through createErrorResult - // (→ buildErrorBody/sanitizeErrorMessage, Rule #12). - if (provider === "clinepass") { - let { body: unwrapped, error: envError } = unwrapClinepassEnvelope(responseBody, provider); - if (envError && /empty/i.test(envError.message || "")) { - log?.warn?.("RETRY", "clinepass returned empty content, retrying once after 2s"); - await new Promise((r) => setTimeout(r, 2000)); - try { - const retryResult = await executeProviderRequest(effectiveModel, false); - if (retryResult?.response?.ok) { - const retryParsed = await parseNonStreamingResponseBody({ - providerResponse: retryResult.response, - upstreamStream: undefined, - providerHeaders: retryResult.headers, - finalBody: retryResult.transformedBody, - targetFormat, - model, - log, - }); - if (retryParsed.kind !== "invalid_sse" && retryParsed.kind !== "invalid_json") { - providerResponse = retryResult.response; - providerUrl = retryResult.url; - providerHeaders = retryResult.headers; - finalBody = providerRequestCapture.body(retryResult.transformedBody); - ({ body: unwrapped, error: envError } = unwrapClinepassEnvelope( - retryParsed.responseBody, - provider - )); - } - } - } catch (retryErr) { - log?.warn?.( - "RETRY", - `clinepass retry failed: ${ - retryErr instanceof Error ? retryErr.message : String(retryErr) - }` - ); - } + if (legResult.kind === "error") { + const err = legResult.result; + const captured = providerRequestCapture.latest?.() ?? null; + finalBody = captured?.body ?? finalBody ?? translatedBody; + if (captured) { + reqLogger.logTargetRequest(captured.url, captured.headers, captured.body); } - if (envError) { - appendRequestLog({ - model, - provider, - connectionId, - status: `FAILED ${HTTP_STATUS.BAD_GATEWAY}`, - }).catch(() => {}); - persistFailureUsage(HTTP_STATUS.BAD_GATEWAY, "clinepass_envelope_error"); - trackPendingRequest(model, provider, connectionId, false); - return createErrorResult(HTTP_STATUS.BAD_GATEWAY, envError.message); - } - if (!isJsonRecord(unwrapped)) { - const invalidEnvelopeMessage = "Invalid JSON response from provider"; - persistFailureUsage(HTTP_STATUS.BAD_GATEWAY, "clinepass_envelope_error"); - trackPendingRequest(model, provider, connectionId, false); - return createErrorResult(HTTP_STATUS.BAD_GATEWAY, invalidEnvelopeMessage); - } - responseBody = unwrapped; - } - responseBody = unwrapClineNonStreamingEnvelope(provider, responseBody); - - // Check for empty content response (fake success) - trigger fallback - if (isEmptyContentResponse(responseBody)) { - appendRequestLog({ - model, - provider, - connectionId, - status: `FAILED ${HTTP_STATUS.BAD_GATEWAY}`, - }).catch(() => {}); - const emptyContentMessage = "Provider returned empty content"; - persistAttemptLogs({ - status: HTTP_STATUS.BAD_GATEWAY, - error: emptyContentMessage, - providerRequest: finalBody || translatedBody, - providerResponse: normalizedProviderPayload, - clientResponse: buildErrorBody(HTTP_STATUS.BAD_GATEWAY, emptyContentMessage), - cacheSource: "upstream", - }); - persistFailureUsage(HTTP_STATUS.BAD_GATEWAY, "empty_content"); - - // Trigger non-recursive fallback for empty content - const nextModel = getNextFamilyFallback(currentModel, triedModels, provider); - if (nextModel) { - triedModels.add(nextModel); - currentModel = nextModel; - translatedBody.model = nextModel; - log?.info?.( - "EMPTY_CONTENT_FALLBACK", - `${model} returned empty content → trying ${nextModel}` + reqLogger.logError(new Error(err.error || "Provider request failed"), finalBody); + const isNetworkThrow = Boolean(err.originalError); + if (err.response && !isNetworkThrow) { + reqLogger.logProviderResponse( + err.status, + err.response.statusText || "Error", + err.response.headers, + err.response ); - try { - const fallbackResult = await executeProviderRequest(nextModel, false); - if (fallbackResult.response.ok) { - const fallbackRaw = await withBodyTimeout(fallbackResult.response.text()); - try { - responseBody = fallbackRaw ? JSON.parse(fallbackRaw) : {}; - providerUrl = fallbackResult.url; - providerHeaders = fallbackResult.headers; - finalBody = providerRequestCapture.body(fallbackResult.transformedBody); - reqLogger.logTargetRequest(providerUrl, providerHeaders, finalBody); - log?.info?.( - "EMPTY_CONTENT_FALLBACK", - `Serving ${nextModel} as fallback for ${model}` - ); - // Fall through — continue processing with the new responseBody - } catch { - trackPendingRequest(model, provider, connectionId, false); - return createErrorResult(HTTP_STATUS.BAD_GATEWAY, emptyContentMessage); - } - } else { - trackPendingRequest(model, provider, connectionId, false); - return createErrorResult(HTTP_STATUS.BAD_GATEWAY, emptyContentMessage); - } - } catch { - trackPendingRequest(model, provider, connectionId, false); - return createErrorResult(HTTP_STATUS.BAD_GATEWAY, emptyContentMessage); - } - } else { - trackPendingRequest(model, provider, connectionId, false); - return createErrorResult(HTTP_STATUS.BAD_GATEWAY, emptyContentMessage); } + appendRequestLog({ + model, + provider, + connectionId, + status: `FAILED ${err.status}`, + }).catch(() => {}); + persistAttemptLogs({ + status: err.status, + error: err.error || "Provider request failed", + providerRequest: finalBody || translatedBody, + providerResponse: isNetworkThrow ? undefined : err.response, + clientResponse: buildErrorBody(err.status, err.error || "Provider request failed"), + cacheSource: "upstream", + }); + persistFailureUsage( + err.status, + err.errorCode || `upstream_${err.status}` + ); + trackPendingRequest(model, provider, connectionId, false); + return err; } - const restoreClaudeNames = sourceFormat === FORMATS.CLAUDE && targetFormat === FORMATS.CLAUDE; - let responseToolNameMap: Map | null; - [responseBody, responseToolNameMap] = restoreNonStreamingToolNames( - responseBody, - toolNameMap, - finalBody, - restoreClaudeNames + pipelineRecovered = true; + const expectedConn = managedLease + ? String(getCurrentConnectionId() || connectionId || "") || undefined + : undefined; + const loopApply = await applyServerOwnedToolLoopIfNeeded({ + enabled: isServerOwnedToolLoopEnabled(), + stream, + isResponsesEndpoint, + sourceFormat, + initialLeg: legResult, + sourceBody: (body || {}) as Record, + skillsModelId: getSkillsModelIdForFormat(sourceFormat), + executionContext: { + apiKeyId: memoryOwnerId || "local", + sessionId: pipelineSessionId, + requestId: skillRequestId, + requestIdentity: derivePostInjectionRequestIdentity({ + apiKeyId: memoryOwnerId || "local", + headers: clientRawRequest?.headers ?? null, + skillRequestId, + postInjectionBody: (body || {}) as Record, + }), + builtinToolNames: injectionResult.builtinToolNames, + injectedCustomSkillNames: injectionResult.injectedCustomSkillNames, + customSkillExecutionEnabled: + Boolean(memoryOwnerId) && memorySettings?.skillsEnabled === true, + executionFenceEnabled: true, + provider, + model: effectiveModel, + }, + abortSignal: clientRawRequest?.signal, + expectedConnectionId: expectedConn, + followUpLeg: async (nextSourceBody) => { + translatedBody = translateRequest( + sourceFormat, + targetFormat, + model, + { ...nextSourceBody }, + false, + credentials, + provider, + reqLogger, + { + normalizeToolCallId: getModelNormalizeToolCallId(provider || "", model || "", sourceFormat), + preserveDeveloperRole: getModelPreserveOpenAIDeveloperRole( + provider || "", + model || "", + sourceFormat + ), + preserveCacheControl, + signatureNamespace: connectionId, + copilotClient: copilotCompatibleReasoning, + reasoningCacheScope, + } + ); + return runNonStreamingProviderLeg( + followUpLegInput( + { + executeProviderRequest: (modelToCall, allowDedup) => + executeProviderRequest(modelToCall, allowDedup), + runProviderExecution: runNonStreamingPipeline, + setRequestWireState: ({ translatedBody: nextBody, effectiveModel: nextModel }) => { + translatedBody = nextBody as typeof translatedBody; + currentModel = nextModel; + triedModels.add(nextModel); + }, + sourceFormat, + targetFormat, + clientResponseFormat, + provider, + model: effectiveModel, + connectionId: String(getCurrentConnectionId() || connectionId || ""), + getCurrentConnectionId: () => getCurrentConnectionId() || undefined, + effectiveModel: currentModel, + translatedBody: translatedBody as Record, + toolNameMap, + requestToolIdentityMap, + reasoningCacheScope, + clientHeaders: clientRawRequest?.headers ?? null, + isClaudeCodeCompatible, + log, + }, + nextSourceBody, + expectedConn + ) + ); + }, + logReceipt: (receipt) => reqLogger.logToolLoopReceipt(receipt), + }); + if (loopApply.kind === "error") { + return await finalizeToolLoopError({ + loop: loopApply.loop, + model, + provider, + connectionId, + providerRequest: loopApply.loop.finalProviderRequest || finalBody || translatedBody, + persistFailureUsage, + persistAttemptLogs, + trackPendingRequest, + }); + } + if (loopApply.kind === "ok") { + toolLoopRan = true; + toolLoopUsage = loopApply.usage; + legResult = loopApply.leg; + } + + if (legResult.upstreamResponse) { + providerResponse = legResult.upstreamResponse; + providerHeaders = normalizeHeaders(legResult.upstreamResponse.headers); + } else { + providerResponse = new Response(null, { + status: 200, + headers: legResult.headers, + }); + providerHeaders = normalizeHeaders(legResult.headers); + } + finalBody = providerRequestCapture.body(legResult.providerRequest || translatedBody); + const capturedOk = providerRequestCapture.latest?.(); + reqLogger.logTargetRequest( + legResult.requestUrl || capturedOk?.url || "", + legResult.requestHeaders || capturedOk?.headers || {}, + capturedOk?.body ?? finalBody ); + const responseBody = legResult.providerBody; + const responsePayloadFormat = legResult.responsePayloadFormat; + const looksLikeSSE = legResult.looksLikeSSE; + let translatedResponse = legResult.response; + const memoryExtractionResponse = legResult.responseForMemoryExtraction; reqLogger.logProviderResponse( - providerResponse.status, - providerResponse.statusText, + 200, + "OK", providerResponse.headers, looksLikeSSE - ? { - _streamed: true, - _format: "sse-json", - summary: responseBody, - } + ? { _streamed: true, _format: "sse-json", summary: responseBody } : responseBody ); - sanitizeUsagePayloadForRequest( - responseBody, - finalBody || translatedBody || body, - responsePayloadFormat - ); effectiveServiceTier = resolveReportedServiceTier(responseBody) ?? effectiveServiceTier; - // Notify success - caller can clear error status if needed if (onRequestSuccess) { await onRequestSuccess(); } @@ -5019,12 +5076,10 @@ export async function handleChatCore({ providerSpecificData: credentials?.providerSpecificData, log, }); - - // Log usage for non-streaming responses - const usage = extractUsageFromResponse(responseBody, provider); + const usage = toolLoopUsage ?? extractUsageFromResponse(responseBody, provider); + const cacheUsageLogMeta = buildCacheUsageLogMeta(usage); if (usage && typeof usage === "object") { attachCompressionUsageReceiptAfterAnalytics(usage as Record, "provider"); - // Track Gemini token consumption for TPM rate-limit pre-check if (provider === "gemini") { const promptTokens = typeof (usage as Record).prompt_tokens === "number" @@ -5033,10 +5088,6 @@ export async function handleChatCore({ if (promptTokens > 0) incrementTokenUsage(model, promptTokens); } } - - // Context Editing telemetry: when the delegated server-side clear actually ran, - // record the provider's cleared-token receipt under engine "context-editing" so - // it surfaces in compression analytics. Best-effort, Claude-only, non-streaming. recordContextEditingTelemetryHook({ contextEditingEnabled, provider, @@ -5051,9 +5102,6 @@ export async function handleChatCore({ tokens: usage, status: "200 OK", }).catch(() => {}); - - // Save structured call log with full payloads - const cacheUsageLogMeta = buildCacheUsageLogMeta(usage); recordNonStreamingUsageStats(usage, { traceEnabled, provider, @@ -5067,104 +5115,6 @@ export async function handleChatCore({ endpoint: endpointPath, }); - // Translate response to client's expected format (usually OpenAI) - // Pass toolNameMap so Claude OAuth proxy_ prefix is stripped in tool_use blocks (#605) - const responseToolSchemas = extractToolSchemaMap(finalBody || translatedBody || body); - let translatedResponse = needsTranslation(responsePayloadFormat, clientResponseFormat) - ? translateNonStreamingResponse( - responseBody, - responsePayloadFormat, - clientResponseFormat, - responseToolNameMap, - responseToolSchemas - ) - : responseBody; - const memoryExtractionResponse = translatedResponse; - - // T26: Strip markdown code blocks if provider format is Claude - if (sourceFormat === "claude" && !stream) { - if (typeof translatedResponse?.choices?.[0]?.message?.content === "string") { - translatedResponse.choices[0].message.content = stripMarkdownCodeFence( - translatedResponse.choices[0].message.content - ) as string; - } - } - - // T18: Normalize finish_reason to 'tool_calls' if tool calls are present - normalizeOpenAIToolFinishReasons(translatedResponse); - - // Reasoning Replay Cache (#1628): Capture reasoning_content from non-streaming responses - // with tool_calls so it can be replayed on subsequent turns (DeepSeek V4, Kimi K2, etc.) - try { - const cacheResponse = translatedResponse?.choices?.[0] - ? translatedResponse - : needsTranslation(responsePayloadFormat, FORMATS.OPENAI) - ? translateNonStreamingResponse( - responseBody, - responsePayloadFormat, - FORMATS.OPENAI, - responseToolNameMap, - responseToolSchemas - ) - : responseBody; - const firstChoice = cacheResponse?.choices?.[0]; - const msg = firstChoice?.message; - const historyMessages = (translatedBody as { messages?: unknown[] } | null | undefined) - ?.messages; - if (requiresReasoningReplay({ provider, model })) { - cacheReasoningFromAssistantMessage(msg, provider, model, { - scope: reasoningCacheScope, - historyMessages: Array.isArray(historyMessages) ? historyMessages : [], - }); - } - } catch { - // Cache capture is non-critical — never block the response - } - // Sanitize response for OpenAI SDK compatibility - // Strips non-standard fields (x_groq, usage_breakdown, service_tier, etc.) - // Extracts and tags into reasoning_content - // Source format determines output shape. If we are outputting OpenAI shape or pseudo-OpenAI shape, sanitize. - if (clientResponseFormat === FORMATS.OPENAI_RESPONSES) { - translatedResponse = sanitizeResponsesApiResponse(translatedResponse); - // Responses-API non-stream path: restore `{namespace, name}` on every - // `function_call` item that was flattened from a namespace sub-tool on - // the request side (#7936 round-trip closure). - const responseOutput = translatedResponse?.output; - if (requestToolIdentityMap && Array.isArray(responseOutput)) { - for (const item of responseOutput) { - if (item?.type !== "function_call") continue; - const identity = requestToolIdentityMap.get(item.name); - if (identity) { - item.namespace = identity.namespace; - item.name = identity.name; - } - } - } - } else if (clientResponseFormat === FORMATS.OPENAI) { - // Port of decolua/9router#517: opt-in `x-omniroute-strip-reasoning` header - // unconditionally drops `reasoning_content` from the final non-streaming - // JSON for clients (e.g. Firecrawl AI SDK) whose JSON parsers break on - // that non-standard field. Reasoning replay cache is captured above this - // sanitize step, so the cache feature is unaffected. - const stripReasoning = isStripReasoningRequested(clientRawRequest?.headers ?? null); - translatedResponse = sanitizeOpenAIResponse(translatedResponse, { - stripReasoning, - parseTextualReasoningTags: shouldParseTextualReasoningTags(provider, model), - }); - } - - // #8331: keep the client-visible metering fields real everywhere except Claude-Code-compatible - // providers, where Claude Code's own context accounting relies on the buffered number — see - // clientUsageBuffer.ts module docstring. - applyClientUsageBuffer( - translatedResponse, - finalBody || translatedBody || body, - clientResponseFormat, - { - preserveContextBudgetInVisibleUsage: isClaudeCodeCompatible, - } - ); - // #12150 P1b surface 3 (fix round 1): a video-bridge-observed request's // request- AND response-derived text both carry the full transcript (the // flattened description on the request side, the model's own reply on @@ -5189,7 +5139,7 @@ export async function handleChatCore({ webFetchFallbackPlan.toolName, ...(memoryOwnerId && memorySettings?.enabled ? MEMORY_BUILTIN_TOOL_NAMES : []), ].filter((name): name is string => Boolean(name)); - if (customSkillExecutionEnabled || builtinToolNames.length > 0) { + if (!toolLoopRan && (customSkillExecutionEnabled || builtinToolNames.length > 0)) { const skillSessionId = pipelineSessionId; translatedResponse = await handleToolCallExecution( @@ -5485,6 +5435,35 @@ export async function handleChatCore({ success: true, response: buildNonStreamingJsonResponse(translatedResponse, responseHeaders), }; + } catch (error) { + trackPendingRequest(model, provider, connectionId, false); + if (isManagedLeaseFenceError(error)) return managedLeaseFenceErrorResult(error); + if (isSemaphoreCapacityError(error)) { + appendRequestLog({ + model, + provider, + connectionId, + status: `FAILED ${error.code}`, + }).catch(() => {}); + const failureMessage = error.message || "Semaphore timeout"; + persistAttemptLogs({ + status: HTTP_STATUS.RATE_LIMITED, + error: failureMessage, + providerRequest: finalBody || translatedBody, + clientResponse: buildErrorBody(HTTP_STATUS.RATE_LIMITED, failureMessage), + claudeCacheMeta: claudePromptCacheLogMeta, + cacheSource: "upstream", + }); + persistFailureUsage(HTTP_STATUS.RATE_LIMITED, error.code); + const result = createErrorResult(HTTP_STATUS.RATE_LIMITED, failureMessage); + return { + ...result, + errorType: "account_semaphore_capacity", + errorCode: error.code, + }; + } + throw error; + } } // Streaming response @@ -5603,7 +5582,7 @@ export async function handleChatCore({ errorCode: streamErrorCode, ttft, itlMs: streamItlMs, - interrupted: streamInterrupted, + interrupted: _streamInterrupted, }) => { const normalizedStreamStatus = streamStatus || 200; if (streamCompletionRecorded) return; diff --git a/open-sse/handlers/chatCore/failureUsage.ts b/open-sse/handlers/chatCore/failureUsage.ts index d9fff0ae33..4e9c95a1e0 100644 --- a/open-sse/handlers/chatCore/failureUsage.ts +++ b/open-sse/handlers/chatCore/failureUsage.ts @@ -23,6 +23,36 @@ export function projectFailureUsageErrorCode(opts: { return errorBody.error.code || String(opts.statusCode); } +export interface FailureUsageAggregate { + prompt_tokens?: number; + completion_tokens?: number; + cache_read_input_tokens?: number; + cache_creation_input_tokens?: number; + reasoning_tokens?: number; +} + +export function toFailureUsageAggregate( + usage: + | { + prompt_tokens?: number; + completion_tokens?: number; + cache_read_input_tokens?: number; + cache_creation_input_tokens?: number; + reasoning_tokens?: number; + } + | null + | undefined +): FailureUsageAggregate | undefined { + if (!usage) return undefined; + return { + prompt_tokens: usage.prompt_tokens, + completion_tokens: usage.completion_tokens, + cache_read_input_tokens: usage.cache_read_input_tokens, + cache_creation_input_tokens: usage.cache_creation_input_tokens, + reasoning_tokens: usage.reasoning_tokens, + }; +} + export function buildFailureUsageRecord(opts: { provider: string | null | undefined; model: string | null | undefined; @@ -35,11 +65,18 @@ export function buildFailureUsageRecord(opts: { errorCode: string | null | undefined; latencyMs: number; endpoint?: string | null | undefined; + aggregate?: FailureUsageAggregate | null; }) { return { provider: opts.provider || "unknown", model: opts.model || "unknown", - tokens: { input: 0, output: 0, cacheRead: 0, cacheCreation: 0, reasoning: 0 }, + tokens: { + input: opts.aggregate?.prompt_tokens ?? 0, + output: opts.aggregate?.completion_tokens ?? 0, + cacheRead: opts.aggregate?.cache_read_input_tokens ?? 0, + cacheCreation: opts.aggregate?.cache_creation_input_tokens ?? 0, + reasoning: opts.aggregate?.reasoning_tokens ?? 0, + }, status: String(opts.statusCode), success: false, latencyMs: opts.latencyMs, diff --git a/open-sse/handlers/chatCore/memorySkillsInjection.ts b/open-sse/handlers/chatCore/memorySkillsInjection.ts index 799c834ff9..c2636cfc73 100644 --- a/open-sse/handlers/chatCore/memorySkillsInjection.ts +++ b/open-sse/handlers/chatCore/memorySkillsInjection.ts @@ -5,7 +5,7 @@ import { toMemoryRetrievalConfig, } from "@/lib/memory/settings"; import { injectMemory, shouldInjectMemory } from "@/lib/memory/injection"; -import { injectSkills } from "@/lib/skills/injection"; +import { injectSkillsWithMetadata } from "@/lib/skills/injection"; import { buildMemoryToolsForProvider } from "@/lib/skills/memoryBuiltins"; import { skillRegistry } from "@/lib/skills/registry"; import { FORMATS } from "../../translator/formats.ts"; @@ -13,6 +13,13 @@ import { detectCachingContext } from "../../services/compression/cachingAware.ts type MemorySkillsLogger = { debug?: (...args: unknown[]) => void } | null | undefined; +export interface MemorySkillsInjectionResult { + body: Record; + memorySettings: { enabled: boolean; skillsEnabled: boolean; maxTokens: number } | null; + builtinToolNames: string[]; + injectedCustomSkillNames: string[]; +} + function getToolName(tool: unknown): string { if (!tool || typeof tool !== "object") return ""; const r = tool as Record; @@ -60,11 +67,14 @@ export async function injectMemoryAndSkills({ targetFormat: string; backgroundReason: string | null; log: MemorySkillsLogger; -}) { +}): Promise { const memorySettings = memoryOwnerId ? await getMemorySettings().catch(() => DEFAULT_MEMORY_SETTINGS) : null; + const builtinOwnerSet: string[] = []; + const injectedCustomSkillNames: string[] = []; + if ( memoryOwnerId && memorySettings && @@ -178,34 +188,48 @@ export async function injectMemoryAndSkills({ return []; }) ); - const memoryTools = buildMemoryToolsForProvider( + const newMemoryTools = buildMemoryToolsForProvider( getSkillsProviderForFormat(sourceFormat) ).filter((tool) => { const record = tool as Record; const name = (record.function as Record | undefined)?.name ?? record.name; return typeof name === "string" && !existingToolNames.has(name); }); - if (memoryTools.length > 0) { + if (newMemoryTools.length > 0) { body = { ...body, - tools: [...existingTools, ...memoryTools], + tools: [...existingTools, ...newMemoryTools], }; + // Track the names of newly injected memory tools for the owner set. + builtinOwnerSet.push( + ...newMemoryTools + .map((tool) => { + const record = tool as Record; + const name = + (record.function as Record | undefined)?.name ?? record.name; + return typeof name === "string" ? name : ""; + }) + .filter(Boolean) + ); log?.debug?.( "MEMORY", - `Injected ${memoryTools.length} memory tool(s) for key=${memoryOwnerId}` + `Injected ${newMemoryTools.length} memory tool(s) for key=${memoryOwnerId}` ); } } - if (memoryOwnerId && memorySettings?.skillsEnabled) { + if (memoryOwnerId && memorySettings?.skillsEnabled && body.stream !== true) { // Ensure the registry cache is warm before listing: on a cold/fresh // process skills that exist only in the DB would be missed (false // negative -> silent skip). loadFromDatabase() is a no-op when the cache // is already warm (TTL = 60 s), so repeated calls are cheap. Mirrors the // pattern in src/lib/skills/interception.ts (#2815). + // Memory builtins and registered Skills are only executed by the + // non-streaming server-owned tool loop; stream clients execute tools + // client-side, so we skip injection for stream requests. await skillRegistry.loadFromDatabase(memoryOwnerId); const existingTools = Array.isArray(body.tools) ? body.tools : []; - const mergedTools = injectSkills({ + const { tools: mergedTools, injectedNames } = injectSkillsWithMetadata({ provider: getSkillsProviderForFormat(sourceFormat), existingTools, apiKeyId: memoryOwnerId, @@ -225,6 +249,7 @@ export async function injectMemoryAndSkills({ ...body, tools: mergedTools, }; + injectedCustomSkillNames.push(...injectedNames); log?.debug?.("SKILLS", `Injected ${mergedTools.length - existingTools.length} skills`); } } @@ -236,5 +261,45 @@ export async function injectMemoryAndSkills({ }; } - return { body, memorySettings }; + return { body, memorySettings, builtinToolNames: builtinOwnerSet, injectedCustomSkillNames }; +} + +interface FallbackPlan { + enabled: boolean; + toolName: string | null; + convertedToolCount: number; +} + +/** + * Pure helper: merge web-search/web-fetch fallback tool names into the + * builtin owner set. Adds a name only when plan.enabled===true, + * plan.convertedToolCount>0, plan.toolName is non-null, and that name + * did not already exist in the pre-conversion client tools (builtinToolNames) + * OR in the original client tool names captured before fallback injection. + * Does not mutate its input; returns a new result. + */ +export function mergeInjectedFallbackOwnerNames( + injectionResult: { builtinToolNames: string[] }, + plans: FallbackPlan[], + preConversionClientToolNames?: string[] +): { builtinToolNames: string[] } { + const existing = new Set(injectionResult.builtinToolNames); + if (preConversionClientToolNames) { + for (const name of preConversionClientToolNames) { + existing.add(name); + } + } + const extraNames: string[] = []; + for (const plan of plans) { + if ( + plan.enabled && + plan.convertedToolCount > 0 && + plan.toolName && + !existing.has(plan.toolName) + ) { + extraNames.push(plan.toolName); + existing.add(plan.toolName); + } + } + return { builtinToolNames: [...injectionResult.builtinToolNames, ...extraNames] }; } diff --git a/open-sse/handlers/chatCore/nonStreamingClientTranslate.ts b/open-sse/handlers/chatCore/nonStreamingClientTranslate.ts new file mode 100644 index 0000000000..79821278f0 --- /dev/null +++ b/open-sse/handlers/chatCore/nonStreamingClientTranslate.ts @@ -0,0 +1,159 @@ +/** + * Client translation for non-streaming responses. + * Extracted from chatCore.ts (lines ~5098-5195) by symbol boundaries. + * + * Handles: translate, tool-name restore, finish-reason normalization, sanitize, + * reasoning replay capture, and client usage buffer application. + * + * Phase distinction: + * - "final": applies applyClientUsageBuffer (normalizes visible usage fields) + * - "intermediate": skips usage buffer (raw usage preserved for aggregation) + */ + +import type { + NonStreamingClientTranslateInput, + NonStreamingClientTranslateResult, +} from "@/lib/skills/toolLoopTypes.ts"; +import { needsTranslation } from "../../translator/index.ts"; +import { FORMATS } from "../../translator/formats.ts"; +import { translateNonStreamingResponse } from "../responseTranslator.ts"; +import { extractToolSchemaMap } from "../../translator/response/openai-responses/toolSchemas.ts"; +import { stripMarkdownCodeFence } from "../../utils/aiSdkCompat.ts"; +import { normalizeOpenAIToolFinishReasons } from "./passthroughToolNames.ts"; +import { + cacheReasoningFromAssistantMessage, + requiresReasoningReplay, +} from "../../services/reasoningCache.ts"; +import { + sanitizeOpenAIResponse, + sanitizeResponsesApiResponse, + shouldParseTextualReasoningTags, +} from "../responseSanitizer.ts"; +import { isStripReasoningRequested } from "./headers.ts"; +import { applyClientUsageBuffer } from "./clientUsageBuffer.ts"; + +export type { NonStreamingClientTranslateInput, NonStreamingClientTranslateResult }; + +/** + * Translate a non-streaming provider response to the client's expected format. + * + * All of: translate, tool-name identity restore, finish-reason normalization, + * and sanitize are applied every round (both intermediate and final). + * `applyClientUsageBuffer` is applied only for `phase === "final"`. + * Reasoning replay capture runs every round. + */ +export function translateNonStreamingClientResponse( + input: NonStreamingClientTranslateInput +): NonStreamingClientTranslateResult { + const { + responseBody, + responsePayloadFormat, + clientResponseFormat, + sourceFormat, + provider, + model, + requestBody, + responseToolNameMap, + requestToolIdentityMap, + reasoningCacheScope, + clientHeaders, + isClaudeCodeCompatible, + phase, + } = input; + + // ── Extract tool schemas for schema-aware translation ────────────────────── + const finalBody = requestBody as Record | null; + const responseToolSchemas = extractToolSchemaMap(finalBody || responseBody); + + // ── Translate response to client's expected format ───────────────────────── + let translatedResponse = needsTranslation(responsePayloadFormat, clientResponseFormat) + ? translateNonStreamingResponse( + responseBody, + responsePayloadFormat, + clientResponseFormat, + responseToolNameMap, + responseToolSchemas + ) + : responseBody; + const responseForMemoryExtraction = translatedResponse; + + // ── T26: Strip markdown code blocks if provider format is Claude ─────────── + if (sourceFormat === "claude") { + if (typeof translatedResponse?.choices?.[0]?.message?.content === "string") { + translatedResponse.choices[0].message.content = stripMarkdownCodeFence( + translatedResponse.choices[0].message.content + ) as string; + } + } + + // ── T18: Normalize finish_reason to 'tool_calls' if tool calls present ───── + normalizeOpenAIToolFinishReasons(translatedResponse); + + // ── Reasoning Replay Cache (#1628) ──────────────────────────────────────── + // Capture reasoning_content from non-streaming responses with tool_calls + // so it can be replayed on subsequent turns. + try { + const cacheResponse = translatedResponse?.choices?.[0] + ? translatedResponse + : needsTranslation(responsePayloadFormat, FORMATS.OPENAI) + ? translateNonStreamingResponse( + responseBody, + responsePayloadFormat, + FORMATS.OPENAI, + responseToolNameMap, + responseToolSchemas + ) + : responseBody; + const firstChoice = cacheResponse?.choices?.[0]; + const msg = firstChoice?.message; + // Prefer explicit historyMessages (parent: translatedBody.messages). Do not + // overload requestBody — Responses-shaped finalBody has `input`, not `messages`. + const historyMessages = Array.isArray(input.historyMessages) + ? input.historyMessages + : (finalBody as { messages?: unknown[] } | null | undefined)?.messages; + if (requiresReasoningReplay({ provider, model })) { + cacheReasoningFromAssistantMessage(msg, provider, model, { + scope: reasoningCacheScope, + historyMessages: Array.isArray(historyMessages) ? historyMessages : [], + }); + } + } catch { + // Cache capture is non-critical — never block the response + } + + // ── Sanitize response for SDK compatibility ──────────────────────────────── + if (clientResponseFormat === FORMATS.OPENAI_RESPONSES) { + translatedResponse = sanitizeResponsesApiResponse(translatedResponse); + // Restore {namespace, name} on function_call items for round-trip closure (#7936) + const responseOutput = translatedResponse?.output; + if (requestToolIdentityMap && Array.isArray(responseOutput)) { + for (const item of responseOutput) { + if (item?.type !== "function_call") continue; + const identity = requestToolIdentityMap.get(item.name); + if (identity) { + item.namespace = identity.namespace; + item.name = identity.name; + } + } + } + } else if (clientResponseFormat === FORMATS.OPENAI) { + const stripReasoning = isStripReasoningRequested(clientHeaders ?? null); + translatedResponse = sanitizeOpenAIResponse(translatedResponse, { + stripReasoning, + parseTextualReasoningTags: shouldParseTextualReasoningTags(provider, model), + }); + } + + // ── Client usage buffer (#8331) ─────────────────────────────────────────── + // Only apply for final phase; intermediate preserves raw usage for aggregation. + if (phase === "final") { + applyClientUsageBuffer(translatedResponse, finalBody || responseBody, clientResponseFormat, { + preserveContextBudgetInVisibleUsage: isClaudeCodeCompatible, + }); + } + + return { + response: translatedResponse, + responseForMemoryExtraction, + }; +} diff --git a/open-sse/handlers/chatCore/nonStreamingFinalization.ts b/open-sse/handlers/chatCore/nonStreamingFinalization.ts new file mode 100644 index 0000000000..7c8325dee4 --- /dev/null +++ b/open-sse/handlers/chatCore/nonStreamingFinalization.ts @@ -0,0 +1,133 @@ +/** + * Request-level finalization for non-streaming chat. + * Success and failure each write usage/cost/quota/attempt/pending once. + */ + +import type { + ChatCoreErrorResult, + ProviderLegUsage, + ServerOwnedToolLoopResult, +} from "@/lib/skills/toolLoopTypes.ts"; +import type { PersistAttemptLogsArgs } from "./attemptLogging.ts"; +import { type FailureUsageAggregate, toFailureUsageAggregate } from "./failureUsage.ts"; + +export type NonStreamingFinalizationPlan = + | { + kind: "success"; + usage: ProviderLegUsage | null; + totalCostUsd: number; + receiptCount: number; + } + | { + kind: "failure"; + error: ChatCoreErrorResult; + usage: ProviderLegUsage | null; + totalCostUsd: number; + receiptCount: number; + }; + +export interface NonStreamingFinalizationDeps { + writeUsage: (plan: NonStreamingFinalizationPlan) => void | Promise; + writeCost: (totalCostUsd: number) => void; + scheduleQuota: (plan: NonStreamingFinalizationPlan) => void | Promise; + writeAttempt: (plan: NonStreamingFinalizationPlan) => void; + finalizePending: (plan: NonStreamingFinalizationPlan) => void; +} + +function missingError(): ChatCoreErrorResult { + return { + success: false, + status: 500, + response: new Response(null, { status: 500 }), + error: "Missing tool-loop error result", + errorCode: "internal_error", + }; +} + +export function buildNonStreamingFinalizationPlan( + loop: ServerOwnedToolLoopResult +): NonStreamingFinalizationPlan { + const usage = loop.cumulativeUsage; + const totalCostUsd = loop.totalCostUsd; + const receiptCount = loop.receipts.length; + if (loop.kind === "error") { + return { + kind: "failure", + error: loop.errorResult ?? missingError(), + usage, + totalCostUsd, + receiptCount, + }; + } + return { + kind: "success", + usage, + totalCostUsd, + receiptCount, + }; +} + +export async function finalizeNonStreamingRequest( + plan: NonStreamingFinalizationPlan, + deps: NonStreamingFinalizationDeps +): Promise { + await deps.writeUsage(plan); + deps.writeCost(plan.totalCostUsd); + if (plan.kind === "success") { + await deps.scheduleQuota(plan); + } + deps.writeAttempt(plan); + deps.finalizePending(plan); +} + +export async function finalizeToolLoopError(input: { + loop: ServerOwnedToolLoopResult; + model: string; + provider: string; + connectionId?: string; + providerRequest?: Record; + persistFailureUsage: ( + status: number, + errorCode: string, + usage?: FailureUsageAggregate | null + ) => void; + persistAttemptLogs: (params: PersistAttemptLogsArgs) => void; + trackPendingRequest: ( + model: string, + provider: string, + connectionId?: string, + isPending?: boolean + ) => void; +}): Promise { + const plan = buildNonStreamingFinalizationPlan(input.loop); + const err = plan.kind === "failure" ? plan.error : missingError(); + await finalizeNonStreamingRequest(plan, { + writeUsage: () => { + input.persistFailureUsage( + err.status, + err.errorCode || `upstream_${err.status}`, + toFailureUsageAggregate(plan.usage) + ); + }, + writeCost: () => {}, + scheduleQuota: () => {}, + writeAttempt: () => { + input.persistAttemptLogs({ + status: err.status, + error: err.error || "Provider request failed", + providerRequest: input.providerRequest, + clientResponse: { + error: { + message: err.error || "Provider request failed", + type: err.errorType || "api_error", + }, + }, + cacheSource: "upstream", + }); + }, + finalizePending: () => { + input.trackPendingRequest(input.model, input.provider, input.connectionId, false); + }, + }); + return err; +} diff --git a/open-sse/handlers/chatCore/nonStreamingProviderLeg.ts b/open-sse/handlers/chatCore/nonStreamingProviderLeg.ts new file mode 100644 index 0000000000..0e8f2e5969 --- /dev/null +++ b/open-sse/handlers/chatCore/nonStreamingProviderLeg.ts @@ -0,0 +1,1144 @@ +/** + * Non-streaming provider leg - one round of provider execution. + * Extracted from chatCore.ts (lines ~3804-5096) by symbol boundaries. + * + * Owns: expectedConnectionId, allowAccountRotation, allowModelFallback. + * Does NOT write terminal side effects (usage_history, cost, memory extraction). + * Returns usage in both ok and error results. + */ + +import type { + ProviderLegUsage, + ProviderLegReceipt, + ChatCoreErrorResult, + NonStreamingProviderLegResult, +} from "@/lib/skills/toolLoopTypes.ts"; +import type { + ProviderExecutionOutcome, + ProviderExecutionPolicy, +} from "./providerExecutionPipeline.ts"; +import { translateNonStreamingClientResponse } from "./nonStreamingClientTranslate.ts"; +import { parseNonStreamingResponseBody, isJsonRecord } from "./nonStreamingResponseParse.ts"; +import { restoreNonStreamingToolNames } from "./passthroughToolNames.ts"; +import { extractUsageFromResponse } from "../usageExtractor.ts"; +import { sanitizeUsagePayloadForRequest } from "../../utils/usageTracking.ts"; +import { createErrorResult, formatProviderError } from "../../utils/error.ts"; +import { unwrapClinepassEnvelope } from "../../utils/clinepassEnvelope.ts"; +import { unwrapClineNonStreamingEnvelope } from "./clineResponseEnvelope.ts"; +import { + isModelUnavailableError, + getNextFamilyFallback, + isContextOverflowError, + findLargerContextModel, + getModelFamily, +} from "../../services/modelFamilyFallback.ts"; +import { isEmptyContentResponse } from "../../services/errorClassifier.ts"; +import { FORMATS } from "../../translator/formats.ts"; + +/* -- exported types -------------------------------------------------------- */ + +export interface ChatCoreExecutorResult { + response: Response; + url: string; + headers: Record; + transformedBody: unknown; + transport?: string; + _executionCredentials?: Record; + _accountSemaphoreRelease?: () => void; +} + +export interface ProviderLegRotationPolicy { + allowAccountRotation?: boolean; +} + +export interface ProviderLegInput { + phase: "initial" | "follow-up"; + sourceBody: Record; + expectedConnectionId?: string; + allowAccountRotation: boolean; + allowModelFallback: boolean; + executeProviderRequest: ( + model: string, + allowDedup: boolean, + policy?: ProviderLegRotationPolicy + ) => Promise; + /** + * Optional 6b seam. When set, the first send goes through the shared + * pipeline (one wire send + account/model recovery). ClinePass empty + * retry and empty-content fallback still use executeProviderRequest. + */ + runProviderExecution?: (input: { + policy: Readonly; + model: string; + translatedBody: Record; + }) => Promise; + setRequestWireState: (state: { + translatedBody: Record; + effectiveModel: string; + }) => void; + sourceFormat?: string; + targetFormat?: string; + clientResponseFormat?: string; + provider?: string; + model?: string; + connectionId?: string; + getCurrentConnectionId?: () => string; + effectiveModel?: string; + translatedBody?: Record; + toolNameMap?: Map | null; + requestToolIdentityMap?: Map | null; + reasoningCacheScope?: string | null; + clientHeaders?: Headers | Record | null; + isClaudeCodeCompatible?: boolean; + sleep?: (ms: number) => Promise; + log?: { + info?: (tag: string, msg: string) => void; + warn?: (tag: string, msg: string) => void; + error?: (tag: string, msg: string) => void; + }; +} + +/* -- helpers --------------------------------------------------------------- */ + +function buildReceipt( + input: ProviderLegInput, + params: { + httpStatus: number; + errorType: string | null; + usage: ProviderLegUsage | null; + termination: string; + latencyMs: number; + startedAt: string; + endedAt: string; + connectionId: string; + model: string; + } +): ProviderLegReceipt { + return { + index: input.phase === "initial" ? 0 : 1, + connectionId: params.connectionId, + provider: input.provider ?? "unknown", + model: params.model, + startedAt: params.startedAt, + endedAt: params.endedAt, + latencyMs: params.latencyMs, + httpStatus: params.httpStatus, + errorType: params.errorType, + usage: params.usage, + serviceTier: null, + computedCostUsd: null, + toolCalls: [], + termination: params.termination, + clientVisible: true, + }; +} + +function extractUsage( + responseBody: Record, + provider: string +): ProviderLegUsage | null { + const raw = extractUsageFromResponse(responseBody, provider); + if (!raw || typeof raw !== "object") return null; + const r = raw as Record; + const pt = typeof r.prompt_tokens === "number" ? r.prompt_tokens : 0; + const ct = typeof r.completion_tokens === "number" ? r.completion_tokens : 0; + return { + prompt_tokens: pt, + completion_tokens: ct, + total_tokens: typeof r.total_tokens === "number" ? r.total_tokens : pt + ct, + cached_tokens: typeof r.cached_tokens === "number" ? r.cached_tokens : undefined, + cache_read_input_tokens: + typeof r.cache_read_input_tokens === "number" ? r.cache_read_input_tokens : undefined, + cache_creation_input_tokens: + typeof r.cache_creation_input_tokens === "number" ? r.cache_creation_input_tokens : undefined, + reasoning_tokens: typeof r.reasoning_tokens === "number" ? r.reasoning_tokens : undefined, + }; +} + +function legError( + status: number, + message: string, + originalError?: unknown, + retryAfterMs?: number | null, + errorCode?: string, + errorType?: string, + opts?: { passthrough?: boolean } +): ChatCoreErrorResult { + // createErrorResult(..., errorCode, errorType, upstreamDetails, opts) + // - opts is the 7th arg, not the 6th. + const result = createErrorResult( + status, + message, + retryAfterMs ?? null, + errorCode, + errorType, + undefined, + opts + ); + (result as ChatCoreErrorResult).originalError = originalError; + return result as ChatCoreErrorResult; +} + +function checkConnectionIdentity( + input: ProviderLegInput, + startMs: number, + startedAt: string, + currentModel: string, + context: string +): NonStreamingProviderLegResult | null { + if (!input.expectedConnectionId || !input.getCurrentConnectionId) return null; + const currentConnId = input.getCurrentConnectionId(); + if (currentConnId !== input.expectedConnectionId) { + const receipt = buildReceipt(input, { + httpStatus: 409, + errorType: "lease_error", + usage: null, + termination: "connection_mismatch", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId: currentConnId, + model: currentModel, + }); + const errorResult = legError( + 409, + `Follow-up connection does not match initial connection (${context})`, + new Error("connection_mismatch"), + null, + "LEASE_CONNECTION_MISMATCH", + "lease_error" + ); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + return null; +} + +function parseRetryAfterMs(response: Response): number | null { + const retryAfterHeader = response.headers?.get?.("retry-after"); + if (!retryAfterHeader) return null; + const retryAfterSec = Number.parseInt(retryAfterHeader, 10); + if (Number.isFinite(retryAfterSec) && retryAfterSec > 0) { + return retryAfterSec * 1000; + } + const retryAfterDate = new Date(retryAfterHeader).getTime(); + if (Number.isFinite(retryAfterDate) && retryAfterDate > Date.now()) { + return retryAfterDate - Date.now(); + } + return null; +} + + +function finishOk( + input: ProviderLegInput, + params: { + responseBody: Record; + responsePayloadFormat: string; + looksLikeSSE: boolean; + requestBody: Record; + transformedBody: unknown; + model: string; + connectionId: string; + headers: Headers; + startMs: number; + startedAt: string; + sourceFormat: string; + targetFormat: string; + clientResponseFormat: string; + provider: string; + upstreamResponse?: Response; + requestHeaders?: Record; + requestUrl?: string; + } +): NonStreamingProviderLegResult { + // F-02: restore + sanitize + translate is the only success tail. + // Fallback/retry must not skip this with responseToolNameMap: null. + const restoreClaudeNames = params.sourceFormat === "claude" && params.targetFormat === "claude"; + let body = params.responseBody; + let responseToolNameMap: Map | null; + [body, responseToolNameMap] = restoreNonStreamingToolNames( + body, + input.toolNameMap ?? null, + params.transformedBody, + restoreClaudeNames + ); + sanitizeUsagePayloadForRequest( + body, + params.transformedBody || params.requestBody, + params.responsePayloadFormat + ); + const usage = extractUsage(body, params.provider); + const clientTranslate = translateNonStreamingClientResponse({ + responseBody: body, + responsePayloadFormat: params.responsePayloadFormat, + clientResponseFormat: params.clientResponseFormat, + sourceFormat: params.sourceFormat, + provider: params.provider, + model: params.model, + requestBody: params.requestBody, + historyMessages: (input.translatedBody as { messages?: unknown[] } | null | undefined) + ?.messages, + responseToolNameMap, + requestToolIdentityMap: input.requestToolIdentityMap ?? null, + reasoningCacheScope: input.reasoningCacheScope ?? null, + clientHeaders: input.clientHeaders ?? null, + isClaudeCodeCompatible: input.isClaudeCodeCompatible ?? false, + phase: input.phase === "initial" ? "final" : "intermediate", + }); + const receipt = buildReceipt(input, { + httpStatus: 200, + errorType: null, + usage, + termination: "completed", + latencyMs: Date.now() - params.startMs, + startedAt: params.startedAt, + endedAt: new Date().toISOString(), + connectionId: params.connectionId, + model: params.model, + }); + return { + kind: "ok", + response: clientTranslate.response, + responseForMemoryExtraction: clientTranslate.responseForMemoryExtraction, + providerBody: body, + providerRequest: input.translatedBody ?? {}, + usage, + responsePayloadFormat: params.responsePayloadFormat, + looksLikeSSE: params.looksLikeSSE, + connectionId: params.connectionId, + headers: params.headers, + upstreamResponse: params.upstreamResponse, + requestHeaders: params.requestHeaders, + requestUrl: params.requestUrl, + receipt, + }; +} + +const DEFAULT_CLINEPASS_sleep = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); + +/* -- main leg function ----------------------------------------------------- */ + +export async function runNonStreamingProviderLeg( + input: ProviderLegInput +): Promise { + const startedAt = new Date().toISOString(); + const startMs = Date.now(); + const effectiveModel = input.effectiveModel ?? input.model ?? "unknown"; + const currentModel = effectiveModel; + const provider = input.provider ?? "unknown"; + const sourceFormat = input.sourceFormat ?? "openai"; + const targetFormat = input.targetFormat ?? "openai"; + const clientResponseFormat = input.clientResponseFormat ?? "openai"; + const connectionId = input.connectionId ?? "unknown"; + const log = input.log; + + // -- Phase policy: follow-up blocks rotation and fallback ------------------- + const allowAccountRotation = input.phase === "follow-up" ? false : input.allowAccountRotation; + const allowModelFallback = input.phase === "follow-up" ? false : input.allowModelFallback; + + // -- Rotation policy for executor ------------------------------------------ + const rotationPolicy: ProviderLegRotationPolicy = { + allowAccountRotation, + }; + + // F-01: lease identity is checked BEFORE any wire send. A follow-up whose + // lease already rotated must not spend an upstream round then 409. + const mismatchAtEntry = checkConnectionIdentity( + input, + startMs, + startedAt, + currentModel, + "before_executor" + ); + if (mismatchAtEntry) return mismatchAtEntry; + + // -- Execute request -------------------------------------------------------- + // Update wire state before every executor call + input.setRequestWireState({ + translatedBody: (input.translatedBody ?? input.sourceBody) as Record, + effectiveModel: currentModel, + }); + + const policy: Readonly = { + allowAccountRotation, + allowModelFallback, + expectedConnectionId: input.expectedConnectionId, + }; + + let executorResult: ChatCoreExecutorResult; + try { + if (input.runProviderExecution) { + const outcome = await input.runProviderExecution({ + policy, + model: currentModel, + translatedBody: (input.translatedBody ?? input.sourceBody) as Record, + }); + if (outcome.kind === "error") { + if ( + outcome.result.status === 409 && + outcome.result.errorCode === "LEASE_CONNECTION_MISMATCH" + ) { + const receipt = buildReceipt(input, { + httpStatus: 409, + errorType: "lease_error", + usage: null, + termination: "connection_mismatch", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId: outcome.connectionId || connectionId, + model: outcome.model || currentModel, + }); + return { + kind: "error", + result: outcome.result, + receipt, + usage: null, + }; + } + const raw = outcome.result.error || "Provider request failed"; + const alreadyPrefixed = /^\[[^\]]+\]:/.test(raw); + const formatted = alreadyPrefixed + ? raw + : formatProviderError( + new Error(raw), + provider, + outcome.model || currentModel, + outcome.result.status + ); + const receipt = buildReceipt(input, { + httpStatus: outcome.result.status, + errorType: outcome.result.errorType ?? null, + usage: outcome.providerUsage, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId: outcome.connectionId || connectionId, + model: outcome.model || currentModel, + }); + return { + kind: "error", + result: { + ...legError( + outcome.result.status, + formatted, + undefined, + null, + outcome.result.errorCode, + outcome.result.errorType, + { passthrough: input.sourceFormat === "claude" } + ), + response: outcome.result.response, + }, + receipt, + usage: outcome.providerUsage, + }; + } + executorResult = { + response: outcome.response, + url: outcome.url, + headers: outcome.headers, + transformedBody: outcome.transformedBody, + }; + } else { + executorResult = await input.executeProviderRequest( + currentModel, + input.phase === "initial", + rotationPolicy + ); + } + } catch (error) { + if ( + !!error && + typeof error === "object" && + ((error as { code?: unknown }).code === "SEMAPHORE_TIMEOUT" || + (error as { code?: unknown }).code === "SEMAPHORE_QUEUE_FULL") + ) { + throw error; + } + const failureStatus = + error instanceof Error && error.name === "AbortError" + ? 499 + : error instanceof Error && error.name === "TimeoutError" + ? 504 + : 502; + const failureMessage = + error instanceof Error + ? formatProviderError(error, provider, currentModel, failureStatus) + : "Provider request failed"; + const receipt = buildReceipt(input, { + httpStatus: failureStatus, + errorType: null, + usage: null, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId, + model: currentModel, + }); + const errorResult = legError(failureStatus, failureMessage, error); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + + const providerResponse = executorResult.response; + const finalBody = executorResult.transformedBody as Record | null; + + // -- Connection mismatch check ---------------------------------------------- + if (input.expectedConnectionId && input.getCurrentConnectionId) { + const currentConnId = input.getCurrentConnectionId(); + if (currentConnId !== input.expectedConnectionId) { + const receipt = buildReceipt(input, { + httpStatus: 409, + errorType: "lease_error", + usage: null, + termination: "connection_mismatch", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId: currentConnId, + model: currentModel, + }); + const errorResult = legError( + 409, + "Follow-up connection does not match initial connection", + new Error("connection_mismatch"), + null, + "LEASE_CONNECTION_MISMATCH", + "lease_error" + ); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + } + + // -- Provider HTTP error classification ------------------------------------- + if (!providerResponse.ok) { + let statusCode = providerResponse.status; + let message = ""; + let retryAfterMs: number | null = null; + let upstreamErrorCode: string | undefined; + let upstreamErrorType: string | undefined; + let parsedErrorBody: Record = {}; + + try { + const errorBodyText = await providerResponse.text(); + try { + parsedErrorBody = JSON.parse(errorBodyText) as Record; + } catch { + // non-JSON error body + } + const errObj = (parsedErrorBody.error ?? parsedErrorBody) as + Record | undefined; + message = + (typeof errObj?.message === "string" ? errObj.message : null) ?? + errorBodyText.slice(0, 200) ?? + "Provider request failed"; + upstreamErrorCode = typeof errObj?.code === "string" ? errObj.code : undefined; + upstreamErrorType = typeof errObj?.type === "string" ? errObj.type : undefined; + } catch { + message = "Provider request failed"; + } + + // Parse Retry-After from upstream response headers + retryAfterMs = parseRetryAfterMs(providerResponse); + + // -- Model-unavailable -> family fallback (initial only) -------------------- + if (allowModelFallback && isModelUnavailableError(statusCode, message, provider)) { + const triedModels = new Set([currentModel]); + const nextModel = getNextFamilyFallback(currentModel, triedModels, provider); + if (nextModel) { + triedModels.add(nextModel); + input.setRequestWireState({ + translatedBody: { ...input.translatedBody, model: nextModel } as Record, + effectiveModel: nextModel, + }); + log?.info?.("MODEL_FALLBACK", `${currentModel} unavailable -> trying ${nextModel}`); + try { + // Connection check before fallback executor + const mismatchBefore = checkConnectionIdentity( + input, + startMs, + startedAt, + nextModel, + "before_fallback" + ); + if (mismatchBefore) return mismatchBefore; + + const fallbackResult = await input.executeProviderRequest(nextModel, false); + + // Connection check after fallback executor + const mismatchAfter = checkConnectionIdentity( + input, + startMs, + startedAt, + nextModel, + "after_fallback" + ); + if (mismatchAfter) return mismatchAfter; + + if (fallbackResult.response.ok) { + const fallbackParsed = await parseNonStreamingResponseBody({ + providerResponse: fallbackResult.response, + upstreamStream: false, + providerHeaders: new Headers(fallbackResult.headers), + finalBody: fallbackResult.transformedBody as Record | null, + targetFormat, + model: nextModel, + log, + }); + if (fallbackParsed.kind !== "invalid_sse" && fallbackParsed.kind !== "invalid_json") { + return finishOk(input, { + responseBody: fallbackParsed.responseBody, + responsePayloadFormat: fallbackParsed.responsePayloadFormat, + looksLikeSSE: fallbackParsed.looksLikeSSE, + requestBody: (fallbackResult.transformedBody || + input.translatedBody || + input.sourceBody) as Record, + transformedBody: fallbackResult.transformedBody, + model: nextModel, + connectionId: input.getCurrentConnectionId?.() ?? connectionId, + headers: new Headers(fallbackResult.headers), + startMs, + startedAt, + sourceFormat, + targetFormat, + clientResponseFormat, + provider, + upstreamResponse: fallbackResult.response, + requestHeaders: fallbackResult.headers, + requestUrl: fallbackResult.url, + }); + } + } + } catch { + // fallback also failed - fall through to standard error + } + } + } + + // -- Context overflow -> family fallback (initial only) --------------------- + if (allowModelFallback && isContextOverflowError(statusCode, message)) { + const triedModels = new Set([currentModel]); + const familyCandidates = getModelFamily(currentModel, provider).filter( + (m) => m !== currentModel && !triedModels.has(m) + ); + const nextModel = + findLargerContextModel(currentModel, familyCandidates, provider) ?? + getNextFamilyFallback(currentModel, triedModels, provider); + if (nextModel) { + triedModels.add(nextModel); + input.setRequestWireState({ + translatedBody: { ...input.translatedBody, model: nextModel } as Record, + effectiveModel: nextModel, + }); + log?.info?.("CONTEXT_OVERFLOW_FALLBACK", `${currentModel} overflow -> trying ${nextModel}`); + try { + // Connection check before fallback executor + const mismatchBefore = checkConnectionIdentity( + input, + startMs, + startedAt, + nextModel, + "before_fallback" + ); + if (mismatchBefore) return mismatchBefore; + + const fallbackResult = await input.executeProviderRequest(nextModel, false); + + // Connection check after fallback executor + const mismatchAfter = checkConnectionIdentity( + input, + startMs, + startedAt, + nextModel, + "after_fallback" + ); + if (mismatchAfter) return mismatchAfter; + + if (fallbackResult.response.ok) { + const fallbackParsed = await parseNonStreamingResponseBody({ + providerResponse: fallbackResult.response, + upstreamStream: false, + providerHeaders: new Headers(fallbackResult.headers), + finalBody: fallbackResult.transformedBody as Record | null, + targetFormat, + model: nextModel, + log, + }); + if (fallbackParsed.kind !== "invalid_sse" && fallbackParsed.kind !== "invalid_json") { + return finishOk(input, { + responseBody: fallbackParsed.responseBody, + responsePayloadFormat: fallbackParsed.responsePayloadFormat, + looksLikeSSE: fallbackParsed.looksLikeSSE, + requestBody: (fallbackResult.transformedBody || + input.translatedBody || + input.sourceBody) as Record, + transformedBody: fallbackResult.transformedBody, + model: nextModel, + connectionId: input.getCurrentConnectionId?.() ?? connectionId, + headers: new Headers(fallbackResult.headers), + startMs, + startedAt, + sourceFormat, + targetFormat, + clientResponseFormat, + provider, + upstreamResponse: fallbackResult.response, + requestHeaders: fallbackResult.headers, + requestUrl: fallbackResult.url, + }); + } + } + } catch { + // fallback also failed - fall through to standard error + } + } + } + + // -- Standard error return ------------------------------------------------- + const errMsg = formatProviderError(new Error(message), provider, currentModel, statusCode); + // Extract usage from error body if present (some providers include usage in error responses) + let usage: ProviderLegUsage | null = null; + try { + const rawUsage = extractUsageFromResponse(parsedErrorBody, provider); + if (rawUsage && typeof rawUsage === "object") { + const r = rawUsage as Record; + const pt = typeof r.prompt_tokens === "number" ? r.prompt_tokens : 0; + const ct = typeof r.completion_tokens === "number" ? r.completion_tokens : 0; + usage = { + prompt_tokens: pt, + completion_tokens: ct, + total_tokens: typeof r.total_tokens === "number" ? r.total_tokens : pt + ct, + }; + } + } catch { + // usage extraction from error body is best-effort + } + const receipt = buildReceipt(input, { + httpStatus: statusCode, + errorType: upstreamErrorCode ?? null, + usage, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId, + model: currentModel, + }); + const errorResult = legError( + statusCode, + errMsg, + new Error(message), + retryAfterMs, + upstreamErrorCode, + upstreamErrorType, + { passthrough: sourceFormat === FORMATS.CLAUDE } + ); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage, + }; + } + + // -- Non-streaming response parsing (body read exactly once) ---------------- + const parsed = await parseNonStreamingResponseBody({ + providerResponse, + upstreamStream: false, + providerHeaders: new Headers(executorResult.headers), + finalBody, + targetFormat, + model: currentModel, + log, + }); + + if (parsed.kind === "invalid_sse") { + const receipt = buildReceipt(input, { + httpStatus: 502, + errorType: "invalid_sse_payload", + usage: null, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId, + model: currentModel, + }); + const errorResult = legError( + 502, + parsed.message, + new Error(parsed.message), + null, + "invalid_sse_payload", + "invalid_sse_payload" + ); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + + if (parsed.kind === "invalid_json") { + const receipt = buildReceipt(input, { + httpStatus: 502, + errorType: "invalid_json_payload", + usage: null, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId, + model: currentModel, + }); + const errorResult = legError( + 502, + parsed.message, + new Error(parsed.message), + null, + "invalid_json_payload", + "invalid_json_payload" + ); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + + let responseBody = parsed.responseBody; + const responsePayloadFormat = parsed.responsePayloadFormat; + const looksLikeSSE = parsed.looksLikeSSE; + + // -- ClinePass envelope unwrap + retry -------------------------------------- + if (provider === "clinepass") { + let { body: unwrapped, error: envError } = unwrapClinepassEnvelope(responseBody, provider); + if (envError && /empty/i.test(envError.message || "")) { + log?.warn?.("RETRY", "clinepass returned empty content, retrying once after 2s"); + const sleepFn = input.sleep ?? DEFAULT_CLINEPASS_sleep; + await sleepFn(2000); + try { + // Connection check before retry executor + const mismatchBefore = checkConnectionIdentity( + input, + startMs, + startedAt, + currentModel, + "before_clinepass_retry" + ); + if (mismatchBefore) return mismatchBefore; + + const retryResult = await input.executeProviderRequest(currentModel, false); + + // Connection check after retry executor + const mismatchAfter = checkConnectionIdentity( + input, + startMs, + startedAt, + currentModel, + "after_clinepass_retry" + ); + if (mismatchAfter) return mismatchAfter; + + if (retryResult?.response?.ok) { + const retryParsed = await parseNonStreamingResponseBody({ + providerResponse: retryResult.response, + upstreamStream: undefined, + providerHeaders: new Headers(retryResult.headers), + finalBody: retryResult.transformedBody as Record | null, + targetFormat, + model: currentModel, + log, + }); + if (retryParsed.kind !== "invalid_sse" && retryParsed.kind !== "invalid_json") { + ({ body: unwrapped, error: envError } = unwrapClinepassEnvelope( + retryParsed.responseBody, + provider + )); + if (!envError && isJsonRecord(unwrapped)) { + responseBody = unwrapped as Record; + } + // If retry succeeded and no envelope error, continue with retry parsed body + if (!envError) { + return finishOk(input, { + responseBody, + responsePayloadFormat: retryParsed.responsePayloadFormat, + looksLikeSSE: retryParsed.looksLikeSSE, + requestBody: (retryResult.transformedBody || + input.translatedBody || + input.sourceBody) as Record, + transformedBody: retryResult.transformedBody, + model: currentModel, + connectionId: input.getCurrentConnectionId?.() ?? connectionId, + headers: new Headers(retryResult.headers), + startMs, + startedAt, + sourceFormat, + targetFormat, + clientResponseFormat, + provider, + upstreamResponse: retryResult.response, + requestHeaders: retryResult.headers, + requestUrl: retryResult.url, + }); + } + } + } + } catch { + // retry failed, fall through + } + } + if (envError) { + const receipt = buildReceipt(input, { + httpStatus: 502, + errorType: "clinepass_envelope_error", + usage: null, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId, + model: currentModel, + }); + const errorResult = legError( + 502, + envError.message, + envError, + null, + "clinepass_envelope_error", + "clinepass_envelope_error" + ); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + if (isJsonRecord(unwrapped)) { + responseBody = unwrapped as Record; + } + } + responseBody = unwrapClineNonStreamingEnvelope(provider, responseBody) as typeof responseBody; + + // -- Empty content -> family fallback (initial only) ------------------------- + if (isEmptyContentResponse(responseBody)) { + const errMsg = "Provider returned empty content"; + if (allowModelFallback) { + const triedModels = new Set([currentModel]); + const nextModel = getNextFamilyFallback(currentModel, triedModels, provider); + if (nextModel) { + triedModels.add(nextModel); + input.setRequestWireState({ + translatedBody: { ...input.translatedBody, model: nextModel } as Record, + effectiveModel: nextModel, + }); + log?.info?.("EMPTY_CONTENT_FALLBACK", `${currentModel} empty -> trying ${nextModel}`); + try { + // Connection check before fallback executor + const mismatchBefore = checkConnectionIdentity( + input, + startMs, + startedAt, + nextModel, + "before_fallback" + ); + if (mismatchBefore) return mismatchBefore; + + const fallbackResult = await input.executeProviderRequest(nextModel, false); + + // Connection check after fallback executor + const mismatchAfter = checkConnectionIdentity( + input, + startMs, + startedAt, + nextModel, + "after_fallback" + ); + if (mismatchAfter) return mismatchAfter; + + if (fallbackResult.response.ok) { + const fallbackParsed = await parseNonStreamingResponseBody({ + providerResponse: fallbackResult.response, + upstreamStream: false, + providerHeaders: new Headers(fallbackResult.headers), + finalBody: fallbackResult.transformedBody as Record | null, + targetFormat, + model: nextModel, + log, + }); + if (fallbackParsed.kind !== "invalid_sse" && fallbackParsed.kind !== "invalid_json") { + responseBody = fallbackParsed.responseBody; + return finishOk(input, { + responseBody, + responsePayloadFormat: fallbackParsed.responsePayloadFormat, + looksLikeSSE: fallbackParsed.looksLikeSSE, + requestBody: (fallbackResult.transformedBody || + input.translatedBody || + input.sourceBody) as Record, + transformedBody: fallbackResult.transformedBody, + model: nextModel, + connectionId: input.getCurrentConnectionId?.() ?? connectionId, + headers: new Headers(fallbackResult.headers), + startMs, + startedAt, + sourceFormat, + targetFormat, + clientResponseFormat, + provider, + upstreamResponse: fallbackResult.response, + requestHeaders: fallbackResult.headers, + requestUrl: fallbackResult.url, + }); + } + const parseCode = + fallbackParsed.kind === "invalid_sse" + ? "invalid_sse_payload" + : "invalid_json_payload"; + const parseMessage = fallbackParsed.message; + const receipt = buildReceipt(input, { + httpStatus: 502, + errorType: parseCode, + usage: null, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId: input.getCurrentConnectionId?.() ?? connectionId, + model: nextModel, + }); + const errorResult = legError( + 502, + parseMessage, + new Error(parseMessage), + null, + parseCode, + parseCode + ); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } else { + const receipt = buildReceipt(input, { + httpStatus: 502, + errorType: "empty_content", + usage: null, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId: input.getCurrentConnectionId?.() ?? connectionId, + model: nextModel, + }); + const errorResult = legError(502, errMsg, new Error(errMsg)); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + } catch { + const receipt = buildReceipt(input, { + httpStatus: 502, + errorType: "empty_content", + usage: null, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId, + model: currentModel, + }); + const errorResult = legError(502, errMsg, new Error(errMsg)); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + } else { + const receipt = buildReceipt(input, { + httpStatus: 502, + errorType: "empty_content", + usage: null, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId, + model: currentModel, + }); + const errorResult = legError(502, errMsg, new Error(errMsg)); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + } else { + const receipt = buildReceipt(input, { + httpStatus: 502, + errorType: "empty_content", + usage: null, + termination: "provider_error", + latencyMs: Date.now() - startMs, + startedAt, + endedAt: new Date().toISOString(), + connectionId, + model: currentModel, + }); + const errorResult = legError(502, errMsg, new Error(errMsg)); + return { + kind: "error", + result: errorResult as ChatCoreErrorResult, + receipt, + usage: null, + }; + } + } + + return finishOk(input, { + responseBody, + responsePayloadFormat, + looksLikeSSE, + requestBody: (finalBody || input.translatedBody || input.sourceBody) as Record, + transformedBody: finalBody, + model: currentModel, + connectionId, + headers: new Headers(executorResult.headers), + startMs, + startedAt, + sourceFormat, + targetFormat, + clientResponseFormat, + provider, + upstreamResponse: executorResult.response, + requestHeaders: executorResult.headers, + requestUrl: executorResult.url, + }); +} diff --git a/open-sse/handlers/chatCore/providerExecutionPipeline.ts b/open-sse/handlers/chatCore/providerExecutionPipeline.ts new file mode 100644 index 0000000000..077ff7cfe1 --- /dev/null +++ b/open-sse/handlers/chatCore/providerExecutionPipeline.ts @@ -0,0 +1,449 @@ +import type { ChatCoreErrorResult, ProviderLegUsage } from "@/lib/skills/toolLoopTypes.ts"; +import type { getProviderCredentials } from "@/sse/services/auth.ts"; +import type { updateFromHeaders, updateFromResponseBody } from "../../services/rateLimitManager.ts"; +import type { writeTerminalStatus } from "@/shared/utils/terminalStatus.ts"; +import type { updateProviderConnection } from "@/lib/db/providers.ts"; +import type { lockModel, recordCoreOwnedAntigravityQuotaState } from "../../services/accountFallback.ts"; +import { createErrorResult } from "../../utils/error.ts"; +import { applyStatusRestatement } from "../../config/upstreamStatusRestatement.ts"; +import { recoverAnthropicThinkingSignature } from "./thinkingSignatureRecovery.ts"; +import { isModelUnavailableError, getNextFamilyFallback as defaultGetNextFamilyFallback } from "../../services/modelFamilyFallback.ts"; +import { COOLDOWN_MS } from "../../config/errorConfig.ts"; +import { normalizeHeaders } from "../../utils/headers.ts"; + +export interface ChatCoreExecutorResult { + response: Response; + url: string; + headers: Record; + transformedBody: unknown; + transport?: string; + _executionCredentials?: Record; + _accountSemaphoreRelease?: () => void; +} + +export interface ProviderExecutionPolicy { + allowAccountRotation: boolean; + allowModelFallback: boolean; + expectedConnectionId?: string; +} + +export type ProviderExecutionOutcome = + | { + kind: "response"; + response: Response; + url: string; + headers: Record; + transformedBody: unknown; + model: string; + connectionId: string; + } + | { + kind: "error"; + result: ChatCoreErrorResult; + providerUsage: ProviderLegUsage | null; + model: string; + connectionId: string; + }; + +export interface PipelineTargetContext { + provider: string; + requestedModel: string; + sourceFormat: string; + targetFormat: string; + stream: boolean; +} + +export interface PipelineConnectionContext { + initialConnectionId: string; + getCurrentConnectionId: () => string | undefined; + getCredentials: () => Record; + replaceCredentials: (next: Record) => void; + onCredentialsRefreshed: (next: Record) => void | Promise; + assertManagedLeaseFence: (connectionId: string) => void; + getProviderCredentials: typeof getProviderCredentials; + refreshCredentials?: ( + credentials: Record + ) => Promise | null>; +} + +export interface PipelineWireState { + body: Record; + currentModel: string; + triedModels: Set; + setBodyAndModel: (body: Record, model: string) => void; +} + +export interface PipelineStateHooks { + updatePendingStage: (stage: string, data?: Record) => void; + recordRateLimitHeaders: typeof updateFromHeaders; + recordRateLimitBody: typeof updateFromResponseBody; + writeTerminalStatus: typeof writeTerminalStatus; + persistConnectionPatch: typeof updateProviderConnection; + setConnectionRateLimitedUntil: ( + connectionId: string, + untilMs: number | null + ) => void | Promise; + lockModel: typeof lockModel; + recordAntigravityQuotaState: typeof recordCoreOwnedAntigravityQuotaState; + markAccountSemaphoreBlocked: (connectionId: string) => void; + isolateProbeFailures: () => boolean | Promise; + onCodexScopeRateLimited?: (params: { + failedConnectionId: string; + model: string | null; + rateLimitedUntil: string; + credentials?: Record | null; + }) => void | Promise; + onClearSessionAffinity?: (params: { failedConnectionId: string }) => void | Promise; + onAuditAccountRotation?: (params: { + action: "codex.account_rotation"; + failedConnectionId: string; + newConnectionId: string; + attempt: number; + retryAfterMs: number | null; + }) => void | Promise; +} + +export interface ProviderExecutionPipelineInput { + policy: Readonly; + target: PipelineTargetContext; + connection: PipelineConnectionContext; + wire: PipelineWireState; + state: PipelineStateHooks; + sendProviderAttempt: (model: string, allowDedup: boolean) => Promise; + getNextFamilyFallback?: ( + currentModel: string, + triedModels: Set, + providerHint?: string | null + ) => string | null; +} + +const LEASE_MISMATCH_STATUS = 409; +const LEASE_MISMATCH_CODE = "LEASE_CONNECTION_MISMATCH"; + +function currentConnectionId(connection: PipelineConnectionContext): string { + return connection.getCurrentConnectionId() ?? connection.initialConnectionId; +} + +function retryAfterMsFrom(attempt: ChatCoreExecutorResult): number | null { + // attempt.headers is the outbound request bag (BaseExecutor finalHeaders). + // Retry-After lives on the upstream Response — same source as the parent + // chatCore rotate path. normalizeHeaders lower-cases keys, so "Retry-After" + // is looked up as "retry-after"; it does not drop the field. + const raw = normalizeHeaders(attempt.response?.headers)["retry-after"]; + if (raw == null || raw === "") return null; + const parsed = Number.parseFloat(String(raw)); + if (!Number.isFinite(parsed) || parsed < 0) return null; + return parsed * 1000; +} + +function leaseMismatch(model: string, connectionId: string): ProviderExecutionOutcome { + const result = createErrorResult( + LEASE_MISMATCH_STATUS, + "Managed lease connection mismatch", + null, + LEASE_MISMATCH_CODE, + "lease_error" + ); + return { + kind: "error", + result: { + success: false, + status: result.status, + response: result.response, + error: result.error, + errorCode: LEASE_MISMATCH_CODE, + errorType: "lease_error", + }, + providerUsage: null, + model, + connectionId, + }; +} + +async function toOutcome( + attempt: ChatCoreExecutorResult, + model: string, + connectionId: string, + provider: string +): Promise { + const status = attempt.response.status; + if (status >= 200 && status < 300) { + return { + kind: "response", + response: attempt.response, + url: attempt.url, + headers: attempt.headers, + transformedBody: attempt.transformedBody, + model, + connectionId, + }; + } + let message = attempt.response.statusText || "upstream error"; + let body: unknown = attempt.transformedBody; + try { + // clone() is the drain. sendProviderAttempt must not cancel() a streaming + // non-2xx body before we get here (BYOP 422 / Codex 429 Retry-After). + body = JSON.parse(await attempt.response.clone().text()); + const err = (body as { error?: { message?: unknown } } | null)?.error; + if (err && typeof err.message === "string" && err.message) message = err.message; + } catch { + // keep statusText + } + const restatement = applyStatusRestatement({ + provider, + status, + message, + body, + retryAfterMs: null, + }); + const result = createErrorResult( + restatement.status, + message, + restatement.retryAfterMs + ); + return { + kind: "error", + result: { + success: false, + status: result.status, + response: attempt.response, + error: result.error, + errorCode: result.errorCode, + errorType: result.errorType, + }, + providerUsage: null, + model, + connectionId, + }; +} + +function assertLease( + policy: Readonly, + connection: PipelineConnectionContext, + model: string +): ProviderExecutionOutcome | null { + const expected = policy.expectedConnectionId; + if (!expected) return null; + const current = connection.getCurrentConnectionId(); + if (current && current !== expected) { + return leaseMismatch(model, current); + } + return null; +} + +function maxAttemptsFor(provider: string): number { + return provider === "codex" ? 3 : 1; +} + +/** + * Shared first-send + provider recovery. Does not read a successful body. + * Account/model retries live here; sendProviderAttempt is one wire send. + */ +export async function runProviderExecutionPipeline( + input: ProviderExecutionPipelineInput +): Promise { + const { policy, target, connection, wire, state, sendProviderAttempt } = input; + const maxAttempts = maxAttemptsFor(target.provider); + const excludedIds: string[] = []; + let attempts = 0; + let lastAttempt: ChatCoreExecutorResult | null = null; + let antigravityByopRotationPending = false; + let authRefreshPending = false; + let authRefreshed = false; + let modelFallbackPending = false; + const resolveFamilyFallback = input.getNextFamilyFallback ?? defaultGetNextFamilyFallback; + + while ( + attempts < maxAttempts || + antigravityByopRotationPending || + authRefreshPending || + modelFallbackPending + ) { + antigravityByopRotationPending = false; + authRefreshPending = false; + modelFallbackPending = false; + const before = assertLease(policy, connection, wire.currentModel); + if (before) return before; + + const attempt = await sendProviderAttempt(wire.currentModel, attempts === 0); + lastAttempt = attempt; + + const after = assertLease(policy, connection, wire.currentModel); + if (after) return after; + + const status = attempt.response.status; + if (status >= 200 && status < 300) { + return toOutcome(attempt, wire.currentModel, currentConnectionId(connection), target.provider); + } + + const isolateProbe = await state.isolateProbeFailures(); + const canRotateAccount = policy.allowAccountRotation && !isolateProbe; + + if ( + canRotateAccount && + target.provider === "codex" && + status === 429 && + attempts < maxAttempts - 1 + ) { + const failedId = currentConnectionId(connection); + const retryAfterMs = retryAfterMsFrom(attempt); + if (failedId && !excludedIds.includes(failedId)) excludedIds.push(failedId); + if (failedId) { + await state.onCodexScopeRateLimited?.({ + failedConnectionId: failedId, + model: wire.currentModel || target.requestedModel || null, + rateLimitedUntil: new Date(Date.now() + (retryAfterMs || 60_000)).toISOString(), + credentials: connection.getCredentials(), + }); + await state.onClearSessionAffinity?.({ failedConnectionId: failedId }); + } + const nextCreds = await connection + .getProviderCredentials("codex", null, null, wire.currentModel, { + excludeConnectionIds: [...excludedIds], + }) + .catch(() => null); + if (nextCreds && !nextCreds.allRateLimited && nextCreds.connectionId) { + await state.onAuditAccountRotation?.({ + action: "codex.account_rotation", + failedConnectionId: failedId, + newConnectionId: String(nextCreds.connectionId), + attempt: attempts + 1, + retryAfterMs, + }); + connection.replaceCredentials(nextCreds as Record); + attempts += 1; + continue; + } + } + + if (canRotateAccount && target.provider === "antigravity" && status === 422) { + // Same drain as toOutcome: clone the Response. A prior body.cancel() + // makes this throw "Body has already been consumed" and skips rotate. + const byopBody = await attempt.response + .clone() + .text() + .catch(() => ""); + if (byopBody.includes("gcp_project_required")) { + const failedId = currentConnectionId(connection); + if (failedId && !excludedIds.includes(failedId)) excludedIds.push(failedId); + if (failedId) { + await state.setConnectionRateLimitedUntil( + failedId, + Date.now() + (COOLDOWN_MS.gcpProjectRequired ?? 24 * 60 * 60 * 1000) + ); + } + const nextCreds = await connection + .getProviderCredentials("antigravity", null, null, wire.currentModel, { + excludeConnectionIds: [...excludedIds], + }) + .catch(() => null); + if (nextCreds && !nextCreds.allRateLimited && nextCreds.connectionId) { + connection.replaceCredentials(nextCreds as Record); + antigravityByopRotationPending = true; + continue; + } + } + } + + if ( + !authRefreshed && + (status === 401 || status === 403) && + typeof connection.refreshCredentials === "function" + ) { + const refreshed = await connection.refreshCredentials(connection.getCredentials()); + if (refreshed && (refreshed.accessToken || refreshed.copilotToken)) { + connection.replaceCredentials({ ...connection.getCredentials(), ...refreshed }); + await connection.onCredentialsRefreshed(refreshed); + authRefreshed = true; + authRefreshPending = true; + continue; + } + } + + { + let signatureMessage = attempt.response.statusText || "upstream error"; + try { + const parsed = JSON.parse(await attempt.response.clone().text()) as { + error?: { message?: unknown }; + }; + if (typeof parsed?.error?.message === "string" && parsed.error.message) { + signatureMessage = parsed.error.message; + } + } catch { + // keep statusText + } + const signatureRecovery = await recoverAnthropicThinkingSignature({ + provider: target.provider, + statusCode: status, + message: signatureMessage, + body: wire.body, + execute: async (recoveryBody) => { + if (recoveryBody && typeof recoveryBody === "object" && !Array.isArray(recoveryBody)) { + wire.setBodyAndModel(recoveryBody as Record, wire.currentModel); + } + return sendProviderAttempt(wire.currentModel, false); + }, + parseError: async (response) => { + let message = response.statusText || "upstream error"; + let responseBody: unknown = null; + try { + responseBody = JSON.parse(await response.clone().text()); + const err = (responseBody as { error?: { message?: unknown } } | null)?.error; + if (typeof err?.message === "string" && err.message) message = err.message; + } catch { + // keep statusText + } + return { + statusCode: response.status, + message, + retryAfterMs: null, + responseBody, + }; + }, + }); + if (signatureRecovery.attempted && signatureRecovery.succeeded && signatureRecovery.execution) { + lastAttempt = { + response: signatureRecovery.execution.response, + url: signatureRecovery.execution.url ?? attempt.url, + headers: (signatureRecovery.execution.headers as Record) ?? attempt.headers, + transformedBody: signatureRecovery.execution.transformedBody ?? attempt.transformedBody, + }; + return toOutcome( + lastAttempt, + wire.currentModel, + currentConnectionId(connection), + target.provider + ); + } + } + + if (policy.allowModelFallback) { + let fallbackMessage = attempt.response.statusText || "upstream error"; + try { + const parsed = JSON.parse(await attempt.response.clone().text()) as { + error?: { message?: unknown }; + }; + if (typeof parsed?.error?.message === "string" && parsed.error.message) { + fallbackMessage = parsed.error.message; + } + } catch { + // keep statusText + } + if (isModelUnavailableError(status, fallbackMessage, target.provider)) { + const nextModel = resolveFamilyFallback(wire.currentModel, wire.triedModels, target.provider); + if (nextModel) { + wire.setBodyAndModel({ ...wire.body, model: nextModel }, nextModel); + modelFallbackPending = true; + continue; + } + } + } + + return toOutcome(attempt, wire.currentModel, currentConnectionId(connection), target.provider); + } + + if (lastAttempt) { + return toOutcome(lastAttempt, wire.currentModel, currentConnectionId(connection), target.provider); + } + return leaseMismatch(wire.currentModel, currentConnectionId(connection)); +} diff --git a/open-sse/handlers/chatCore/serverOwnedToolLoopGate.ts b/open-sse/handlers/chatCore/serverOwnedToolLoopGate.ts new file mode 100644 index 0000000000..29b2a6925b --- /dev/null +++ b/open-sse/handlers/chatCore/serverOwnedToolLoopGate.ts @@ -0,0 +1,15 @@ +import { FORMATS } from "../../translator/formats.ts"; + +export function shouldRunServerOwnedToolLoop(input: { + enabled: boolean; + stream: boolean; + isResponsesEndpoint: boolean; + sourceFormat: string; +}): boolean { + if (!input.enabled) return false; + if (input.stream) return false; + if (input.isResponsesEndpoint) return false; + if (input.sourceFormat === FORMATS.OPENAI) return true; + if (input.sourceFormat === FORMATS.CLAUDE) return true; + return false; +} diff --git a/open-sse/handlers/chatCore/serverOwnedToolLoopWire.ts b/open-sse/handlers/chatCore/serverOwnedToolLoopWire.ts new file mode 100644 index 0000000000..e79cdcff26 --- /dev/null +++ b/open-sse/handlers/chatCore/serverOwnedToolLoopWire.ts @@ -0,0 +1,175 @@ +import type { + ExecutionContext, + NonStreamingProviderLegResult, + ProviderLegUsage, + ServerOwnedToolLoopResult, + ToolCall, +} from "@/lib/skills/toolLoopTypes.ts"; +import { executeServerOwned } from "@/lib/skills/interception"; +import { runServerOwnedToolLoop, LOOP_BUDGET_MS } from "@/lib/skills/serverOwnedToolLoop.ts"; +import { deriveToolRequestIdentity } from "@/lib/skills/stableJson.ts"; +import { getIdempotencyKey } from "@/lib/idempotencyLayer"; +import { runNonStreamingProviderLeg } from "./nonStreamingProviderLeg.ts"; +import type { ProviderLegInput } from "./nonStreamingProviderLeg.ts"; +import { shouldRunServerOwnedToolLoop } from "./serverOwnedToolLoopGate.ts"; +import { FORMATS } from "../../translator/formats.ts"; + +export function derivePostInjectionRequestIdentity(input: { + apiKeyId: string; + headers: unknown; + skillRequestId: string; + postInjectionBody: Record; +}): string { + const stableClientRequestId = getIdempotencyKey(input.headers as never); + return deriveToolRequestIdentity({ + apiKeyId: input.apiKeyId, + stableClientRequestId, + skillRequestId: input.skillRequestId, + postInjectionBody: input.postInjectionBody, + }); +} + +export async function continueServerOwnedToolLoop(input: { + initialLeg: NonStreamingProviderLegResult & { kind: "ok" }; + sourceBody: Record; + sourceFormat: "openai" | "claude"; + skillsModelId: string; + executionContext: ExecutionContext; + abortSignal?: AbortSignal; + deadlineAtMs: number; + expectedConnectionId?: string; + followUpLeg: (nextSourceBody: Record) => Promise; + executeServerOwned?: ( + calls: ToolCall[], + context: ExecutionContext + ) => Promise; +}): Promise { + const runOwned = input.executeServerOwned ?? executeServerOwned; + return runServerOwnedToolLoop({ + initialLeg: input.initialLeg, + sourceBody: input.sourceBody, + sourceFormat: input.sourceFormat, + skillsModelId: input.skillsModelId, + executionContext: input.executionContext, + abortSignal: input.abortSignal, + deadlineAtMs: input.deadlineAtMs, + executeServerOwned: (calls: ToolCall[], context: ExecutionContext) => runOwned(calls, context), + resumeUpstream: async (nextSourceBody, expectedConnectionId) => { + if ( + expectedConnectionId && + input.expectedConnectionId && + expectedConnectionId !== input.expectedConnectionId + ) { + return { + kind: "error", + result: { + success: false, + status: 409, + response: new Response(null, { status: 409 }), + error: "Follow-up connection mismatch", + errorCode: "LEASE_CONNECTION_MISMATCH", + }, + receipt: input.initialLeg.receipt, + usage: null, + }; + } + return input.followUpLeg(nextSourceBody); + }, + }); +} + +export function followUpLegInput( + base: Omit< + ProviderLegInput, + "phase" | "allowAccountRotation" | "allowModelFallback" | "sourceBody" + >, + nextSourceBody: Record, + expectedConnectionId?: string +): ProviderLegInput { + return { + ...base, + phase: "follow-up", + sourceBody: nextSourceBody, + expectedConnectionId, + allowAccountRotation: false, + allowModelFallback: false, + }; +} + +export function mergeLoopIntoOkLeg( + leg: NonStreamingProviderLegResult & { kind: "ok" }, + loop: ServerOwnedToolLoopResult +): NonStreamingProviderLegResult & { kind: "ok" } { + return { + ...leg, + response: loop.response ?? leg.response, + responseForMemoryExtraction: + loop.responseForMemoryExtraction ?? leg.responseForMemoryExtraction, + providerBody: loop.finalProviderBody ?? leg.providerBody, + providerRequest: loop.finalProviderRequest ?? leg.providerRequest, + usage: loop.cumulativeUsage, + }; +} + +export type ToolLoopApplyResult = + | { kind: "skip" } + | { + kind: "ok"; + leg: NonStreamingProviderLegResult & { kind: "ok" }; + usage: ProviderLegUsage | null; + loop: ServerOwnedToolLoopResult; + } + | { kind: "error"; loop: ServerOwnedToolLoopResult }; + +export async function applyServerOwnedToolLoopIfNeeded(input: { + enabled: boolean; + stream: boolean; + isResponsesEndpoint: boolean; + sourceFormat: string; + initialLeg: NonStreamingProviderLegResult; + sourceBody: Record; + skillsModelId: string; + executionContext: ExecutionContext; + abortSignal?: AbortSignal; + expectedConnectionId?: string; + followUpLeg: (nextSourceBody: Record) => Promise; + logReceipt: (receipt: ServerOwnedToolLoopResult["receipts"][number]) => void; + executeServerOwned?: ( + calls: ToolCall[], + context: ExecutionContext + ) => Promise; +}): Promise { + if ( + input.initialLeg.kind !== "ok" || + !shouldRunServerOwnedToolLoop({ + enabled: input.enabled, + stream: input.stream, + isResponsesEndpoint: input.isResponsesEndpoint, + sourceFormat: input.sourceFormat, + }) + ) { + return { kind: "skip" }; + } + const loop = await continueServerOwnedToolLoop({ + initialLeg: input.initialLeg, + sourceBody: input.sourceBody, + sourceFormat: input.sourceFormat === FORMATS.CLAUDE ? "claude" : "openai", + skillsModelId: input.skillsModelId, + executionContext: input.executionContext, + abortSignal: input.abortSignal, + deadlineAtMs: Date.now() + LOOP_BUDGET_MS, + expectedConnectionId: input.expectedConnectionId, + followUpLeg: input.followUpLeg, + executeServerOwned: input.executeServerOwned, + }); + for (const receipt of loop.receipts) input.logReceipt(receipt); + if (loop.kind === "error") return { kind: "error", loop }; + return { + kind: "ok", + leg: mergeLoopIntoOkLeg(input.initialLeg, loop), + usage: loop.cumulativeUsage, + loop, + }; +} + +export { LOOP_BUDGET_MS, runNonStreamingProviderLeg }; diff --git a/open-sse/utils/requestLogger.ts b/open-sse/utils/requestLogger.ts index f6722ac674..f8e715ed00 100644 --- a/open-sse/utils/requestLogger.ts +++ b/open-sse/utils/requestLogger.ts @@ -19,6 +19,7 @@ export type RequestPipelinePayloads = { providerResponse?: JsonRecord; clientResponse?: JsonRecord; error?: JsonRecord; + toolLoop?: { legs: JsonRecord[] }; streamChunks?: { provider?: string[]; openai?: string[]; @@ -48,6 +49,7 @@ type RequestLogger = { logConvertedResponse: (body: unknown) => void; appendConvertedChunk: (chunk: string) => void; logError: (error: unknown, requestBody?: unknown) => void; + logToolLoopReceipt: (receipt: unknown) => void; getPipelinePayloads: () => RequestPipelinePayloads | null; }; @@ -74,6 +76,7 @@ const MAX_LOG_STRING_LENGTH = 64 * 1024; // existing plain-constant shape; CHAT_LOG_ARRAY_TAIL_ITEMS still overrides it. export const MAX_LOG_ARRAY_ITEMS = getChatLogArrayTailItems(); const MAX_LOG_OBJECT_KEYS = 80; +const MAX_TOOL_LOOP_LEGS = 4; function maskSensitiveHeaders(headers: HeaderInput): Record { if (!headers) return {}; @@ -282,7 +285,16 @@ function compactPipelinePayloads( continue; } - result[key as keyof RequestPipelinePayloads] = value; + if (key === "toolLoop" && value && typeof value === "object") { + const legs = (value as { legs?: unknown }).legs; + if (Array.isArray(legs) && legs.length > 0) { + result.toolLoop = { legs: legs as JsonRecord[] }; + } + continue; + } + + const payloadKey = key as Exclude; + result[payloadKey] = value as JsonRecord; } return hasOwnValues(result) ? result : null; @@ -384,6 +396,7 @@ export async function createRequestLogger( logConvertedResponse() {}, appendConvertedChunk: chunkMethods.appendConvertedChunk, logError() {}, + logToolLoopReceipt() {}, getPipelinePayloads() { return routeDecision ? { routeDecision } : null; }, @@ -468,6 +481,14 @@ export async function createRequestLogger( }; }, + logToolLoopReceipt(receipt) { + const legs = payloads.toolLoop?.legs ?? []; + if (legs.length >= MAX_TOOL_LOOP_LEGS) return; + const cloned = cloneBoundedForLog(receipt); + if (!cloned || typeof cloned !== "object" || Array.isArray(cloned)) return; + payloads.toolLoop = { legs: [...legs, cloned as JsonRecord] }; + }, + getPipelinePayloads() { return compactPipelinePayloads(payloads); }, diff --git a/src/i18n/messages/en.json b/src/i18n/messages/en.json index ab2feb8761..61429ab690 100644 --- a/src/i18n/messages/en.json +++ b/src/i18n/messages/en.json @@ -13006,6 +13006,10 @@ }, "SKILLS_SANDBOX_NETWORK_ENABLED": { "description": "Enable network access in the skills sandbox." + }, + "SERVER_OWNED_TOOL_LOOP_ENABLED": { + "label": "Server-Owned Tool Loop", + "description": "Continue non-streaming server-owned tool calls until the model returns a client-usable response." } } }, diff --git a/src/i18n/messages/pt-BR.json b/src/i18n/messages/pt-BR.json index ad64d3f9a5..7e76593cf7 100644 --- a/src/i18n/messages/pt-BR.json +++ b/src/i18n/messages/pt-BR.json @@ -13007,6 +13007,10 @@ }, "SKILLS_SANDBOX_NETWORK_ENABLED": { "description": "Ativa o acesso à rede no sandbox de skills." + }, + "SERVER_OWNED_TOOL_LOOP_ENABLED": { + "label": "Server-Owned Tool Loop", + "description": "Continue chamadas de ferramentas do servidor (server-owned) em não-streaming até que o modelo retorne uma resposta utilizável pelo cliente." } } }, diff --git a/src/i18n/messages/vi.json b/src/i18n/messages/vi.json index adc5ac0fe1..59c7fa5123 100644 --- a/src/i18n/messages/vi.json +++ b/src/i18n/messages/vi.json @@ -13007,6 +13007,10 @@ }, "SKILLS_SANDBOX_NETWORK_ENABLED": { "description": "Cho phép môi trường sandbox của kỹ năng truy cập mạng." + }, + "SERVER_OWNED_TOOL_LOOP_ENABLED": { + "label": "Vòng lặp công cụ do máy chủ sở hữu", + "description": "Tiếp tục các lời gọi công cụ do máy chủ sở hữu ở chế độ không streaming cho đến khi mô hình trả về phản hồi mà máy khách dùng được." } } }, diff --git a/src/lib/db/migrations/174_server_tool_executions.sql b/src/lib/db/migrations/174_server_tool_executions.sql new file mode 100644 index 0000000000..fb371e0ef9 --- /dev/null +++ b/src/lib/db/migrations/174_server_tool_executions.sql @@ -0,0 +1,25 @@ +-- Migration 174: Durable server tool execution fence table. +-- Tracks claim/result state for server-owned tool calls to prevent duplicate execution +-- across retries and concurrent requests. Independent of skill_executions. + +CREATE TABLE IF NOT EXISTS server_tool_executions ( + id TEXT PRIMARY KEY, + api_key_id TEXT NOT NULL, + request_identity TEXT NOT NULL, + tool_call_id TEXT NOT NULL, + tool_name TEXT NOT NULL, + input_digest TEXT NOT NULL, + output TEXT, + status TEXT NOT NULL CHECK(status IN ('running', 'success', 'error', 'timeout')), + error_message TEXT, + duration_ms INTEGER, + claim_expires_at TEXT NOT NULL, + completed_at TEXT, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + UNIQUE(api_key_id, request_identity, tool_call_id) +); + +CREATE INDEX IF NOT EXISTS idx_server_tool_executions_status_expiry + ON server_tool_executions(status, claim_expires_at); +CREATE INDEX IF NOT EXISTS idx_server_tool_executions_created + ON server_tool_executions(created_at); diff --git a/src/lib/db/skillExecutionFence.ts b/src/lib/db/skillExecutionFence.ts new file mode 100644 index 0000000000..ed399dbd05 --- /dev/null +++ b/src/lib/db/skillExecutionFence.ts @@ -0,0 +1,241 @@ +import { randomUUID } from "node:crypto"; +import type { SqliteAdapter } from "./adapters/types"; +import { getDbInstance } from "./core"; +import { + sanitizeErrorMessage, + sanitizeUpstreamDetails, +} from "@omniroute/open-sse/utils/errorSanitization"; + +const MAX_PERSISTED_OUTPUT_CHARS = 32_768; +const MAX_PERSISTED_ERROR_CHARS = 4_096; + +export type ServerToolClaim = + | { kind: "claimed"; executionId: string } + | { + kind: "replay"; + executionId: string; + status: "success" | "error" | "timeout"; + output: unknown; + errorMessage: string | null; + } + | { kind: "in_progress"; executionId: string } + | { kind: "unknown"; executionId: string } + | { kind: "identity_conflict"; executionId: string }; + +function isUniqueConstraintError(err: unknown): boolean { + if (!err || typeof err !== "object") return false; + const code = String((err as { code?: unknown }).code ?? ""); + const msg = String((err as { message?: unknown }).message ?? ""); + // Accept only extended UNIQUE/PRIMARYKEY codes — NOT generic SQLITE_CONSTRAINT + // (NOT NULL/CHECK constraints must NOT be swallowed) + if (/SQLITE_CONSTRAINT_PRIMARYKEY/i.test(code)) return true; + if (/SQLITE_CONSTRAINT_UNIQUE/i.test(code)) return true; + // Fallback: message must explicitly reference UNIQUE (not just CONSTRAINT) + if (/UNIQUE constraint failed/i.test(msg)) return true; + if (/UNIQUE constraint violation/i.test(msg)) return true; + return false; +} + +function sanitizeAndBounded(value: unknown, maxChars: number): string | null { + if (value === null || value === undefined) return null; + // Strings: sanitize credentials/stack, store as bounded plain text (error_message is TEXT) + if (typeof value === "string") { + const sanitized = sanitizeErrorMessage(value); + return sanitized.length > maxChars ? sanitized.slice(0, maxChars) : sanitized; + } + // Objects/arrays: recursive sanitization, JSON.stringify, then guaranteed-valid truncation + const sanitized = sanitizeUpstreamDetails(value); + let text: string; + try { + text = JSON.stringify(sanitized); + } catch { + text = JSON.stringify({ error: "Value is not JSON-serializable" }); + } + if (text.length <= maxChars) return text; + // Truncation envelope: progressively shorten preview until serialized envelope fits + const envelope = { truncated: true, preview: "" as string }; + for (let len = maxChars - 40; len > 0; len -= 20) { + envelope.preview = text.slice(0, len); + const candidate = JSON.stringify(envelope); + if (candidate.length <= maxChars) return candidate; + } + envelope.preview = text.slice(0, 20); + return JSON.stringify(envelope); +} + +export function readRow( + db: SqliteAdapter, + executionId: string +): { + id: string; + tool_name: string; + input_digest: string; + status: string; + output: string | null; + error_message: string | null; +} | null { + const row = db + .prepare( + "SELECT id, tool_name, input_digest, status, output, error_message FROM server_tool_executions WHERE id = ?" + ) + .get(executionId) as + | { + id: string; + tool_name: string; + input_digest: string; + status: string; + output: string | null; + error_message: string | null; + } + | undefined; + return row ?? null; +} + +function readExistingByIdentity( + db: SqliteAdapter, + apiKeyId: string, + requestIdentity: string, + toolCallId: string +): (ReturnType & { claim_expires_at: string }) | null { + const row = db + .prepare( + `SELECT id, tool_name, input_digest, status, output, error_message, claim_expires_at + FROM server_tool_executions + WHERE api_key_id = ? AND request_identity = ? AND tool_call_id = ?` + ) + .get(apiKeyId, requestIdentity, toolCallId) as + | { + id: string; + tool_name: string; + input_digest: string; + status: string; + output: string | null; + error_message: string | null; + claim_expires_at: string; + } + | undefined; + return row ?? null; +} + +function buildReplayClaim(row: { + id: string; + status: string; + output: string | null; + error_message: string | null; +}): ServerToolClaim { + let parsedOutput: unknown = null; + if (row.output !== null) { + try { + parsedOutput = JSON.parse(row.output); + } catch { + parsedOutput = row.output; + } + } + return { + kind: "replay", + executionId: row.id, + status: row.status as "success" | "error" | "timeout", + output: parsedOutput, + errorMessage: row.error_message, + }; +} + +export function claimServerToolExecution( + input: { + apiKeyId: string; + requestIdentity: string; + toolCallId: string; + toolName: string; + inputDigest: string; + leaseExpiresAt: string; + }, + db: SqliteAdapter = getDbInstance(), + nowMs?: number +): ServerToolClaim { + const executionId = randomUUID(); + const now = nowMs ?? Date.now(); + + // Try INSERT in a transaction + const tryInsert = db.transaction(() => { + db.prepare( + `INSERT INTO server_tool_executions + (id, api_key_id, request_identity, tool_call_id, tool_name, input_digest, status, claim_expires_at) + VALUES (?, ?, ?, ?, ?, ?, 'running', ?)` + ).run( + executionId, + input.apiKeyId, + input.requestIdentity, + input.toolCallId, + input.toolName, + input.inputDigest, + input.leaseExpiresAt + ); + }); + + try { + tryInsert(); + return { kind: "claimed", executionId }; + } catch (err: unknown) { + if (!isUniqueConstraintError(err)) { + throw err; + } + } + + // UNIQUE conflict — re-read in a fresh transaction + const existing = db.transaction(() => { + return readExistingByIdentity(db, input.apiKeyId, input.requestIdentity, input.toolCallId); + })(); + + if (!existing) { + return { kind: "unknown", executionId }; + } + + // Identity conflict: different name or digest + if (existing.tool_name !== input.toolName || existing.input_digest !== input.inputDigest) { + return { kind: "identity_conflict", executionId: existing.id }; + } + + // Terminal status → replay + if ( + existing.status === "success" || + existing.status === "error" || + existing.status === "timeout" + ) { + return buildReplayClaim(existing); + } + + // Running status — check STORED lease expiry (not input) + const storedExpiresAt = new Date(existing.claim_expires_at).getTime(); + if (storedExpiresAt <= now) { + return { kind: "unknown", executionId: existing.id }; + } + + // Running + unexpired → in_progress + return { kind: "in_progress", executionId: existing.id }; +} + +export function finalizeServerToolExecution( + input: { + executionId: string; + status: "success" | "error" | "timeout"; + output: unknown | null; + errorMessage: string | null; + durationMs: number; + }, + db: SqliteAdapter = getDbInstance() +): boolean { + const safeOutput = sanitizeAndBounded(input.output, MAX_PERSISTED_OUTPUT_CHARS); + const safeError = sanitizeAndBounded(input.errorMessage, MAX_PERSISTED_ERROR_CHARS); + + const result = db.transaction(() => { + return db + .prepare( + `UPDATE server_tool_executions + SET status = ?, output = ?, error_message = ?, duration_ms = ?, completed_at = datetime('now') + WHERE id = ? AND status = 'running'` + ) + .run(input.status, safeOutput, safeError, input.durationMs, input.executionId); + })(); + + return result.changes > 0; +} diff --git a/src/lib/skills/executor.ts b/src/lib/skills/executor.ts index ac958f1a54..06cc96a646 100644 --- a/src/lib/skills/executor.ts +++ b/src/lib/skills/executor.ts @@ -218,6 +218,8 @@ class SkillExecutor { throw new Error(`Skill not found: ${skillName}`); } + // Check enabled/disabled BEFORE creating a DB row (preserves pre-Task-3 + // behavior: disabled/missing skills never write to skill_executions). if (!skill.enabled) { throw new Error(`Skill is disabled: ${skillName}`); } @@ -242,41 +244,10 @@ class SkillExecutor { new Date().toISOString() ); - let handler = this.handlers.get(skill.handler); - if (!handler) { - // Builtin handlers are registered by instrumentation-node at startup, - // but Next.js may compile this module into multiple chunks (each with - // its own SkillExecutor singleton). Fall back to the builtin registry - // so `POST /api/skills/executions` works regardless of which chunk the - // route is served from. - const builtin = builtinSkills[skill.handler]; - if (builtin) { - this.handlers.set(skill.handler, builtin); - handler = builtin; - } - } - if (!handler) { - throw new Error(`Handler not found: ${skill.handler}`); - } - - let output: Record | null = null; - let errorMessage: string | null = null; - let status = SkillStatus.SUCCESS; - - try { - const result = await this.executeWithTimeout( - handler(input, { apiKeyId: context.apiKeyId, sessionId: context.sessionId || "" }) - ); - const resultIsFailure = isSkillFailureOutput(result); - output = projectSkillOutputForBoundary(result); - if (resultIsFailure) { - errorMessage = skillFailureMessage(result); - status = SkillStatus.ERROR; - } - } catch (err) { - errorMessage = toSafeSkillErrorMessage(err); - status = SkillStatus.ERROR; - } + const { output, errorMessage, status } = await this.runHandler(skillName, input, { + apiKeyId: context.apiKeyId, + sessionId: context.sessionId || "", + }); const durationMs = Date.now() - startTime; @@ -323,6 +294,119 @@ class SkillExecutor { ]); } + /** + * Shared handler lookup + execute + output projection used by both + * `execute()` (which writes history) and `executeClaimed()` (which does not). + */ + private async runHandler( + skillName: string, + input: Record, + context: { apiKeyId: string; sessionId: string } + ): Promise<{ + output: Record | null; + errorMessage: string | null; + status: SkillStatus; + }> { + const skill = skillRegistry.getSkill(skillName, context.apiKeyId); + if (!skill) { + throw new Error(`Skill not found: ${skillName}`); + } + if (!skill.enabled) { + throw new Error(`Skill is disabled: ${skillName}`); + } + + let handler = this.handlers.get(skill.handler); + if (!handler) { + const builtin = builtinSkills[skill.handler]; + if (builtin) { + this.handlers.set(skill.handler, builtin); + handler = builtin; + } + } + if (!handler) { + throw new Error(`Handler not found: ${skill.handler}`); + } + + let output: Record | null = null; + let errorMessage: string | null = null; + let status = SkillStatus.SUCCESS; + + try { + const result = await this.executeWithTimeout( + handler(input, { apiKeyId: context.apiKeyId, sessionId: context.sessionId || "" }) + ); + const resultIsFailure = isSkillFailureOutput(result); + output = projectSkillOutputForBoundary(result); + if (resultIsFailure) { + errorMessage = skillFailureMessage(result); + status = SkillStatus.ERROR; + } + } catch (err) { + errorMessage = toSafeSkillErrorMessage(err); + status = SkillStatus.ERROR; + } + + return { output, errorMessage, status }; + } + + /** + * Execute a claimed skill call for the server-owned tool loop. + * Reuses the same handler lookup/timeout/projection as `execute()` + * but does NOT write to `skill_executions` — the fence table owns + * persistence for claimed executions. + */ + async executeClaimed( + skillName: string, + input: Record, + context: { apiKeyId: string; sessionId: string }, + executionId: string + ): Promise { + const settings = await getSettings(); + if (settings.skillsEnabled === false) { + throw new Error("Skills execution is disabled. Enable Skills in Settings > AI."); + } + + const skill = skillRegistry.getSkill(skillName, context.apiKeyId); + if (!skill) { + throw new Error(`Skill not found: ${skillName}`); + } + + const startTime = Date.now(); + log.info("skills.executor.claimed_start", { + skillId: skill.id, + skillName, + apiKeyId: context.apiKeyId, + executionId, + }); + + const { output, errorMessage, status } = await this.runHandler(skillName, input, context); + const durationMs = Date.now() - startTime; + + if (status !== SkillStatus.SUCCESS) { + throw new Error(`Skill execution failed: ${errorMessage ?? "unknown error"}`); + } + + log.info("skills.executor.claimed_complete", { + skillId: skill.id, + success: status === SkillStatus.SUCCESS, + durationMs, + executionId, + }); + + return { + id: executionId, + skillId: skill.id, + apiKeyId: context.apiKeyId, + sessionId: context.sessionId || "", + input, + output, + status, + errorMessage, + durationMs, + createdAt: new Date(), + }; + } + getExecution(executionId: string): SkillExecution | undefined { const db = getDbInstance(); const row = db.prepare("SELECT * FROM skill_executions WHERE id = ?").get(executionId) as any; diff --git a/src/lib/skills/followUpTranscript.ts b/src/lib/skills/followUpTranscript.ts new file mode 100644 index 0000000000..47bb8a97fd --- /dev/null +++ b/src/lib/skills/followUpTranscript.ts @@ -0,0 +1,412 @@ +import type { + BoundedToolResult, + BuildFollowUpTranscriptInput, + ExecutedToolResult, + ToolCall, +} from "./toolLoopTypes"; + +/** + * Pure transcript builder for the server-owned tool loop (spec §5.2). + * + * `serializeBoundedToolResult` serializes one executed tool result to JSON text + * and, when it exceeds a UTF-8 byte budget, truncates it at code-point + * boundaries with a `[TRUNCATED N BYTES BY OMNIROUTE]` marker whose own bytes + * count against the budget. + * + * `buildFollowUpSourceBody` appends the assistant tool-call turn and the + * bounded tool results to the source-format `messages` array. It never mutates + * its inputs, rejects orphan/mixed calls before building anything, and consumes + * the total output budget per-tool in result order. + */ + +export const MAX_RESULT_BYTES_PER_TOOL = 32_768; +export const MAX_RESULT_BYTES_TOTAL = 65_536; + +const NON_SERIALIZABLE_ERROR = "Tool result is not JSON-serializable"; + +function assertValidBudget(name: string, value: number): void { + if (!Number.isFinite(value) || !Number.isInteger(value) || value < 0) { + throw new RangeError(`${name} must be a non-negative finite integer, got ${value}`); + } +} + +function markerFor(droppedBytes: number): string { + return `[TRUNCATED ${droppedBytes} BYTES BY OMNIROUTE]`; +} + +/** + * Returns the longest code-point-aligned prefix of `text` whose UTF-8 byte + * length does not exceed `maxBytes`. Iterating `for...of` over a string yields + * full code points, so surrogate pairs (astral CJK, emoji) are never split and + * the result is always valid UTF-8. + */ +function truncateToCodePointBoundary(text: string, maxBytes: number): string { + let out = ""; + let bytes = 0; + for (const codePoint of text) { + const codePointBytes = Buffer.byteLength(codePoint, "utf8"); + if (bytes + codePointBytes > maxBytes) break; + out += codePoint; + bytes += codePointBytes; + } + return out; +} + +function projectSerializable(value: unknown): unknown { + if (value === undefined) return null; + if (value instanceof Error) return { error: value.message }; + if (typeof value === "bigint") return value.toString(); + return value; +} + +export function serializeBoundedToolResult(value: unknown, maxBytes: number): BoundedToolResult { + assertValidBudget("maxBytes", maxBytes); + + const projected = projectSerializable(value); + + let serialized: string | undefined; + try { + const raw = JSON.stringify(projected); + serialized = typeof raw === "string" ? raw : undefined; + } catch { + serialized = undefined; + } + + // Top-level function/symbol (JSON.stringify returns undefined, not a string) + // and serialization exceptions (cycles, nested BigInt, exotic objects) all + // resolve to the fixed non-serializable error shape. + if (serialized === undefined) { + serialized = JSON.stringify({ error: NON_SERIALIZABLE_ERROR }); + } + + const originalBytes = Buffer.byteLength(serialized, "utf8"); + if (originalBytes <= maxBytes) { + return { text: serialized, truncated: false, originalBytes }; + } + + const fullMarker = markerFor(originalBytes); + const fullMarkerBytes = Buffer.byteLength(fullMarker, "utf8"); + + // The full marker alone does not fit: return its code-point-safe UTF-8 prefix + // (maxBytes 0 yields the empty string). + if (maxBytes <= fullMarkerBytes) { + return { + text: truncateToCodePointBoundary(fullMarker, maxBytes), + truncated: true, + originalBytes, + }; + } + + // Reserve marker space first, then take the longest valid prefix of the text. + // `dropped <= originalBytes` so markerFor(dropped) is never longer than the + // reserved marker; the total therefore stays within maxBytes. + const prefixBudget = maxBytes - fullMarkerBytes; + const prefix = truncateToCodePointBoundary(serialized, prefixBudget); + const droppedBytes = originalBytes - Buffer.byteLength(prefix, "utf8"); + + return { + text: prefix + markerFor(droppedBytes), + truncated: true, + originalBytes, + }; +} + +function safeStringifyArguments(value: Record): string { + try { + const raw = JSON.stringify(value); + return typeof raw === "string" ? raw : "{}"; + } catch { + return "{}"; + } +} + +function validateCallsAndResults(toolCalls: ToolCall[], results: ExecutedToolResult[]): void { + if (toolCalls.length !== results.length) { + throw new Error( + `buildFollowUpSourceBody requires toolCalls (${toolCalls.length}) and results (${results.length}) to have the same length` + ); + } + + const callIds = toolCalls.map((call) => call.id); + const resultIds = results.map((result) => result.id); + + if (new Set(callIds).size !== callIds.length) { + throw new Error("buildFollowUpSourceBody requires unique tool call ids"); + } + if (new Set(resultIds).size !== resultIds.length) { + throw new Error("buildFollowUpSourceBody requires unique tool result ids"); + } + + const resultIdSet = new Set(resultIds); + for (const id of callIds) { + if (!resultIdSet.has(id)) { + throw new Error( + `buildFollowUpSourceBody requires every tool call to have a matching result (missing: ${id})` + ); + } + } + + for (const result of results) { + const call = toolCalls.find((c) => c.id === result.id); + if (call && call.name !== result.name) { + throw new Error( + `buildFollowUpSourceBody: result name "${result.name}" for id "${result.id}" does not match tool call name "${call.name}"` + ); + } + } +} + +function extractOpenAIMessage(response: Record): Record | null { + const choice = Array.isArray(response.choices) ? (response.choices[0] as unknown) : null; + if (choice && typeof choice === "object") { + const message = (choice as Record).message; + if (message && typeof message === "object" && !Array.isArray(message)) { + return message as Record; + } + } + if ( + response.message && + typeof response.message === "object" && + !Array.isArray(response.message) + ) { + return response.message as Record; + } + return null; +} + +/** + * Returns the assistant tool_calls exactly as the previous response carried + * them (provenance preserved), restricted to calls with a matching result. + * Falls back to reconstructing the OpenAI wire shape from the parsed + * `toolCalls` when the response has no tool_calls of its own. + * + * When the previous response does carry tool_calls, validates that each + * call ID appears exactly once with the correct name (fails closed on + * partial, duplicate, or mismatched entries). + */ +function resolveOpenAIAssistantToolCalls( + previousResponse: Record, + toolCalls: ToolCall[], + matchedIds: Set +): unknown[] { + let originalToolCalls: unknown[] = []; + const prevMessage = extractOpenAIMessage(previousResponse); + if (prevMessage && Array.isArray(prevMessage.tool_calls)) { + originalToolCalls = prevMessage.tool_calls as unknown[]; + } + if (originalToolCalls.length === 0 && toolCalls.length > 0) { + originalToolCalls = toolCalls.map((call) => ({ + id: call.id, + type: "function", + function: { name: call.name, arguments: safeStringifyArguments(call.arguments) }, + })); + } + + // Validate when previous response carries its own tool_calls. + const hasOwnToolCalls = + prevMessage && Array.isArray(prevMessage.tool_calls) && prevMessage.tool_calls.length > 0; + if (hasOwnToolCalls) { + const idCounts = new Map(); + const idNames = new Map(); + for (const raw of originalToolCalls) { + if (!raw || typeof raw !== "object") continue; + const rec = raw as Record; + const id = typeof rec.id === "string" ? rec.id : undefined; + if (id === undefined) continue; + idCounts.set(id, (idCounts.get(id) ?? 0) + 1); + if (idNames.has(id)) continue; + const fn = rec.function; + if ( + fn && + typeof fn === "object" && + typeof (fn as Record).name === "string" + ) { + idNames.set(id, (fn as Record).name as string); + } + } + for (const call of toolCalls) { + if (!matchedIds.has(call.id)) continue; + const count = idCounts.get(call.id) ?? 0; + if (count !== 1) { + throw new Error( + `buildFollowUpSourceBody: previous response tool_calls must contain exactly one match per call ID (id "${call.id}" has ${count})` + ); + } + const prevName = idNames.get(call.id); + if (prevName !== undefined && prevName !== call.name) { + throw new Error( + `buildFollowUpSourceBody: previous response tool_call "${call.id}" name "${prevName}" does not match tool call name "${call.name}"` + ); + } + } + } + + return originalToolCalls.filter((call) => { + if (!call || typeof call !== "object") return false; + const record = call as Record; + const id = typeof record.id === "string" ? record.id : record.call_id; + return typeof id === "string" && matchedIds.has(id); + }); +} + +/** + * Resolves tool_use blocks from the Claude previous response. When content + * blocks are present, validates that each call ID appears exactly once with + * the correct name (fails closed on partial, duplicate, or mismatched). + * Non-tool content blocks (text, thinking, reasoning) are preserved in their + * original order; only tool_use blocks are filtered to matched IDs. + */ +function resolveClaudeToolUseBlocks( + previousResponse: Record, + toolCalls: ToolCall[], + matchedIds: Set +): unknown[] { + if (Array.isArray(previousResponse.content)) { + const blocks = previousResponse.content as unknown[]; + + // Validate when content has tool_use blocks for matched IDs. + const idCounts = new Map(); + const idNames = new Map(); + for (const block of blocks) { + if (!block || typeof block !== "object") continue; + const rec = block as Record; + if (rec.type !== "tool_use" || typeof rec.id !== "string") continue; + if (!matchedIds.has(rec.id)) continue; + idCounts.set(rec.id, (idCounts.get(rec.id) ?? 0) + 1); + if (idNames.has(rec.id)) continue; + if (typeof rec.name === "string") { + idNames.set(rec.id, rec.name); + } + } + for (const call of toolCalls) { + if (!matchedIds.has(call.id)) continue; + const count = idCounts.get(call.id) ?? 0; + if (count !== 1) { + throw new Error( + `buildFollowUpSourceBody: previous response content must contain exactly one tool_use match per call ID (id "${call.id}" has ${count})` + ); + } + const prevName = idNames.get(call.id); + if (prevName !== undefined && prevName !== call.name) { + throw new Error( + `buildFollowUpSourceBody: previous response tool_use "${call.id}" name "${prevName}" does not match tool call name "${call.name}"` + ); + } + } + + // Match prepareClaudeRequest: keep all thinking/signature blocks, keep + // ordinary content only before the first tool_use, and retain only the + // tool_use blocks whose results are being replayed. + let foundToolUse = false; + const replayBlocks: unknown[] = []; + for (const block of blocks) { + if (!block || typeof block !== "object") continue; + const rec = block as Record; + if (rec.type === "tool_use") { + foundToolUse = true; + if (typeof rec.id === "string" && matchedIds.has(rec.id)) replayBlocks.push(block); + continue; + } + if (rec.type === "thinking" || rec.type === "redacted_thinking" || !foundToolUse) { + replayBlocks.push(block); + } + } + return replayBlocks; + } + + return toolCalls.map((call) => ({ + type: "tool_use", + id: call.id, + name: call.name, + input: call.arguments, + })); +} + +export function buildFollowUpSourceBody( + input: BuildFollowUpTranscriptInput +): Record { + const { sourceBody, previousResponse, toolCalls, results, sourceFormat } = input; + const maxResultBytes = input.maxResultBytes ?? MAX_RESULT_BYTES_PER_TOOL; + const maxTotalResultBytes = input.maxTotalResultBytes ?? MAX_RESULT_BYTES_TOTAL; + + if (sourceFormat !== "openai" && sourceFormat !== "claude") { + throw new Error('sourceFormat must be "openai" or "claude"'); + } + + assertValidBudget("maxResultBytes", maxResultBytes); + assertValidBudget("maxTotalResultBytes", maxTotalResultBytes); + + if (!Array.isArray(sourceBody.messages)) { + throw new Error("buildFollowUpSourceBody requires sourceBody.messages array"); + } + + validateCallsAndResults(toolCalls, results); + + // When a pre-serialized map is provided (from the loop's own budget pass), + // use its text verbatim instead of re-serializing. Otherwise serialize with + // the standard budget logic. + const serializedResultTextById = input.serializedResultTextById; + let boundedResults: BoundedToolResult[]; + if (serializedResultTextById) { + boundedResults = results.map((result) => { + const text = serializedResultTextById.get(result.id) ?? ""; + return { text, truncated: false, originalBytes: Buffer.byteLength(text, "utf8") }; + }); + } else { + let remainingBytes = maxTotalResultBytes; + boundedResults = results.map((result) => { + const itemMaxBytes = Math.min(maxResultBytes, remainingBytes); + const bounded = serializeBoundedToolResult(result.result, itemMaxBytes); + remainingBytes -= Buffer.byteLength(bounded.text, "utf8"); + return bounded; + }); + } + + const matchedIds = new Set(results.map((result) => result.id)); + const messages = [...(sourceBody.messages as unknown[])]; + + if (sourceFormat === "openai") { + const prevMessage = extractOpenAIMessage(previousResponse); + const previousContent = prevMessage && "content" in prevMessage ? prevMessage.content : null; + + const assistantMessage: Record = { + role: "assistant", + content: previousContent ?? null, + }; + const assistantToolCalls = resolveOpenAIAssistantToolCalls( + previousResponse, + toolCalls, + matchedIds + ); + if (assistantToolCalls.length > 0) { + assistantMessage.tool_calls = assistantToolCalls; + } + messages.push(assistantMessage); + + boundedResults.forEach((bounded, index) => { + messages.push({ + role: "tool", + tool_call_id: results[index].id, + content: bounded.text, + }); + }); + + return { ...sourceBody, messages, stream: false }; + } + + // Claude Messages: the original assistant tool_use turn, then a separate + // user tool_result message. A tool_result must never share the assistant + // content — Anthropic rejects it (openai-to-claude.ts:323). + const toolUseBlocks = resolveClaudeToolUseBlocks(previousResponse, toolCalls, matchedIds); + messages.push({ role: "assistant", content: toolUseBlocks }); + messages.push({ + role: "user", + content: boundedResults.map((bounded, index) => ({ + type: "tool_result", + tool_use_id: results[index].id, + content: bounded.text, + })), + }); + + return { ...sourceBody, messages }; +} diff --git a/src/lib/skills/injection.ts b/src/lib/skills/injection.ts index a0d65b5e5b..881a039189 100644 --- a/src/lib/skills/injection.ts +++ b/src/lib/skills/injection.ts @@ -253,6 +253,28 @@ function scoreAutoSkill( } export function injectSkills(options: InjectionOptions): unknown[] { + return injectSkillsWithMetadata(options).tools; +} + +export interface InjectSkillsWithMetadataResult { + tools: unknown[]; + injectedNames: string[]; +} + +function getToolNameFromDef(tool: unknown): string { + if (!tool || typeof tool !== "object") return ""; + const r = tool as Record; + if (typeof r.name === "string") return r.name; + if (r.function && typeof r.function === "object") { + const fn = r.function as Record; + if (typeof fn.name === "string") return fn.name; + } + return ""; +} + +export function injectSkillsWithMetadata( + options: InjectionOptions +): InjectSkillsWithMetadataResult { const contextText = buildContextText(options); const contextTokens = extractTokens(contextText); const backgroundTokens = extractTokens(toLowerText(options.backgroundReason)); @@ -295,7 +317,7 @@ export function injectSkills(options: InjectionOptions): unknown[] { apiKeyId: options.apiKeyId, reason: "no_enabled_skills", }); - return options.existingTools || []; + return { tools: options.existingTools || [], injectedNames: [] }; } log.info("skills.injection.injected", { @@ -317,11 +339,30 @@ export function injectSkills(options: InjectionOptions): unknown[] { } }); - if (options.existingTools && options.existingTools.length > 0) { - return [...injectedTools, ...options.existingTools]; + // Compute the set of existing tool names to exclude client collisions. + const existingToolNames = new Set( + (options.existingTools || []).map((t) => getToolNameFromDef(t)).filter(Boolean) + ); + + // Filter out skills whose encoded name collides with a client-declared tool. + const nonCollidingTools = injectedTools.filter((tool) => { + const name = getToolNameFromDef(tool); + return name && !existingToolNames.has(name); + }); + + const injectedNames: string[] = []; + for (const tool of nonCollidingTools) { + const name = getToolNameFromDef(tool); + if (name) { + injectedNames.push(name); + } } - return injectedTools; + if (options.existingTools && options.existingTools.length > 0) { + return { tools: [...nonCollidingTools, ...options.existingTools], injectedNames }; + } + + return { tools: nonCollidingTools, injectedNames }; } export function injectSkillTools( diff --git a/src/lib/skills/interception.ts b/src/lib/skills/interception.ts index 43c83c2d25..e91cb536f6 100644 --- a/src/lib/skills/interception.ts +++ b/src/lib/skills/interception.ts @@ -6,10 +6,29 @@ import { detectProvider, decodeSkillToolName } from "./injection"; import { OMNIROUTE_WEB_SEARCH_FALLBACK_TOOL_NAME } from "@omniroute/open-sse/services/webSearchFallback.ts"; import { OMNIROUTE_WEB_FETCH_FALLBACK_TOOL_NAME } from "@omniroute/open-sse/services/webFetchInterception.ts"; import { sanitizeErrorMessage } from "@omniroute/open-sse/utils/errorSanitization.ts"; +import { runWithServerToolFence } from "./toolExecutionFence"; +import type { ExecutedToolResult, ToolCall, ExecutionContext } from "./toolLoopTypes"; import { logger } from "../../../open-sse/utils/logger.ts"; const log = logger("SKILLS_INTERCEPTION"); +/** + * Typed error for server-owned tool execution control-flow states + * (in_progress, unknown, identity_conflict). These must never be fed + * back to the model as tool results — they represent infrastructure + * conditions that should surface as HTTP-level errors. + */ +export class ServerOwnedExecutionError extends Error { + readonly code: string; + readonly httpStatus: number; + constructor(message: string, code: string, httpStatus: number) { + super(message); + this.name = "ServerOwnedExecutionError"; + this.code = code; + this.httpStatus = httpStatus; + } +} + function toSafeSkillErrorMessage(value: unknown): string { try { const raw = value instanceof Error ? value.message : value; @@ -24,24 +43,7 @@ function projectSkillResultForPublicResponse(result: unknown): unknown { return projectSkillOutputForBoundary(result as Record); } -interface ToolCall { - id: string; - name: string; - arguments: Record; -} - -interface ExecutionContext { - apiKeyId: string; - sessionId: string; - requestId: string; - builtinToolNames?: string[]; - customSkillExecutionEnabled?: boolean; - // #7339: threaded through to the web_fetch builtin so it can resolve a per-model - // pinned fetch backend (interceptionRules.fetchBackend). Optional — every other - // builtin/skill ignores these. - provider?: string; - model?: string; -} +// ToolCall and ExecutionContext types are imported from ./toolLoopTypes.ts const BUILTIN_TOOL_ALIASES: Record = { [OMNIROUTE_WEB_SEARCH_FALLBACK_TOOL_NAME]: "web_search", @@ -200,8 +202,19 @@ export async function interceptToolCalls( return results; } -export function extractToolCalls(response: any, modelId: string): ToolCall[] { - const provider = detectProvider(modelId); +export function extractToolCalls(response: any, modelIdOrSourceFormat: string): ToolCall[] { + // Accept either a sourceFormat ("openai" | "claude") or a model ID string. + // Map known sourceFormat values; fall back to detectProvider for model IDs. + const format = + modelIdOrSourceFormat === "openai" || modelIdOrSourceFormat === "claude" + ? modelIdOrSourceFormat + : undefined; + const provider = + format === "claude" + ? "anthropic" + : format === "openai" + ? "openai" + : detectProvider(modelIdOrSourceFormat); switch (provider) { case "openai": { @@ -446,3 +459,349 @@ export async function handleToolCallExecution( return response; } } + +// ─── Task 3: ownership classifier ──────────────────────────────────────────── + +export async function classifyServerOwnedCalls( + toolCalls: ToolCall[], + context: ExecutionContext +): Promise<{ serverOwned: ToolCall[]; clientNative: ToolCall[] }> { + const builtinSet = new Set(context.builtinToolNames || []); + const customSet = new Set(context.injectedCustomSkillNames || []); + + const serverOwned: ToolCall[] = []; + const clientNative: ToolCall[] = []; + + for (const call of toolCalls) { + if (builtinSet.has(call.name)) { + serverOwned.push(call); + } else if (context.customSkillExecutionEnabled && customSet.has(call.name)) { + serverOwned.push(call); + } else { + clientNative.push(call); + } + } + + return { serverOwned, clientNative }; +} + +// ─── Task 3: executeServerOwned — fence-gated execution callback ──────────── + +const LEASE_DURATION_MS = 120_000; + +/** + * Runtime seam: overridable fence function for testing. + * When null, the real `runWithServerToolFence` is used. + * Tests may set this to inject controlled fence outcomes. + */ +let _fenceFn: typeof runWithServerToolFence | null = null; + +export function setFenceFnForTesting(fn: typeof runWithServerToolFence | null): void { + _fenceFn = fn; +} + +export async function executeServerOwned( + calls: ToolCall[], + context: ExecutionContext, + fenceFn?: typeof runWithServerToolFence +): Promise { + if (context.executionFenceEnabled && !context.requestIdentity) { + throw new Error( + "executeServerOwned requires context.requestIdentity when executionFenceEnabled is true" + ); + } + + const results: ExecutedToolResult[] = []; + + for (const call of calls) { + const builtinHandlerName = resolveBuiltinHandlerName(call.name, context); + const isMemoryBuiltin = builtinHandlerName && MEMORY_TOOL_NAMES.has(builtinHandlerName); + const isOrdinaryBuiltin = builtinHandlerName && builtinHandlerName in builtinSkills; + const isCustomSkill = + !builtinHandlerName && + context.customSkillExecutionEnabled && + context.injectedCustomSkillNames?.includes(call.name); + + const executeFn = async (executionId: string): Promise => { + if (isMemoryBuiltin) { + const handlerName = builtinHandlerName as keyof typeof memoryBuiltinHandlers; + return memoryBuiltinHandlers[handlerName](call.arguments, { + apiKeyId: context.apiKeyId, + sessionId: context.sessionId, + }); + } + if (isOrdinaryBuiltin) { + const handlerName = builtinHandlerName as keyof typeof builtinSkills; + return builtinSkills[handlerName](call.arguments, { + apiKeyId: context.apiKeyId, + sessionId: context.sessionId, + provider: context.provider, + model: context.model, + }); + } + if (isCustomSkill) { + const decodedName = decodeSkillToolName(call.name); + const [name, version] = decodedName.includes("@") + ? decodedName.split("@", 2) + : [decodedName, "latest"]; + const skillName = version === "latest" ? name : `${name}@${version}`; + const execution = await skillExecutor.executeClaimed( + skillName, + call.arguments, + { + apiKeyId: context.apiKeyId, + sessionId: context.sessionId, + }, + executionId + ); + return ( + execution.output ?? + (execution.errorMessage + ? { error: toSafeSkillErrorMessage(execution.errorMessage) } + : { error: "Skill execution returned no output" }) + ); + } + throw new Error(`No handler for tool: ${call.name}`); + }; + + if (context.executionFenceEnabled && context.requestIdentity) { + const activeFenceFn = fenceFn ?? _fenceFn ?? runWithServerToolFence; + const fenceResult = await activeFenceFn({ + apiKeyId: context.apiKeyId, + requestIdentity: context.requestIdentity, + toolCallId: call.id, + toolName: call.name, + arguments: call.arguments, + leaseDurationMs: LEASE_DURATION_MS, + execute: executeFn, + }); + + switch (fenceResult.kind) { + case "executed": + results.push({ + id: call.id, + name: call.name, + result: projectSkillResultForPublicResponse(fenceResult.value), + replayed: false, + }); + break; + case "replayed": { + if (fenceResult.status === "error") { + throw new ServerOwnedExecutionError( + fenceResult.errorMessage ?? "Tool execution failed", + "TOOL_EXECUTION_ERROR", + 500 + ); + } + if (fenceResult.status === "timeout") { + throw new ServerOwnedExecutionError( + fenceResult.errorMessage ?? "Tool execution timed out", + "TOOL_EXECUTION_TIMEOUT", + 504 + ); + } + results.push({ + id: call.id, + name: call.name, + result: projectSkillResultForPublicResponse(fenceResult.value), + replayed: true, + }); + break; + } + case "in_progress": + throw new ServerOwnedExecutionError( + "Tool execution in progress", + "TOOL_IN_PROGRESS", + 409 + ); + case "unknown": + throw new ServerOwnedExecutionError( + "Tool execution state unknown", + "TOOL_STATE_UNKNOWN", + 500 + ); + case "identity_conflict": + throw new ServerOwnedExecutionError( + "Tool execution identity conflict", + "IDENTITY_CONFLICT", + 409 + ); + } + } else { + // Flag-off path: no fence, dispatch directly. + try { + const value = await executeFn(""); + results.push({ + id: call.id, + name: call.name, + result: projectSkillResultForPublicResponse(value), + replayed: false, + }); + } catch (err) { + results.push({ + id: call.id, + name: call.name, + result: { error: toSafeSkillErrorMessage(err) }, + replayed: false, + }); + } + } + } + + return results; +} + +// ─── Task 3: pure escape formatter ─────────────────────────────────────────── + +function extractOpenAIToolCalls( + response: Record +): Array<{ id: string; function: { name: string; arguments: string } }> { + const rootToolCalls = Array.isArray(response.tool_calls) ? response.tool_calls : []; + const choiceToolCalls = Array.isArray(response.choices) + ? (response.choices as any[]).flatMap((choice: any) => + Array.isArray(choice?.message?.tool_calls) ? choice.message.tool_calls : [] + ) + : []; + return rootToolCalls.length > 0 ? rootToolCalls : choiceToolCalls; +} + +function getOpenAIResponseOutput( + response: Record +): { target: Record; output: unknown[] } | null { + if (Array.isArray(response.output)) { + return { target: response, output: response.output }; + } + if ( + response.response && + typeof response.response === "object" && + !Array.isArray(response.response) && + Array.isArray((response.response as Record).output) + ) { + return { + target: response.response as Record, + output: (response.response as Record).output as unknown[], + }; + } + return null; +} + +export function formatEscapeHatchResponse( + response: Record, + serverCalls: ToolCall[], + results: ExecutedToolResult[], + clientCalls: ToolCall[], + sourceFormat: "openai" | "claude", + serializedResultTextById?: Map +): Record { + const serverIds = new Set(serverCalls.map((c) => c.id)); + + if (sourceFormat === "openai") { + // Check for Responses API format. + const responsesOutput = getOpenAIResponseOutput(response); + if (responsesOutput) { + // For Responses, append function_call_output for server calls. + const functionOutputs = results + .filter((r) => serverIds.has(r.id)) + .map((r) => ({ + type: "function_call_output", + call_id: r.id, + output: serializedResultTextById?.get(r.id) ?? JSON.stringify(r.result), + })); + return { + ...response, + response: + responsesOutput.target !== response + ? { + ...(response.response as Record), + output: [...responsesOutput.output, ...functionOutputs], + } + : undefined, + output: + responsesOutput.target === response + ? [...responsesOutput.output, ...functionOutputs] + : response.output, + }; + } + + // Chat Completions format. + const originalToolCalls = extractOpenAIToolCalls(response); + const remainingToolCalls = originalToolCalls.filter( + (tc: any) => !serverIds.has(tc.id || tc.call_id) + ); + + // Build result text from server results. + const resultTexts = results + .filter((r) => serverIds.has(r.id)) + .map( + (r) => + `[${r.name} result]\n${serializedResultTextById?.get(r.id) ?? JSON.stringify(r.result)}` + ) + .join("\n\n"); + + const existingContent = + typeof response.choices?.[0]?.message?.content === "string" + ? response.choices[0].message.content + : ""; + const newContent = existingContent ? `${existingContent}\n\n${resultTexts}` : resultTexts; + + // Clone response to avoid mutation. + const formatted = JSON.parse(JSON.stringify(response)); + if (formatted.choices?.[0]?.message) { + formatted.choices[0].message.content = newContent; + formatted.choices[0].message.tool_calls = + remainingToolCalls.length > 0 ? remainingToolCalls : undefined; + } + + // Mixed → keep tool_calls finish_reason; all-server → stop. + if (remainingToolCalls.length === 0 && clientCalls.length === 0) { + formatted.choices[0].finish_reason = "stop"; + } + + return formatted; + } + + if (sourceFormat === "claude") { + const remainingContent = (Array.isArray(response.content) ? response.content : []).filter( + (block: any) => !(block?.type === "tool_use" && serverIds.has(block.id)) + ); + + // Build result text blocks. + const resultTextBlocks = results + .filter((r) => serverIds.has(r.id)) + .map((r) => ({ + type: "text", + text: `[${r.name} result]\n${serializedResultTextById?.get(r.id) ?? JSON.stringify(r.result)}`, + })); + + // Insert result text blocks before the first remaining tool_use. + const firstRemainingIndex = remainingContent.findIndex( + (block: any) => block?.type === "tool_use" + ); + + let newContent: unknown[]; + if (firstRemainingIndex === -1) { + newContent = [...remainingContent, ...resultTextBlocks]; + } else { + newContent = [ + ...remainingContent.slice(0, firstRemainingIndex), + ...resultTextBlocks, + ...remainingContent.slice(firstRemainingIndex), + ]; + } + + const formatted: Record = { ...response, content: newContent }; + + // All-server → end_turn; mixed → keep original stop_reason. + const remainingToolUseCount = remainingContent.filter( + (b: any) => b?.type === "tool_use" + ).length; + if (remainingToolUseCount === 0 && clientCalls.length === 0) { + formatted.stop_reason = "end_turn"; + formatted.stop_sequence = null; + } + + return formatted; + } + + return response; +} diff --git a/src/lib/skills/serverOwnedToolLoop.ts b/src/lib/skills/serverOwnedToolLoop.ts new file mode 100644 index 0000000000..a25523ff90 --- /dev/null +++ b/src/lib/skills/serverOwnedToolLoop.ts @@ -0,0 +1,466 @@ +/** + * Server-owned tool loop state machine (spec §5.5). + * + * Handles the complete lifecycle: + * 1. Classify ownership of tool calls + * 2. Execute server-owned tools + * 3. Serialize results within UTF-8 byte budgets + * 4. Build accumulated transcript and resume upstream + * 5. Terminate on any boundary condition + */ + +import type { + ServerOwnedToolLoopOptions, + ServerOwnedToolLoopResult, + NonStreamingProviderLegResult, + ProviderLegUsage, + ProviderLegReceipt, + ChatCoreErrorResult, + ExecutedToolResult, +} from "./toolLoopTypes.ts"; +import { + classifyServerOwnedCalls, + extractToolCalls, + formatEscapeHatchResponse, + ServerOwnedExecutionError, +} from "./interception.ts"; +import { + buildFollowUpSourceBody, + serializeBoundedToolResult, + MAX_RESULT_BYTES_PER_TOOL, + MAX_RESULT_BYTES_TOTAL, +} from "./followUpTranscript.ts"; +import { createErrorResult } from "@omniroute/open-sse/utils/error.ts"; + +// ─── Constants ──────────────────────────────────────────────────────────────── + +export const MAX_FOLLOW_UPS = 3; +export const LOOP_BUDGET_MS = 120_000; +export const MIN_REMAINING_FOR_FOLLOW_UP_MS = 10_000; + +// ─── Usage aggregation ──────────────────────────────────────────────────────── + +export function aggregateProviderLegUsage( + usages: Array +): ProviderLegUsage { + let prompt_tokens = 0; + let completion_tokens = 0; + let cached_tokens: number | undefined; + let cache_read_input_tokens: number | undefined; + let cache_creation_input_tokens: number | undefined; + let reasoning_tokens: number | undefined; + let cost_in_usd_ticks: number | undefined; + + for (const u of usages) { + if (!u) continue; + prompt_tokens += u.prompt_tokens; + completion_tokens += u.completion_tokens; + if (u.cached_tokens !== undefined) { + cached_tokens = (cached_tokens ?? 0) + u.cached_tokens; + } + if (u.cache_read_input_tokens !== undefined) { + cache_read_input_tokens = (cache_read_input_tokens ?? 0) + u.cache_read_input_tokens; + } + if (u.cache_creation_input_tokens !== undefined) { + cache_creation_input_tokens = + (cache_creation_input_tokens ?? 0) + u.cache_creation_input_tokens; + } + if (u.reasoning_tokens !== undefined) { + reasoning_tokens = (reasoning_tokens ?? 0) + u.reasoning_tokens; + } + if (u.cost_in_usd_ticks !== undefined) { + cost_in_usd_ticks = (cost_in_usd_ticks ?? 0) + u.cost_in_usd_ticks; + } + } + + return { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + ...(cached_tokens !== undefined ? { cached_tokens } : {}), + ...(cache_read_input_tokens !== undefined ? { cache_read_input_tokens } : {}), + ...(cache_creation_input_tokens !== undefined ? { cache_creation_input_tokens } : {}), + ...(reasoning_tokens !== undefined ? { reasoning_tokens } : {}), + ...(cost_in_usd_ticks !== undefined ? { cost_in_usd_ticks } : {}), + }; +} + +function aggregateUsageOrNull(usages: Array): ProviderLegUsage | null { + const hasNonNull = usages.some((u) => u !== null); + if (!hasNonNull) return null; + return aggregateProviderLegUsage(usages); +} + +// ─── Execution error → errorResult mapping ──────────────────────────────────── + +function mapExecutionError(err: ServerOwnedExecutionError): ChatCoreErrorResult { + return createErrorResult( + err.httpStatus, + err.message, + null, + err.code, + "server_tool_execution_error" + ) as unknown as ChatCoreErrorResult; +} + +// ─── State machine ──────────────────────────────────────────────────────────── + +export async function runServerOwnedToolLoop( + options: ServerOwnedToolLoopOptions +): Promise { + const now = options.now ?? performance.now.bind(performance); + const loopStartedAtMs = now(); + const loopDeadlineAtMs = Math.min(options.deadlineAtMs, loopStartedAtMs + LOOP_BUDGET_MS); + + const maxFollowUps = options.maxFollowUps ?? MAX_FOLLOW_UPS; + const maxResultBytes = options.maxResultBytes ?? MAX_RESULT_BYTES_PER_TOOL; + const maxTotalResultBytes = options.maxTotalResultBytes ?? MAX_RESULT_BYTES_TOTAL; + + // State + let currentSourceBody = options.sourceBody; + let currentLeg: NonStreamingProviderLegResult = options.initialLeg; + let cumulativeOutputBytes = 0; + let followUps = 0; + const receipts: ProviderLegReceipt[] = [options.initialLeg.receipt]; + const usages: Array = [options.initialLeg.usage]; + + // Accumulate initial leg receipt cost + let totalCostUsd = options.initialLeg.receipt.computedCostUsd ?? 0; + + const errorResult: (msg: string, status?: number, code?: string) => ServerOwnedToolLoopResult = ( + msg, + status = 500, + code = "internal_error" + ) => ({ + kind: "error", + errorResult: createErrorResult(status, msg, null, code) as unknown as ChatCoreErrorResult, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination: "provider_error", + }); + + // Main loop + while (true) { + // Extract tool calls from current leg response + const response = currentLeg.kind === "ok" ? currentLeg.response : undefined; + if (!response) { + return errorResult("No response from provider leg"); + } + + // Extract tool calls by actual source format, not model alias heuristic + const toolCalls = extractToolCalls(response, options.sourceFormat); + + // Classify ownership + const { serverOwned, clientNative } = await classifyServerOwnedCalls( + toolCalls, + options.executionContext + ); + + // No server-owned calls → done + if (serverOwned.length === 0) { + const termination = clientNative.length > 0 ? "client_tools" : "completed"; + return { + kind: "ok", + response: currentLeg.kind === "ok" ? currentLeg.response : undefined, + responseForMemoryExtraction: + currentLeg.kind === "ok" ? currentLeg.responseForMemoryExtraction : undefined, + finalProviderBody: currentLeg.kind === "ok" ? currentLeg.providerBody : undefined, + finalProviderRequest: currentLeg.kind === "ok" ? currentLeg.providerRequest : undefined, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination, + }; + } + + // Mixed tools → escape hatch, no follow-up + if (clientNative.length > 0) { + // Check abort before execute + if (options.abortSignal?.aborted) { + return { + kind: "error", + errorResult: createErrorResult( + 499, + "Client closed request", + null, + "client_closed_request", + "invalid_request_error" + ) as unknown as ChatCoreErrorResult, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination: "client_abort", + }; + } + + // Execute server calls, then format with bounded serialization + const execResults = await options.executeServerOwned(serverOwned, options.executionContext); + + // Build serialized text map using bounded serialization + const serMap = new Map(); + let remainingBytes = maxTotalResultBytes; + for (const r of execResults) { + const itemMaxBytes = Math.min(maxResultBytes, remainingBytes); + const bounded = serializeBoundedToolResult(r.result, itemMaxBytes); + serMap.set(r.id, bounded.text); + remainingBytes -= Buffer.byteLength(bounded.text, "utf8"); + } + + const formatted = formatEscapeHatchResponse( + response, + serverOwned, + execResults, + clientNative, + options.sourceFormat, + serMap + ); + + return { + kind: "ok", + response: formatted, + responseForMemoryExtraction: + currentLeg.kind === "ok" ? currentLeg.responseForMemoryExtraction : undefined, + finalProviderBody: currentLeg.kind === "ok" ? currentLeg.providerBody : undefined, + finalProviderRequest: currentLeg.kind === "ok" ? currentLeg.providerRequest : undefined, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination: "mixed_tools", + }; + } + + // All server-owned: check abort before execute + if (options.abortSignal?.aborted) { + return { + kind: "error", + errorResult: createErrorResult( + 499, + "Client closed request", + null, + "client_closed_request", + "invalid_request_error" + ) as unknown as ChatCoreErrorResult, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination: "client_abort", + }; + } + + // Execute server-owned calls + let execResults: ExecutedToolResult[]; + try { + execResults = await options.executeServerOwned(serverOwned, options.executionContext); + } catch (err: unknown) { + if (err instanceof ServerOwnedExecutionError) { + const termMap: Record = { + TOOL_IN_PROGRESS: "execution_in_progress", + TOOL_STATE_UNKNOWN: "execution_unknown", + IDENTITY_CONFLICT: "execution_identity_conflict", + TOOL_EXECUTION_ERROR: "execution_error", + TOOL_EXECUTION_TIMEOUT: "execution_timeout", + }; + const termination = termMap[err.code] ?? "execution_error"; + return { + kind: "error", + errorResult: mapExecutionError(err), + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination, + }; + } + throw err; + } + + // Serialize results with cumulative UTF-8 budget + const serMap = new Map(); + let anyTruncated = false; + let remainingBytes = maxTotalResultBytes - cumulativeOutputBytes; + if (remainingBytes < 0) remainingBytes = 0; + + for (const r of execResults) { + const itemMaxBytes = Math.min(maxResultBytes, remainingBytes); + const bounded = serializeBoundedToolResult(r.result, itemMaxBytes); + serMap.set(r.id, bounded.text); + const bytes = Buffer.byteLength(bounded.text, "utf8"); + cumulativeOutputBytes += bytes; + remainingBytes -= bytes; + if (bounded.truncated) { + anyTruncated = true; + } + } + + // Check budget exhaustion — any truncated result terminates + if (anyTruncated || cumulativeOutputBytes >= maxTotalResultBytes) { + const formatted = formatEscapeHatchResponse( + response, + serverOwned, + execResults, + [], + options.sourceFormat, + serMap + ); + + return { + kind: "ok", + response: formatted, + responseForMemoryExtraction: + currentLeg.kind === "ok" ? currentLeg.responseForMemoryExtraction : undefined, + finalProviderBody: currentLeg.kind === "ok" ? currentLeg.providerBody : undefined, + finalProviderRequest: currentLeg.kind === "ok" ? currentLeg.providerRequest : undefined, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination: "tool_output_budget", + }; + } + + // Check follow-up limit + if (followUps >= maxFollowUps) { + const formatted = formatEscapeHatchResponse( + response, + serverOwned, + execResults, + [], + options.sourceFormat, + serMap + ); + + return { + kind: "ok", + response: formatted, + responseForMemoryExtraction: + currentLeg.kind === "ok" ? currentLeg.responseForMemoryExtraction : undefined, + finalProviderBody: currentLeg.kind === "ok" ? currentLeg.providerBody : undefined, + finalProviderRequest: currentLeg.kind === "ok" ? currentLeg.providerRequest : undefined, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination: "max_followups", + }; + } + + // Check deadline before resume + const remainingMs = loopDeadlineAtMs - now(); + if (remainingMs < MIN_REMAINING_FOR_FOLLOW_UP_MS) { + const formatted = formatEscapeHatchResponse( + response, + serverOwned, + execResults, + [], + options.sourceFormat, + serMap + ); + + return { + kind: "ok", + response: formatted, + responseForMemoryExtraction: + currentLeg.kind === "ok" ? currentLeg.responseForMemoryExtraction : undefined, + finalProviderBody: currentLeg.kind === "ok" ? currentLeg.providerBody : undefined, + finalProviderRequest: currentLeg.kind === "ok" ? currentLeg.providerRequest : undefined, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination: "deadline", + }; + } + + // Check abort before resume + if (options.abortSignal?.aborted) { + return { + kind: "error", + errorResult: createErrorResult( + 499, + "Client closed request", + null, + "client_closed_request", + "invalid_request_error" + ) as unknown as ChatCoreErrorResult, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps, + termination: "client_abort", + }; + } + + // Build accumulated transcript + const nextSourceBody = buildFollowUpSourceBody({ + sourceBody: currentSourceBody, + previousResponse: response, + toolCalls: serverOwned, + results: execResults, + sourceFormat: options.sourceFormat, + maxResultBytes, + maxTotalResultBytes: maxTotalResultBytes - cumulativeOutputBytes, + serializedResultTextById: serMap, + }); + + // Resume upstream + const nextLeg = await options.resumeUpstream( + nextSourceBody, + options.initialLeg.connectionId, + loopDeadlineAtMs + ); + + // Provider error → preserve identity + if (nextLeg.kind === "error") { + receipts.push(nextLeg.receipt); + usages.push(nextLeg.usage); + totalCostUsd += nextLeg.receipt.computedCostUsd ?? 0; + + return { + kind: "error", + errorResult: nextLeg.result, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps: followUps + 1, + termination: "provider_error", + }; + } + + // Connection mismatch check + if (nextLeg.connectionId !== options.initialLeg.connectionId) { + receipts.push(nextLeg.receipt); + usages.push(nextLeg.usage); + totalCostUsd += nextLeg.receipt.computedCostUsd ?? 0; + + return { + kind: "error", + errorResult: createErrorResult( + 409, + "Follow-up connection does not match initial connection", + null, + "LEASE_CONNECTION_MISMATCH", + "lease_error" + ) as unknown as ChatCoreErrorResult, + cumulativeUsage: aggregateUsageOrNull(usages), + totalCostUsd, + receipts, + followUps: followUps + 1, + termination: "connection_mismatch", + }; + } + + // Update state for next iteration + currentSourceBody = nextSourceBody; + currentLeg = nextLeg; + followUps++; + receipts.push(nextLeg.receipt); + usages.push(nextLeg.usage); + totalCostUsd += nextLeg.receipt.computedCostUsd ?? 0; + } +} diff --git a/src/lib/skills/stableJson.ts b/src/lib/skills/stableJson.ts new file mode 100644 index 0000000000..acfebdad59 --- /dev/null +++ b/src/lib/skills/stableJson.ts @@ -0,0 +1,121 @@ +import { createHash } from "node:crypto"; + +const REJECT_MSG = "Value cannot be represented as canonical JSON"; + +function codePointCompare(a: string, b: string): number { + const aLen = a.length; + const bLen = b.length; + let i = 0; + let j = 0; + while (i < aLen && j < bLen) { + const aCode = a.charCodeAt(i); + const bCode = b.charCodeAt(j); + // Check if either is a surrogate pair + if (aCode >= 0xd800 && aCode <= 0xdbff && i + 1 < aLen) { + const aFull = (aCode - 0xd800) * 0x400 + (a.charCodeAt(i + 1) - 0xdc00) + 0x10000; + if (bCode >= 0xd800 && bCode <= 0xdbff && j + 1 < bLen) { + const bFull = (bCode - 0xd800) * 0x400 + (b.charCodeAt(j + 1) - 0xdc00) + 0x10000; + if (aFull !== bFull) return aFull - bFull; + i += 2; + j += 2; + } else { + // astral vs BMP + return 1; + } + } else if (bCode >= 0xd800 && bCode <= 0xdbff && j + 1 < bLen) { + return -1; + } else { + if (aCode !== bCode) return aCode - bCode; + i++; + j++; + } + } + return aLen - bLen; +} + +function canonicalStringify(value: unknown, seen: Set): string { + if (value === undefined) throw new TypeError(REJECT_MSG); + if (typeof value === "bigint") throw new TypeError(REJECT_MSG); + if (typeof value === "symbol") throw new TypeError(REJECT_MSG); + if (typeof value === "function") throw new TypeError(REJECT_MSG); + + if (typeof value === "number") { + // Normalize -0 to 0 + const normalized = Object.is(value, -0) ? 0 : value; + if (!Number.isFinite(normalized)) throw new TypeError(REJECT_MSG); + return String(normalized); + } + + if (typeof value === "string") return JSON.stringify(value); + if (typeof value === "boolean") return String(value); + if (value === null) return "null"; + + if (typeof value === "object") { + if (seen.has(value)) throw new TypeError(REJECT_MSG); + seen.add(value); + + if (Array.isArray(value)) { + // Reject sparse arrays + const len = (value as unknown[]).length; + for (let i = 0; i < len; i++) { + if (!(i in (value as unknown[]))) { + throw new TypeError(REJECT_MSG); + } + } + // Reject arrays with getters + for (let i = 0; i < len; i++) { + const desc = Object.getOwnPropertyDescriptor(value, i); + if (desc && (desc.get || desc.set)) { + throw new TypeError(REJECT_MSG); + } + } + const items = (value as unknown[]).map((v) => canonicalStringify(v, seen)); + seen.delete(value); + return `[${items.join(",")}]`; + } + + // Reject non-plain objects: Date, Map, Set, class instances, etc. + const proto = Object.getPrototypeOf(value); + if (proto !== Object.prototype && proto !== null) { + throw new TypeError(REJECT_MSG); + } + + // Plain object: check for accessors on any own key, then sort by code point + const keys = Object.keys(value); + for (const k of keys) { + const desc = Object.getOwnPropertyDescriptor(value, k); + if (desc && (desc.get || desc.set)) { + throw new TypeError(REJECT_MSG); + } + } + keys.sort(codePointCompare); + const pairs = keys.map( + (k) => + `${JSON.stringify(k)}:${canonicalStringify((value as Record)[k], seen)}` + ); + seen.delete(value); + return `{${pairs.join(",")}}`; + } + + throw new TypeError(REJECT_MSG); +} + +export function canonicalJson(value: unknown): string { + return canonicalStringify(value, new Set()); +} + +export function canonicalJsonSha256(value: unknown): string { + const json = canonicalJson(value); + return createHash("sha256").update(json, "utf8").digest("hex"); +} + +export function deriveToolRequestIdentity(input: { + apiKeyId: string; + stableClientRequestId: string | null; + skillRequestId: string; + postInjectionBody: Record; +}): string { + const stableKey = input.stableClientRequestId ?? input.skillRequestId; + const bodyDigest = canonicalJsonSha256(input.postInjectionBody); + return `${input.apiKeyId}:${stableKey}:${bodyDigest}`; +} diff --git a/src/lib/skills/toolExecutionFence.ts b/src/lib/skills/toolExecutionFence.ts new file mode 100644 index 0000000000..6f7558b6b6 --- /dev/null +++ b/src/lib/skills/toolExecutionFence.ts @@ -0,0 +1,241 @@ +import { canonicalJsonSha256 } from "./stableJson"; +import { + claimServerToolExecution, + finalizeServerToolExecution, + readRow, +} from "../db/skillExecutionFence"; +import type { SqliteAdapter } from "../db/adapters/types"; +import { getDbInstance } from "../db/core"; + +export type RunWithServerToolFenceResult = + | { kind: "executed"; value: T } + | { kind: "replayed"; value: T; status: "success"; errorMessage: null } + | { + kind: "replayed"; + value: unknown | null; + status: "error" | "timeout"; + errorMessage: string | null; + } + | { kind: "in_progress" } + | { kind: "unknown" } + | { kind: "identity_conflict" }; + +type ExecutionKey = string; + +const activePromises = new Map< + ExecutionKey, + { promise: Promise; status: "pending" | "resolved" | "rejected" } +>(); + +const POLL_INTERVAL_MS = 50; +const MAX_POLL_MS = 2_000; + +function buildExecutionKey( + apiKeyId: string, + requestIdentity: string, + toolCallId: string +): ExecutionKey { + return `${apiKeyId}:${requestIdentity}:${toolCallId}`; +} + +export interface RunWithServerToolFenceOptions { + apiKeyId: string; + requestIdentity: string; + toolCallId: string; + toolName: string; + arguments: Record; + leaseDurationMs: number; + execute: (executionId: string) => Promise; + now?: () => number; + sleep?: (ms: number) => Promise; + db?: SqliteAdapter; +} + +export async function runWithServerToolFence( + input: RunWithServerToolFenceOptions +): Promise> { + const db = input.db ?? getDbInstance(); + const now = input.now ?? (() => Date.now()); + const sleep = input.sleep ?? ((ms: number) => new Promise((resolve) => setTimeout(resolve, ms))); + const inputDigest = canonicalJsonSha256(input.arguments); + const leaseExpiresAt = new Date(now() + input.leaseDurationMs).toISOString(); + + const claim = claimServerToolExecution( + { + apiKeyId: input.apiKeyId, + requestIdentity: input.requestIdentity, + toolCallId: input.toolCallId, + toolName: input.toolName, + inputDigest, + leaseExpiresAt, + }, + db, + now() + ); + + switch (claim.kind) { + case "claimed": { + const key = buildExecutionKey(input.apiKeyId, input.requestIdentity, input.toolCallId); + const claimStartTime = now(); + const wrapperPromise = (async () => { + try { + const value = await input.execute(claim.executionId); + const durationMs = now() - claimStartTime; + finalizeServerToolExecution( + { + executionId: claim.executionId, + status: "success", + output: value, + errorMessage: null, + durationMs, + }, + db + ); + return { kind: "success" as const, value, errorMessage: null }; + } catch (err: unknown) { + const message = err instanceof Error ? err.message : String(err); + const safeMessage = message.replace(/\bat\s+\/[^\s"']+/g, "[stack-redacted]"); + const durationMs = now() - claimStartTime; + finalizeServerToolExecution( + { + executionId: claim.executionId, + status: "error", + output: null, + errorMessage: safeMessage, + durationMs, + }, + db + ); + return { kind: "error" as const, value: null, errorMessage: safeMessage }; + } + })(); + + const entry: { promise: Promise; status: "pending" | "resolved" | "rejected" } = { + promise: wrapperPromise as Promise, + status: "pending", + }; + activePromises.set(key, entry); + wrapperPromise.then( + () => { + entry.status = "resolved"; + }, + () => { + entry.status = "rejected"; + } + ); + + try { + const result = await wrapperPromise; + if (result.kind === "success") { + return { kind: "executed", value: result.value as T }; + } + // Handler errored — finalize already done, propagate + throw new Error(result.errorMessage ?? "tool execution failed"); + } finally { + activePromises.delete(key); + } + } + + case "replay": { + if (claim.status === "success") { + return { + kind: "replayed", + value: claim.output as T, + status: "success", + errorMessage: null, + }; + } + return { + kind: "replayed", + value: claim.output ?? null, + status: claim.status, + errorMessage: claim.errorMessage, + }; + } + + case "in_progress": { + const key = buildExecutionKey(input.apiKeyId, input.requestIdentity, input.toolCallId); + const deadline = now() + MAX_POLL_MS; + while (now() < deadline) { + // Check process-internal promise first + const active = activePromises.get(key); + if (active) { + if (active.status === "resolved") { + try { + const result = await active.promise; + if (result && typeof result === "object" && "status" in result) { + const r = result as { value: unknown; status: string; errorMessage: string | null }; + if (r.status === "success") { + return { + kind: "replayed", + value: r.value as T, + status: "success", + errorMessage: null, + }; + } + return { + kind: "replayed", + value: r.value ?? null, + status: r.status as "error" | "timeout", + errorMessage: r.errorMessage, + }; + } + return { + kind: "replayed", + value: result as T, + status: "success", + errorMessage: null, + }; + } catch { + return { kind: "unknown" }; + } + } + if (active.status === "rejected") { + // Rejected in-process promise — handler failed, return unknown + return { kind: "unknown" }; + } + } + // Also check DB — another process may have finalized + const row = readRow(db, claim.executionId); + if (row && row.status !== "running") { + if (row.status === "success" || row.status === "error" || row.status === "timeout") { + let parsedOutput: unknown = null; + if (row.output !== null) { + try { + parsedOutput = JSON.parse(row.output); + } catch { + parsedOutput = row.output; + } + } + if (row.status === "success") { + return { + kind: "replayed", + value: parsedOutput as T, + status: "success", + errorMessage: null, + }; + } + return { + kind: "replayed", + value: parsedOutput ?? null, + status: row.status as "error" | "timeout", + errorMessage: row.error_message, + }; + } + return { kind: "unknown" }; + } + await sleep(POLL_INTERVAL_MS); + } + return { kind: "in_progress" }; + } + + case "unknown": + return { kind: "unknown" }; + + case "identity_conflict": + return { kind: "identity_conflict" }; + + default: + return { kind: "unknown" }; + } +} diff --git a/src/lib/skills/toolLoopTypes.ts b/src/lib/skills/toolLoopTypes.ts new file mode 100644 index 0000000000..5ebf575518 --- /dev/null +++ b/src/lib/skills/toolLoopTypes.ts @@ -0,0 +1,200 @@ +/** + * Shared types for the server-owned tool loop. + * All consumers use `import type` — no runtime imports. + */ + +// ─── §5.4 Provider Leg ───────────────────────────────────────────────────── + +export interface ProviderLegUsage { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; + cached_tokens?: number; + cache_read_input_tokens?: number; + cache_creation_input_tokens?: number; + reasoning_tokens?: number; + cost_in_usd_ticks?: number; +} + +export interface ProviderLegReceipt { + index: number; + connectionId: string; + provider: string; + model: string; + startedAt: string; + endedAt: string; + latencyMs: number; + httpStatus: number; + errorType: string | null; + usage: ProviderLegUsage | null; + serviceTier: string | null; + computedCostUsd: number | null; + toolCalls: Array<{ id: string; name: string }>; + termination: string; + clientVisible: boolean; +} + +export interface ChatCoreErrorResult { + success: false; + status: number; + response: Response; + error?: string; + errorCode?: string; + errorType?: string; + retryAfterMs?: number; + originalError?: unknown; + rawMessage?: string; +} + +export type NonStreamingProviderLegResult = + | { + kind: "ok"; + response: Record; + responseForMemoryExtraction: Record; + providerBody: Record; + providerRequest: Record; + usage: ProviderLegUsage | null; + responsePayloadFormat: string; + looksLikeSSE: boolean; + connectionId: string; + headers: Headers; + requestHeaders?: Record; + requestUrl?: string; + upstreamResponse?: Response; + receipt: ProviderLegReceipt; + } + | { + kind: "error"; + result: ChatCoreErrorResult; + receipt: ProviderLegReceipt; + usage: ProviderLegUsage | null; + }; + +// ─── §5.5 Tool Loop ──────────────────────────────────────────────────────── + +export interface ServerOwnedToolLoopOptions { + initialLeg: NonStreamingProviderLegResult & { kind: "ok" }; + sourceBody: Record; + sourceFormat: "openai" | "claude"; + skillsModelId: string; + executionContext: ExecutionContext; + abortSignal?: AbortSignal; + now?: () => number; + executeServerOwned: ( + calls: ToolCall[], + context: ExecutionContext + ) => Promise; + resumeUpstream: ( + nextSourceBody: Record, + expectedConnectionId: string, + deadlineAtMs: number + ) => Promise; + maxFollowUps?: number; + maxResultBytes?: number; + maxTotalResultBytes?: number; + deadlineAtMs: number; +} + +export interface ServerOwnedToolLoopResult { + kind: "ok" | "error"; + response?: Record; + responseForMemoryExtraction?: Record; + finalProviderBody?: Record; + finalProviderRequest?: Record; + errorResult?: ChatCoreErrorResult; + cumulativeUsage: ProviderLegUsage | null; + totalCostUsd: number; + receipts: ProviderLegReceipt[]; + followUps: number; + termination: + | "completed" + | "client_tools" + | "mixed_tools" + | "max_followups" + | "tool_output_budget" + | "deadline" + | "client_abort" + | "provider_error" + | "connection_mismatch" + | "execution_in_progress" + | "execution_unknown" + | "execution_identity_conflict" + | "execution_error" + | "execution_timeout"; +} + +// ─── §5.1 Shared Context ─────────────────────────────────────────────────── + +export interface ToolCall { + id: string; + name: string; + arguments: Record; +} + +export interface ExecutionContext { + apiKeyId: string; + sessionId: string; + requestId: string; + requestIdentity?: string; + builtinToolNames?: string[]; + injectedCustomSkillNames?: string[]; + customSkillExecutionEnabled?: boolean; + executionFenceEnabled?: boolean; + provider?: string; + model?: string; +} + +export interface ExecutedToolResult { + id: string; + name: string; + result: unknown; + replayed: boolean; +} + +// ─── §5.2 Transcript Builder ─────────────────────────────────────────────── + +export interface BuildFollowUpTranscriptInput { + sourceBody: Record; + previousResponse: Record; + toolCalls: ToolCall[]; + results: ExecutedToolResult[]; + sourceFormat: "openai" | "claude"; + maxResultBytes: number; + maxTotalResultBytes?: number; + serializedResultTextById?: Map; +} + +export interface BoundedToolResult { + text: string; + truncated: boolean; + originalBytes: number; +} + +// ─── §5.3 Client Translate ───────────────────────────────────────────────── + +export interface NonStreamingClientTranslateInput { + responseBody: Record; + responsePayloadFormat: string; + clientResponseFormat: string; + sourceFormat: string; + provider: string; + model: string; + requestBody: Record; + /** + * Transcript used for no-tool_calls reasoning replay (#1628). + * Must be the client-translated Chat `messages` (parent: `translatedBody.messages`), + * not `finalBody` — Responses-shaped `finalBody` has `input`, not `messages`. + */ + historyMessages?: unknown[] | null; + responseToolNameMap: Map | null; + requestToolIdentityMap: Map | null; + reasoningCacheScope: string | null; + clientHeaders: Headers | Record | null; + isClaudeCodeCompatible: boolean; + phase: "intermediate" | "final"; +} + +export interface NonStreamingClientTranslateResult { + response: Record; + responseForMemoryExtraction: Record; +} diff --git a/src/shared/constants/featureFlagDefinitions.ts b/src/shared/constants/featureFlagDefinitions.ts index a73d4b7a25..88ace23020 100644 --- a/src/shared/constants/featureFlagDefinitions.ts +++ b/src/shared/constants/featureFlagDefinitions.ts @@ -558,6 +558,18 @@ export const FEATURE_FLAG_DEFINITIONS: FeatureFlagDefinition[] = [ requiresRestart: false, warningLevel: "info", }, + { + key: "SERVER_OWNED_TOOL_LOOP_ENABLED", + label: "Server-Owned Tool Loop", + description: + "Continue non-streaming server-owned tool calls until the model returns a client-usable response.", + descriptionI18nKey: "featureFlagServerOwnedToolLoopDescription", + category: "runtime", + defaultValue: "false", + type: "boolean", + requiresRestart: false, + warningLevel: "caution", + }, // ──────────────── CLI (5) ──────────────── { diff --git a/src/shared/utils/featureFlags.ts b/src/shared/utils/featureFlags.ts index 84f44f9565..22bb1caed8 100644 --- a/src/shared/utils/featureFlags.ts +++ b/src/shared/utils/featureFlags.ts @@ -172,3 +172,17 @@ export function isNetworkRotationSharedEgressGuardEnabled(): boolean { return true; } } + +export function isServerOwnedToolLoopEnabled( + reader: (key: string) => boolean = isFeatureFlagEnabled +): boolean { + try { + return reader("SERVER_OWNED_TOOL_LOOP_ENABLED"); + } catch (error) { + console.error( + "[featureFlags] Failed to resolve SERVER_OWNED_TOOL_LOOP_ENABLED, defaulting to disabled:", + error instanceof Error ? error.message : error + ); + return false; + } +} diff --git a/stryker.conf.json b/stryker.conf.json index c282b84214..6f074e038c 100644 --- a/stryker.conf.json +++ b/stryker.conf.json @@ -238,6 +238,7 @@ "tests/unit/cursor-renewal.test.ts", "tests/unit/custom-model-target-format.test.ts", "tests/unit/db-reset-module-state.test.ts", + "tests/unit/db-server-tool-executions-migration.test.ts", "tests/unit/db/stats-dbstat-optional.test.ts", "tests/unit/ddg-circuit-breaker-null-content-6999-7000.test.ts", "tests/unit/domain-persistence.test.ts", @@ -254,6 +255,7 @@ "tests/unit/executor-devin-cli-agentic-acp.test.ts", "tests/unit/executor-web-cookie-sweep.test.ts", "tests/unit/false-terminal-401-quota.test.ts", + "tests/unit/follow-up-transcript.test.ts", "tests/unit/format-provider-error-cause.test.ts", "tests/unit/forwarded-header-budget.test.ts", "tests/unit/fusion-vision-panel-3378.test.ts", @@ -297,6 +299,8 @@ "tests/unit/no-memory-header.test.ts", "tests/unit/noauth-autocombo-lockout-7623.test.ts", "tests/unit/ollama-404-model-lockout-11071.test.ts", + "tests/unit/non-streaming-client-translate.test.ts", + "tests/unit/non-streaming-provider-leg.test.ts", "tests/unit/non-streaming-sse-terminal-typescan-4459.test.ts", "tests/unit/nvidia-410-model-scope.test.ts", "tests/unit/nvidia-passthrough-models-6773.test.ts", @@ -324,6 +328,7 @@ "tests/unit/probe-testall-isolation.test.ts", "tests/unit/provider-breaker-halfopen-recovery.test.ts", "tests/unit/provider-error-rules.test.ts", + "tests/unit/provider-execution-pipeline.test.ts", "tests/unit/provider-health-matrix.test.ts", "tests/unit/provider-request-failure-pipeline.test.ts", "tests/unit/providers-route-codex-account-pool.test.ts", @@ -373,6 +378,10 @@ "tests/unit/serial/combo-quota-share-cooldown-wait-timing.test.ts", "tests/unit/serial/combo-strategy-fallbacks-half-open-timing.test.ts", "tests/unit/serial/provider-health-autopilot.test.ts", + "tests/unit/server-owned-tool-loop-flag.test.ts", + "tests/unit/server-owned-tool-loop.test.ts", + "tests/unit/skill-execution-fence.test.ts", + "tests/unit/skills-interception-server-owned.test.ts", "tests/unit/service-combo-metrics.test.ts", "tests/unit/service-provider-plugin-registry.test.ts", "tests/unit/services-branch-hardening.test.ts", @@ -386,6 +395,7 @@ "tests/unit/sse-auth-exclusive-leases.test.ts", "tests/unit/sse-auth-resource-404.test.ts", "tests/unit/sse-auth.test.ts", + "tests/unit/stable-json.test.ts", "tests/unit/stream-early-eof-breaker.test.ts", "tests/unit/stream-readiness.test.ts", "tests/unit/strict-random-deck.test.ts", @@ -396,6 +406,7 @@ "tests/unit/thundering-herd.test.ts", "tests/unit/token-refresh-race-comprehensive.test.ts", "tests/unit/token-refresh-service.test.ts", + "tests/unit/tool-loop-usage.test.ts", "tests/unit/tools-filter-anthropic-format.test.ts", "tests/unit/tproxy-route.test.ts", "tests/unit/trae-publiccred.test.ts", diff --git a/tests/integration/chat-pipeline.test.ts b/tests/integration/chat-pipeline.test.ts index 2b0bf96d33..191462d93c 100644 --- a/tests/integration/chat-pipeline.test.ts +++ b/tests/integration/chat-pipeline.test.ts @@ -1269,6 +1269,10 @@ test("chat pipeline returns current no-credentials contract when no provider con test("chat pipeline surfaces upstream 500 responses as structured errors", async () => { await seedConnection("openai", { apiKey: "sk-openai-500" }); + await settingsDb.updateSettings({ + requestRetry: 0, + maxRetryIntervalSec: 0, + }); globalThis.fetch = async () => new Response(JSON.stringify({ error: { message: "provider exploded" } }), { diff --git a/tests/integration/server-owned-tool-loop-pipeline.test.ts b/tests/integration/server-owned-tool-loop-pipeline.test.ts new file mode 100644 index 0000000000..7f0058dcba --- /dev/null +++ b/tests/integration/server-owned-tool-loop-pipeline.test.ts @@ -0,0 +1,259 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { encodeSkillToolName } from "../../src/lib/skills/injection.ts"; + +import { createChatPipelineHarness } from "./_chatPipelineHarness.ts"; + +// Split out of skills-pipeline.test.ts (#12867): the three server-owned tool loop +// cases pushed that file to 1338 lines, past the 1200-line test cap. Same harness, +// its own instance so the two files stay independently runnable. +const harness = await createChatPipelineHarness("server-owned-tool-loop-pipeline"); +const { + BaseExecutor, + buildOpenAIResponse, + buildOpenAIToolCallResponse, + buildRequest, + handleChat, + resetStorage, + seedApiKey, + seedConnection, + settingsDb, + skillExecutor, + skillRegistry, +} = harness; + +test.beforeEach(async () => { + BaseExecutor.RETRY_CONFIG.delayMs = 0; + await resetStorage(); +}); + +test.afterEach(async () => { + BaseExecutor.RETRY_CONFIG.delayMs = harness.originalRetryDelayMs; + await resetStorage(); +}); + +test.after(async () => { + await harness.cleanup(); +}); + +async function enableSkills() { + await settingsDb.updateSettings({ skillsEnabled: true }); +} + +async function registerSkill({ + apiKeyId, + name, + version = "1.0.0", + handler, + enabled = true, + description = "Test skill", + mode, + tags, + installCount, +}) { + return skillRegistry.register({ + apiKeyId, + name, + version, + description, + schema: { + input: { + type: "object", + properties: { + location: { type: "string" }, + path: { type: "string" }, + }, + }, + output: { + type: "object", + }, + }, + handler, + enabled, + mode, + tags, + installCount, + }); +} + +test("server-owned tool loop completes end-to-end for OpenAI without returning tool_results (issue #12696)", async () => { + await seedConnection("openai", { apiKey: "test-openai-key" }); + const apiKey = await seedApiKey(); + await enableSkills(); + + skillExecutor.registerHandler("weather-handler-loop-openai", async (input) => ({ + forecast: `Sunny in ${input.location}`, + })); + await registerSkill({ + apiKeyId: apiKey.id, + name: "lookupWeather", + handler: "weather-handler-loop-openai", + }); + + const prevFlag = process.env.SERVER_OWNED_TOOL_LOOP_ENABLED; + process.env.SERVER_OWNED_TOOL_LOOP_ENABLED = "true"; + + const fetchCalls: Array<{ url: string; body: Record }> = []; + globalThis.fetch = async (input, init) => { + const bodyStr = typeof init?.body === "string" ? init.body : "{}"; + const body = JSON.parse(bodyStr); + fetchCalls.push({ url: String(input), body }); + + if (fetchCalls.length === 1) { + return buildOpenAIToolCallResponse({ + toolCallId: "call_weather_loop_1", + toolName: encodeSkillToolName("lookupWeather", "1.0.0"), + argumentsObject: { location: "Tokyo" }, + }); + } + + return buildOpenAIResponse("The weather in Tokyo is 18C and sunny.", "gpt-4o-mini"); + }; + + try { + const response = await handleChat( + buildRequest({ + authKey: apiKey.key, + body: { + model: "openai/gpt-4o-mini", + stream: false, + messages: [{ role: "user", content: "What is the weather in Tokyo?" }], + }, + }) + ); + + assert.equal(response.status, 200); + const json = (await response.json()) as Record; + + assert.equal(fetchCalls.length, 2, "should have dispatched 2 provider legs"); + assert.equal( + json.tool_results, + undefined, + "gateway should not return tool_results when loop is enabled" + ); + assert.equal(json.choices[0].finish_reason, "stop"); + assert.equal(json.choices[0].message.content, "The weather in Tokyo is 18C and sunny."); + assert.equal(fetchCalls[1].body.messages.length, 3); + assert.equal(fetchCalls[1].body.messages[1].role, "assistant"); + assert.equal(fetchCalls[1].body.messages[2].role, "tool"); + assert.equal(fetchCalls[1].body.messages[2].tool_call_id, "call_weather_loop_1"); + } finally { + process.env.SERVER_OWNED_TOOL_LOOP_ENABLED = prevFlag; + } +}); + +test("server-owned tool loop follow-up failure propagates error cleanly", async () => { + await seedConnection("openai", { apiKey: "test-openai-key" }); + const apiKey = await seedApiKey(); + await enableSkills(); + + skillExecutor.registerHandler("weather-handler-loop-fail", async (input) => ({ + forecast: `Sunny in ${input.location}`, + })); + await registerSkill({ + apiKeyId: apiKey.id, + name: "lookupWeather", + handler: "weather-handler-loop-fail", + }); + + const prevFlag = process.env.SERVER_OWNED_TOOL_LOOP_ENABLED; + process.env.SERVER_OWNED_TOOL_LOOP_ENABLED = "true"; + + let callCount = 0; + globalThis.fetch = async () => { + callCount++; + if (callCount === 1) { + return buildOpenAIToolCallResponse({ + callId: "call_weather_fail_1", + toolName: encodeSkillToolName("lookupWeather", "1.0.0"), + argumentsObject: { location: "Tokyo" }, + }); + } + return new Response(JSON.stringify({ error: { message: "Rate limit exceeded" } }), { + status: 429, + headers: { "content-type": "application/json" }, + }); + }; + + try { + const response = await handleChat( + buildRequest({ + authKey: apiKey.key, + body: { + model: "openai/gpt-4o-mini", + stream: false, + messages: [{ role: "user", content: "What is the weather in Tokyo?" }], + }, + }) + ); + + assert.ok(callCount >= 2, `expected at least 2 fetch calls, got ${callCount}`); + assert.equal(response.status, 429); + } finally { + process.env.SERVER_OWNED_TOOL_LOOP_ENABLED = prevFlag; + } +}); + +test("server-owned tool loop completes end-to-end for Claude messages client without returning tool_results", async () => { + await seedConnection("openai", { apiKey: "test-openai-key" }); + const apiKey = await seedApiKey(); + await enableSkills(); + + skillExecutor.registerHandler("weather-handler-loop-claude-client", async (input) => ({ + forecast: `Sunny in ${input.location}`, + })); + await registerSkill({ + apiKeyId: apiKey.id, + name: "lookupWeather", + handler: "weather-handler-loop-claude-client", + }); + + const prevFlag = process.env.SERVER_OWNED_TOOL_LOOP_ENABLED; + process.env.SERVER_OWNED_TOOL_LOOP_ENABLED = "true"; + + const fetchCalls: Array<{ url: string; body: Record }> = []; + globalThis.fetch = async (input, init) => { + const bodyStr = typeof init?.body === "string" ? init.body : "{}"; + const body = JSON.parse(bodyStr); + fetchCalls.push({ url: String(input), body }); + + if (fetchCalls.length === 1) { + return buildOpenAIToolCallResponse({ + toolCallId: "call_weather_claude_1", + toolName: encodeSkillToolName("lookupWeather", "1.0.0"), + argumentsObject: { location: "Tokyo" }, + }); + } + + return buildOpenAIResponse("The weather in Tokyo is 18C and sunny.", "gpt-4o-mini"); + }; + + try { + const response = await handleChat( + buildRequest({ + url: "http://localhost/v1/messages", + authKey: apiKey.key, + body: { + model: "openai/gpt-4o-mini", + stream: false, + max_tokens: 256, + messages: [ + { role: "user", content: [{ type: "text", text: "What is the weather in Tokyo?" }] }, + ], + }, + }) + ); + + assert.equal(response.status, 200); + const json = (await response.json()) as Record; + + assert.equal(fetchCalls.length, 2, "should have dispatched 2 provider legs"); + assert.equal(json.type, "message"); + assert.equal(json.role, "assistant"); + assert.equal(json.stop_reason, "end_turn"); + assert.equal(json.content[0].text, "The weather in Tokyo is 18C and sunny."); + assert.equal(json.tool_results, undefined); + } finally { + process.env.SERVER_OWNED_TOOL_LOOP_ENABLED = prevFlag; + } +}); diff --git a/tests/unit/antigravity-byop-account-rotation.test.ts b/tests/unit/antigravity-byop-account-rotation.test.ts index e73b22c184..b044ee9d4a 100644 --- a/tests/unit/antigravity-byop-account-rotation.test.ts +++ b/tests/unit/antigravity-byop-account-rotation.test.ts @@ -186,6 +186,117 @@ test("Antigravity BYOP 422 rotates to a sibling account and the request succeeds } }); +test("streaming Antigravity BYOP 422 still rotates — execute must not cancel the error body", async () => { + const byopAccount = await createAntigravityAccount({ + name: "antigravity-byop-stream-a", + email: "byop-stream-a@example.test", + accessToken: "fake-byop-account-a-token", + refreshToken: "fake-byop-account-a-refresh", + priority: 1, + }); + const healthyAccount = await createAntigravityAccount({ + name: "antigravity-healthy-stream-b", + email: "byop-stream-b@example.test", + accessToken: "fake-healthy-account-b-token", + refreshToken: "fake-healthy-account-b-refresh", + priority: 2, + }); + + let onboardCallsForA = 0; + const modelCalls: Array<{ token: string }> = []; + const originalFetch = globalThis.fetch; + globalThis.fetch = async (input, init) => { + const request = input instanceof Request ? input : new Request(input, init); + if (request.url.startsWith("https://oauth2.googleapis.com/token")) { + const form = await request.text().catch(() => ""); + const refreshMatch = form.match(/refresh_token=([^&]+)/); + const refreshToken = refreshMatch ? decodeURIComponent(refreshMatch[1]) : ""; + const accessToken = + refreshToken === "fake-byop-account-a-refresh" + ? "fake-byop-account-a-token" + : "fake-healthy-account-b-token"; + return new Response(JSON.stringify({ access_token: accessToken, expires_in: 3600 }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (request.url.endsWith(":loadCodeAssist")) { + const token = (request.headers.get("authorization") || "").replace(/^Bearer\s+/i, ""); + if (token === "fake-healthy-account-b-token") { + return new Response( + JSON.stringify({ cloudaicompanionProject: "projects/healthy-b-project" }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + } + return new Response("{}", { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (request.url.endsWith(":onboardUser")) { + const token = (request.headers.get("authorization") || "").replace(/^Bearer\s+/i, ""); + if (token === "fake-byop-account-a-token") { + onboardCallsForA += 1; + return new Response(JSON.stringify({ done: true }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response( + JSON.stringify({ + done: true, + cloudaicompanionProject: { name: "projects/healthy-b-project" }, + }), + { status: 200, headers: { "Content-Type": "application/json" } } + ); + } + if (new URL(request.url).hostname === "cloudcode-pa.googleapis.com") { + modelCalls.push({ + token: (request.headers.get("authorization") || "").replace(/^Bearer\s+/i, ""), + }); + return new Response( + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"ok from account B"}]},"finishReason":"STOP"}]}}\n\n', + { status: 200, headers: { "Content-Type": "text/event-stream" } } + ); + } + throw new Error(`Unexpected external fetch: ${request.url}`); + }; + + try { + const response = await handleChat( + buildRequest({ + body: { + model: "antigravity/gemini-2.5-flash", + stream: true, + messages: [{ role: "user", content: "hello" }], + }, + }) + ); + + assert.equal(response.status, 200); + const bodyText = await response.text().catch(() => ""); + assert.match(bodyText, /ok from account B/); + assert.ok(modelCalls.length >= 1, "model call should have been made"); + assert.equal(modelCalls[0].token, "fake-healthy-account-b-token"); + assert.equal(onboardCallsForA, 1); + const updatedA = await providersDb.getProviderConnectionById(byopAccount.id); + assert.ok( + updatedA && Number(updatedA.rateLimitedUntil) > Date.now(), + "BYOP account should be excluded from selection" + ); + const updatedB = await providersDb.getProviderConnectionById(healthyAccount.id); + assert.ok( + !updatedB || + !Number(updatedB.rateLimitedUntil) || + Number(updatedB.rateLimitedUntil) <= Date.now(), + "healthy sibling account must not be excluded" + ); + } finally { + globalThis.fetch = originalFetch; + clearAntigravityProjectCache(); + } +}); + test("Antigravity BYOP with no sibling account surfaces the actionable 422 and excludes the connection", async () => { const byopAccount = await createAntigravityAccount({ name: "antigravity-byop-only", diff --git a/tests/unit/chatcore-failure-usage.test.ts b/tests/unit/chatcore-failure-usage.test.ts index cc20cf7824..17918db5b0 100644 --- a/tests/unit/chatcore-failure-usage.test.ts +++ b/tests/unit/chatcore-failure-usage.test.ts @@ -64,16 +64,74 @@ test("applies the unknown/undefined fallbacks", () => { test("combo strategy is included only for combo requests", () => { const combo = buildFailureUsageRecord({ - provider: "x", model: "y", connectionId: null, apiKeyInfo: null, - effectiveServiceTier: "standard", isCombo: true, comboStrategy: "round-robin", - statusCode: 500, errorCode: "boom", latencyMs: 1, + provider: "x", + model: "y", + connectionId: null, + apiKeyInfo: null, + effectiveServiceTier: "standard", + isCombo: true, + comboStrategy: "round-robin", + statusCode: 500, + errorCode: "boom", + latencyMs: 1, }); assert.equal(combo.comboStrategy, "round-robin"); const comboNoStrategy = buildFailureUsageRecord({ - provider: "x", model: "y", connectionId: null, apiKeyInfo: null, - effectiveServiceTier: "standard", isCombo: true, comboStrategy: null, - statusCode: 500, errorCode: "boom", latencyMs: 1, + provider: "x", + model: "y", + connectionId: null, + apiKeyInfo: null, + effectiveServiceTier: "standard", + isCombo: true, + comboStrategy: null, + statusCode: 500, + errorCode: "boom", + latencyMs: 1, }); assert.equal(comboNoStrategy.comboStrategy, undefined); }); + +test("maps aggregate usage onto failure tokens instead of zeros", () => { + const r = buildFailureUsageRecord({ + provider: "openai", + model: "gpt-4o", + connectionId: "conn-1", + apiKeyInfo: { id: "key-1", name: "My Key" }, + effectiveServiceTier: "priority", + isCombo: false, + comboStrategy: null, + statusCode: 429, + errorCode: "rate_limited", + latencyMs: 50, + aggregate: { + prompt_tokens: 100, + completion_tokens: 20, + cache_read_input_tokens: 10, + reasoning_tokens: 5, + }, + }); + assert.deepEqual(r.tokens, { + input: 100, + output: 20, + cacheRead: 10, + cacheCreation: 0, + reasoning: 5, + }); +}); + +test("keeps zeroed tokens when aggregate is absent", () => { + const r = buildFailureUsageRecord({ + provider: "openai", + model: "gpt-4o", + connectionId: null, + apiKeyInfo: null, + effectiveServiceTier: "standard", + isCombo: false, + comboStrategy: null, + statusCode: 502, + errorCode: null, + latencyMs: 7, + }); + assert.deepEqual(r.tokens, { input: 0, output: 0, cacheRead: 0, cacheCreation: 0, reasoning: 0 }); +}); diff --git a/tests/unit/chatcore-memory-skills-injection.test.ts b/tests/unit/chatcore-memory-skills-injection.test.ts index e9991e6d52..08126c9ee5 100644 --- a/tests/unit/chatcore-memory-skills-injection.test.ts +++ b/tests/unit/chatcore-memory-skills-injection.test.ts @@ -9,12 +9,24 @@ import path from "node:path"; const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-mem-skills-")); process.env.DATA_DIR = TEST_DATA_DIR; -const { getSkillsProviderForFormat, injectMemoryAndSkills, sortToolsByName } = - await import("../../open-sse/handlers/chatCore/memorySkillsInjection.ts"); +const { + getSkillsProviderForFormat, + injectMemoryAndSkills, + sortToolsByName, + mergeInjectedFallbackOwnerNames, +} = await import("../../open-sse/handlers/chatCore/memorySkillsInjection.ts"); const { FORMATS } = await import("../../open-sse/translator/formats.ts"); const core = await import("../../src/lib/db/core.ts"); +const { skillRegistry } = await import("../../src/lib/skills/registry.ts"); + +function resetSkillsRegistry() { + skillRegistry["registeredSkills"].clear(); + skillRegistry["versionCache"].clear(); + skillRegistry.invalidateCache(); +} test.after(() => { + resetSkillsRegistry(); core.resetDbInstance(); fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); }); @@ -258,3 +270,334 @@ test("injectMemoryAndSkills does not inject memory tools when memory is disabled invalidateMemorySettingsCache(); }); + +// ─── Task 3: owner-set provenance + stream gate RED tests ──────────────────── + +test("stream:true + skills enabled + registry has items → no custom skill tool injected, injectedCustomSkillNames=[]", async () => { + const { updateSettings } = await import("../../src/lib/db/settings.ts"); + const { invalidateMemorySettingsCache: inv2 } = await import("../../src/lib/memory/settings.ts"); + + await updateSettings({ memoryEnabled: true, memoryMaxTokens: 2000, skillsEnabled: true }); + inv2(); + resetSkillsRegistry(); + + await skillRegistry.register({ + name: "test-skill", + version: "1.0.0", + description: "test skill for stream gate", + schema: { input: {}, output: {} }, + handler: "test-handler", + enabled: true, + apiKeyId: "owner-stream-skills", + mode: "on", + }); + + const body: Record = { + model: "gpt-4o", + stream: true, + messages: [{ role: "user", content: "hello" }], + }; + + const result = await injectMemoryAndSkills({ + body, + memoryOwnerId: "owner-stream-skills", + provider: "openai", + effectiveModel: "gpt-4o", + sourceFormat: FORMATS.OPENAI, + targetFormat: FORMATS.OPENAI, + backgroundReason: null, + log: { debug: () => {} }, + }); + + const toolNames = ( + (result.body.tools as { function?: { name?: string }; name?: string }[] | undefined) ?? [] + ).map((t) => t.function?.name ?? t.name); + const hasCustomSkill = toolNames.some( + (n) => typeof n === "string" && (n.includes("test-skill") || n.startsWith("omr_skill_")) + ); + assert.equal(hasCustomSkill, false, "stream:true must not inject custom skill tools"); + + assert.deepEqual( + (result as Record).injectedCustomSkillNames, + [], + "injectedCustomSkillNames must be empty for stream requests" + ); + + resetSkillsRegistry(); + inv2(); +}); + +test("memory actual injection → builtinToolNames equals exactly the newly added memory tool names", async () => { + const { updateSettings } = await import("../../src/lib/db/settings.ts"); + const { invalidateMemorySettingsCache: inv3 } = await import("../../src/lib/memory/settings.ts"); + const { MEMORY_BUILTIN_TOOL_NAMES } = await import("../../src/lib/skills/memoryBuiltins.ts"); + + await updateSettings({ memoryEnabled: true, memoryMaxTokens: 2000 }); + inv3(); + resetSkillsRegistry(); + + const body: Record = { + model: "gpt-4o", + messages: [{ role: "user", content: "hello" }], + tools: [{ type: "function", function: { name: "some_client_tool", description: "x" } }], + }; + + const result = await injectMemoryAndSkills({ + body, + memoryOwnerId: "owner-builtin-own", + provider: "openai", + effectiveModel: "gpt-4o", + sourceFormat: FORMATS.OPENAI, + targetFormat: FORMATS.OPENAI, + backgroundReason: null, + log: { debug: () => {} }, + }); + + const builtinToolNames = (result as Record).builtinToolNames as + string[] | undefined; + assert.ok(builtinToolNames, "builtinToolNames must be present in result"); + + const expectedNewMemoryNames = [...MEMORY_BUILTIN_TOOL_NAMES]; + assert.deepEqual( + builtinToolNames.sort(), + expectedNewMemoryNames.sort(), + "builtinToolNames must equal exactly the newly added memory tool names" + ); + + resetSkillsRegistry(); + inv3(); +}); + +test("client already has memory_search → not injected, not in builtinToolNames", async () => { + const { updateSettings } = await import("../../src/lib/db/settings.ts"); + const { invalidateMemorySettingsCache: inv4 } = await import("../../src/lib/memory/settings.ts"); + const { MEMORY_SEARCH_TOOL_NAME } = await import("../../src/lib/skills/memoryBuiltins.ts"); + + await updateSettings({ memoryEnabled: true, memoryMaxTokens: 2000 }); + inv4(); + resetSkillsRegistry(); + + const body: Record = { + model: "gpt-4o", + messages: [{ role: "user", content: "hello" }], + tools: [ + { + type: "function", + function: { name: MEMORY_SEARCH_TOOL_NAME, description: "client memory" }, + }, + ], + }; + + const result = await injectMemoryAndSkills({ + body, + memoryOwnerId: "owner-client-mem", + provider: "openai", + effectiveModel: "gpt-4o", + sourceFormat: FORMATS.OPENAI, + targetFormat: FORMATS.OPENAI, + backgroundReason: null, + log: { debug: () => {} }, + }); + + const toolNames = ( + (result.body.tools as { function?: { name?: string }[] | undefined }) ?? [] + ).map((t: { function?: { name?: string } }) => t.function?.name); + + const memorySearchCount = toolNames.filter((n) => n === MEMORY_SEARCH_TOOL_NAME).length; + assert.equal(memorySearchCount, 1, "only one memory_search (client's) must exist"); + + const builtinToolNames = (result as Record).builtinToolNames as + string[] | undefined; + assert.ok(builtinToolNames, "builtinToolNames must be present"); + assert.equal( + builtinToolNames.includes(MEMORY_SEARCH_TOOL_NAME), + false, + "client-owned memory_search must NOT be in builtinToolNames" + ); + + resetSkillsRegistry(); + inv4(); +}); + +test("custom skill client collision: client has same encoded skill name → not injected, not in injectedCustomSkillNames", async () => { + const { updateSettings } = await import("../../src/lib/db/settings.ts"); + const { invalidateMemorySettingsCache: inv5 } = await import("../../src/lib/memory/settings.ts"); + const { encodeSkillToolName } = await import("../../src/lib/skills/injection.ts"); + + await updateSettings({ memoryEnabled: true, memoryMaxTokens: 2000, skillsEnabled: true }); + inv5(); + resetSkillsRegistry(); + + await skillRegistry.register({ + name: "collision-skill", + version: "1.0.0", + description: "skill that collides", + schema: { input: {}, output: {} }, + handler: "collision-handler", + enabled: true, + apiKeyId: "owner-collision", + mode: "on", + }); + + const encodedName = encodeSkillToolName("collision-skill", "1.0.0"); + + const body: Record = { + model: "gpt-4o", + messages: [{ role: "user", content: "hello" }], + tools: [{ type: "function", function: { name: encodedName, description: "client collision" } }], + }; + + const result = await injectMemoryAndSkills({ + body, + memoryOwnerId: "owner-collision", + provider: "openai", + effectiveModel: "gpt-4o", + sourceFormat: FORMATS.OPENAI, + targetFormat: FORMATS.OPENAI, + backgroundReason: null, + log: { debug: () => {} }, + }); + + const toolNames = ( + (result.body.tools as { function?: { name?: string }[] | undefined }) ?? [] + ).map((t: { function?: { name?: string } }) => t.function?.name); + + const count = toolNames.filter((n) => n === encodedName).length; + assert.equal(count, 1, "only one instance of encoded name must exist (client's)"); + + const injectedCustomSkillNames = (result as Record).injectedCustomSkillNames as + string[] | undefined; + assert.ok(injectedCustomSkillNames, "injectedCustomSkillNames must be present"); + assert.equal( + injectedCustomSkillNames.includes(encodedName), + false, + "client-owned skill name must NOT be in injectedCustomSkillNames" + ); + + resetSkillsRegistry(); + inv5(); +}); + +test("web-search fallback: client has same tool name → not added to builtinToolNames", async () => { + const { updateSettings } = await import("../../src/lib/db/settings.ts"); + const { invalidateMemorySettingsCache: inv6 } = await import("../../src/lib/memory/settings.ts"); + const { OMNIROUTE_WEB_SEARCH_FALLBACK_TOOL_NAME } = + await import("../../open-sse/services/webSearchFallback.ts"); + + await updateSettings({ memoryEnabled: true, memoryMaxTokens: 2000 }); + inv6(); + resetSkillsRegistry(); + + const body: Record = { + model: "gpt-4o", + messages: [{ role: "user", content: "hello" }], + tools: [ + { + type: "function", + function: { name: OMNIROUTE_WEB_SEARCH_FALLBACK_TOOL_NAME, description: "client search" }, + }, + ], + }; + + const result = await injectMemoryAndSkills({ + body, + memoryOwnerId: "owner-websearch", + provider: "openai", + effectiveModel: "gpt-4o", + sourceFormat: FORMATS.OPENAI, + targetFormat: FORMATS.OPENAI, + backgroundReason: null, + log: { debug: () => {} }, + }); + + const builtinToolNames = (result as Record).builtinToolNames as + string[] | undefined; + assert.ok(builtinToolNames, "builtinToolNames must be present"); + assert.equal( + builtinToolNames.includes(OMNIROUTE_WEB_SEARCH_FALLBACK_TOOL_NAME), + false, + "client-owned web search tool must NOT be in builtinToolNames" + ); + + resetSkillsRegistry(); + inv6(); +}); + +// ─── Fix Round 2: Defect 5 — mergeInjectedFallbackOwnerNames + provenance ─── + +test("mergeInjectedFallbackOwnerNames: adds name only when enabled=true, convertedToolCount>0, toolName non-null, and not already in client tools", () => { + const result = mergeInjectedFallbackOwnerNames({ builtinToolNames: ["memory_search"] }, [ + { enabled: true, toolName: "omniroute_web_search", convertedToolCount: 2 }, + { enabled: true, toolName: null, convertedToolCount: 1 }, + { enabled: false, toolName: "omniroute_web_fetch", convertedToolCount: 3 }, + { enabled: true, toolName: "omniroute_web_fetch", convertedToolCount: 0 }, + ]); + + assert.deepEqual(result.builtinToolNames, ["memory_search", "omniroute_web_search"]); +}); + +test("mergeInjectedFallbackOwnerNames: does not mutate input injectionResult", () => { + const input = { builtinToolNames: ["original"] }; + const plans = [{ enabled: true, toolName: "omniroute_web_search", convertedToolCount: 1 }]; + const result = mergeInjectedFallbackOwnerNames(input, plans); + + // input must be unchanged + assert.deepEqual(input.builtinToolNames, ["original"]); + // result is a new object + assert.notEqual(result, input); + assert.deepEqual(result.builtinToolNames, ["original", "omniroute_web_search"]); +}); + +test("mergeInjectedFallbackOwnerNames: skips name already present in pre-conversion client tools", () => { + const result = mergeInjectedFallbackOwnerNames({ builtinToolNames: ["omniroute_web_search"] }, [ + { enabled: true, toolName: "omniroute_web_search", convertedToolCount: 2 }, + ]); + + // Must not duplicate — omniroute_web_search already present + assert.deepEqual(result.builtinToolNames, ["omniroute_web_search"]); +}); + +// ─── Fix Round 3: Defect 3 — pre-conversion collision guard ───────────────── + +test("mergeInjectedFallbackOwnerNames: client has omniroute_web_search → not added to builtinToolNames even if enabled=true", () => { + // Scenario: client sends {type:"web_search"} plus function named omniroute_web_search. + // prepareWebSearchFallbackBody emits enabled=true, convertedToolCount=2 (from the + // builtin conversion) but the synthetic tool was NOT added because client already has it. + // mergeInjectedFallbackOwnerNames must check pre-conversion client names. + const result = mergeInjectedFallbackOwnerNames( + { builtinToolNames: [] }, + [{ enabled: true, toolName: "omniroute_web_search", convertedToolCount: 2 }], + ["omniroute_web_search"] + ); + + // Must NOT add omniroute_web_search — client already owns it + assert.deepEqual(result.builtinToolNames, []); +}); + +test("mergeInjectedFallbackOwnerNames: client has omniroute_web_fetch → not added to builtinToolNames", () => { + const result = mergeInjectedFallbackOwnerNames( + { builtinToolNames: [] }, + [{ enabled: true, toolName: "omniroute_web_fetch", convertedToolCount: 1 }], + ["omniroute_web_fetch"] + ); + + assert.deepEqual(result.builtinToolNames, []); +}); + +test("mergeInjectedFallbackOwnerNames: client does NOT have the fallback name → added to builtinToolNames", () => { + const result = mergeInjectedFallbackOwnerNames( + { builtinToolNames: [] }, + [{ enabled: true, toolName: "omniroute_web_search", convertedToolCount: 2 }], + ["some_other_tool"] + ); + + assert.deepEqual(result.builtinToolNames, ["omniroute_web_search"]); +}); + +test("mergeInjectedFallbackOwnerNames: no preConversionClientToolNames provided → falls back to existing behavior", () => { + const result = mergeInjectedFallbackOwnerNames({ builtinToolNames: [] }, [ + { enabled: true, toolName: "omniroute_web_search", convertedToolCount: 2 }, + ]); + + assert.deepEqual(result.builtinToolNames, ["omniroute_web_search"]); +}); diff --git a/tests/unit/chatcore-stream-error-result.test.ts b/tests/unit/chatcore-stream-error-result.test.ts index 3ebe5821e7..e125dcb216 100644 --- a/tests/unit/chatcore-stream-error-result.test.ts +++ b/tests/unit/chatcore-stream-error-result.test.ts @@ -67,3 +67,26 @@ test("getUpstreamErrorIdentifier returns a non-empty string code or undefined", assert.equal(getUpstreamErrorIdentifier(null), undefined); assert.equal(getUpstreamErrorIdentifier("ECONNRESET"), undefined); }); + +test("non-streaming runNonStreamingProviderLeg is inside a try that maps semaphore errors", async () => { + const fs = await import("node:fs"); + const src = fs.readFileSync("open-sse/handlers/chatCore.ts", "utf8"); + // `let`, not `const`, since 6077b9dd (#12867) made the finalization step reassign + // legResult. The guard is about the try/catch that wraps the call, not the keyword. + const idx = src.search(/(?:const|let) legResult = await runNonStreamingProviderLeg/); + assert.ok(idx >= 0, "non-streaming branch must exist"); + const start = src.lastIndexOf("if (!stream)", idx); + const end = src.indexOf("// Streaming response", idx); + assert.ok(start >= 0 && end > start, "non-stream block bounds"); + const block = src.slice(start, end); + assert.match( + block, + /try\s*\{[\s\S]*runNonStreamingProviderLeg/, + "non-stream leg must sit in a try so SEMAPHORE_TIMEOUT cannot escape handleChatCore" + ); + assert.match( + block, + /isSemaphoreCapacityError/, + "same catch that maps stream semaphore errors must cover the non-stream leg" + ); +}); diff --git a/tests/unit/db-server-tool-executions-migration.test.ts b/tests/unit/db-server-tool-executions-migration.test.ts new file mode 100644 index 0000000000..0291d952ae --- /dev/null +++ b/tests/unit/db-server-tool-executions-migration.test.ts @@ -0,0 +1,242 @@ +process.env.DISABLE_SQLITE_AUTO_BACKUP = "true"; + +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import path from "node:path"; +import os from "node:os"; +import { createRequire } from "node:module"; + +const require_ = createRequire(import.meta.url); +const BetterSqlite3 = require_("better-sqlite3") as typeof import("better-sqlite3"); + +import { createBetterSqliteAdapter } from "../../src/lib/db/adapters/betterSqliteAdapter"; +import type { SqliteAdapter } from "../../src/lib/db/adapters/types"; + +const MIGRATION_174_PATH = path.resolve( + import.meta.dirname ?? ".", + "../../src/lib/db/migrations/174_server_tool_executions.sql" +); +const MIGRATION_174_SQL = fs.readFileSync(MIGRATION_174_PATH, "utf8"); + +// Minimal pre-174 fixture: only skills + skill_executions tables. +// No SCHEMA_SQL import from core.ts, no runMigrations. +const PRE_174_FIXTURE = ` + CREATE TABLE IF NOT EXISTS skills ( + id TEXT PRIMARY KEY, + api_key_id TEXT NOT NULL, + name TEXT NOT NULL, + version TEXT NOT NULL DEFAULT '1.0.0', + description TEXT, + schema TEXT NOT NULL, + handler TEXT NOT NULL, + enabled INTEGER NOT NULL DEFAULT 1, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')) + ); + CREATE TABLE IF NOT EXISTS skill_executions ( + id TEXT PRIMARY KEY, + skill_id TEXT NOT NULL, + api_key_id TEXT NOT NULL, + session_id TEXT, + input TEXT NOT NULL, + output TEXT, + status TEXT NOT NULL CHECK(status IN ('pending', 'running', 'success', 'error', 'timeout')), + error_message TEXT, + duration_ms INTEGER, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + FOREIGN KEY (skill_id) REFERENCES skills(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_skill_executions_skill ON skill_executions(skill_id); + CREATE INDEX IF NOT EXISTS idx_skill_executions_api_key ON skill_executions(api_key_id); +`; + +function makeTempDb(): { + adapter: SqliteAdapter; + dir: string; + raw: import("better-sqlite3").Database; +} { + const dir = fs.mkdtempSync(path.join(os.tmpdir(), "migration-174-test-")); + const dbPath = path.join(dir, "test.db"); + const raw = new BetterSqlite3(dbPath); + raw.pragma("journal_mode = WAL"); + raw.pragma("busy_timeout = 2000"); + // Create migrations tracking table + raw.exec(` + CREATE TABLE IF NOT EXISTS _omniroute_migrations ( + version TEXT PRIMARY KEY, + name TEXT NOT NULL, + applied_at TEXT NOT NULL DEFAULT (datetime('now')) + ); + `); + // Apply minimal pre-174 fixture (skills + skill_executions only) + raw.exec(PRE_174_FIXTURE); + // Apply migration 174 via raw.exec (real SQL, twice for idempotency) + raw.exec(MIGRATION_174_SQL); + raw.exec(MIGRATION_174_SQL); + const adapter = createBetterSqliteAdapter(raw); + return { adapter, dir, raw }; +} + +function getColumnInfo(raw: import("better-sqlite3").Database, table: string) { + return raw.pragma(`table_info(${table})`) as Array<{ + name: string; + type: string; + notnull: number; + pk: number; + }>; +} + +function getIndexInfo(raw: import("better-sqlite3").Database, table: string) { + return raw.pragma(`index_list(${table})`) as Array<{ + name: string; + unique: number; + }>; +} + +// ── Migration 174 tests ── + +test("migration 174: server_tool_executions table exists with correct columns", (_, done) => { + const { raw, dir } = makeTempDb(); + try { + const tables = raw.pragma("table_list") as Array<{ name: string }>; + const names = tables.map((t) => t.name); + assert.ok( + names.includes("server_tool_executions"), + `Expected server_tool_executions in: ${names.join(", ")}` + ); + const cols = getColumnInfo(raw, "server_tool_executions"); + const colNames = cols.map((c) => c.name); + assert.ok(colNames.includes("id"), "must have id column"); + assert.ok(colNames.includes("api_key_id"), "must have api_key_id column"); + assert.ok(colNames.includes("request_identity"), "must have request_identity column"); + assert.ok(colNames.includes("tool_call_id"), "must have tool_call_id column"); + assert.ok(colNames.includes("tool_name"), "must have tool_name column"); + assert.ok(colNames.includes("input_digest"), "must have input_digest column"); + assert.ok(colNames.includes("status"), "must have status column"); + assert.ok(colNames.includes("claim_expires_at"), "must have claim_expires_at column"); + assert.ok(colNames.includes("duration_ms"), "must have duration_ms column"); + } finally { + raw.close(); + fs.rmSync(dir, { recursive: true, force: true }); + done(); + } +}); + +test("migration 174: UNIQUE constraint on (api_key_id, request_identity, tool_call_id)", (_, done) => { + const { raw, dir } = makeTempDb(); + try { + raw.exec(` + INSERT INTO server_tool_executions + (id, api_key_id, request_identity, tool_call_id, tool_name, input_digest, status, claim_expires_at) + VALUES ('e1','k1','r1','c1','tool_a','d1','running',datetime('now')) + `); + assert.throws(() => { + raw.exec(` + INSERT INTO server_tool_executions + (id, api_key_id, request_identity, tool_call_id, tool_name, input_digest, status, claim_expires_at) + VALUES ('e2','k1','r1','c1','tool_a','d1','running',datetime('now')) + `); + }, /UNIQUE/i); + } finally { + raw.close(); + fs.rmSync(dir, { recursive: true, force: true }); + done(); + } +}); + +test("migration 174: two indexes exist on server_tool_executions", (_, done) => { + const { raw, dir } = makeTempDb(); + try { + const indexes = getIndexInfo(raw, "server_tool_executions"); + const names = indexes.map((i) => i.name); + assert.ok( + names.some((n) => n.includes("status_expiry")), + `Expected status_expiry index, got: ${names.join(", ")}` + ); + assert.ok( + names.some((n) => n.includes("created")), + `Expected created index, got: ${names.join(", ")}` + ); + } finally { + raw.close(); + fs.rmSync(dir, { recursive: true, force: true }); + done(); + } +}); + +test("migration 174: existing skill_executions data preserved after migration", (_, done) => { + const { raw, dir } = makeTempDb(); + try { + // Insert a skill to satisfy FK + raw.exec(` + INSERT INTO skills (id, api_key_id, name, version, schema, handler) + VALUES ('s1','k1','test','1.0.0','{}','h.js') + `); + raw.exec(` + INSERT INTO skill_executions (id, skill_id, api_key_id, input, status) + VALUES ('old_exec','s1','k1','{"q":"test"}','success') + `); + const rows = raw.prepare("SELECT * FROM skill_executions WHERE id = 'old_exec'").all(); + assert.equal(rows.length, 1, "old row should exist after migration 174"); + } finally { + raw.close(); + fs.rmSync(dir, { recursive: true, force: true }); + done(); + } +}); + +test("migration 174: skill_executions still enforces skill_id NOT NULL", (_, done) => { + const { raw, dir } = makeTempDb(); + try { + assert.throws(() => { + raw.exec( + `INSERT INTO skill_executions (id, api_key_id, input, status) + VALUES ('bad','k1','{}','running')` + ); + }, /NOT NULL/i); + } finally { + raw.close(); + fs.rmSync(dir, { recursive: true, force: true }); + done(); + } +}); + +test("migration 174: custom skill execution write still works after migration", (_, done) => { + const { raw, dir } = makeTempDb(); + try { + // Insert a skill to satisfy FK + raw.exec(` + INSERT INTO skills (id, api_key_id, name, version, schema, handler) + VALUES ('s2','k1','test2','1.0.0','{}','h.js') + `); + raw.exec(` + INSERT INTO skill_executions (id, skill_id, api_key_id, input, status) + VALUES ('new_exec','s2','k1','{"q":"test2"}','success') + `); + const allRows = raw.prepare("SELECT * FROM skill_executions").all(); + assert.ok(allRows.length >= 1, "should read skill_executions after migration 174"); + } finally { + raw.close(); + fs.rmSync(dir, { recursive: true, force: true }); + done(); + } +}); + +test("migration 174: idempotent — running twice does not error or duplicate", (_, done) => { + const { raw, dir } = makeTempDb(); + try { + // makeTempDb already runs the SQL twice; verify no error and table exists + const tables = raw.pragma("table_list") as Array<{ name: string }>; + const names = tables.map((t) => t.name); + assert.ok(names.includes("server_tool_executions"), "table must exist after double-apply"); + // Verify no duplicate columns or indexes from double-apply + const indexes = getIndexInfo(raw, "server_tool_executions"); + const statusIdx = indexes.filter((i) => i.name.includes("status_expiry")); + assert.equal(statusIdx.length, 1, "must have exactly one status_expiry index"); + } finally { + raw.close(); + fs.rmSync(dir, { recursive: true, force: true }); + done(); + } +}); diff --git a/tests/unit/feature-flags-settings.test.ts b/tests/unit/feature-flags-settings.test.ts index 962f7ff60b..9f0184f41c 100644 --- a/tests/unit/feature-flags-settings.test.ts +++ b/tests/unit/feature-flags-settings.test.ts @@ -39,7 +39,7 @@ const { // OMNIROUTE_DISABLE_THINKING_LEVEL_VARIANTS bumped it from 53 to 54; // the dead ONEPROXY_ENABLED (readerless since the 1proxy purge, #12091) // brought it back to 53. UNIVERSAL_CONTEXT_HANDOFF_ENABLED bumped it to 54. -const EXPECTED_FEATURE_FLAG_COUNT = 54; +const EXPECTED_FEATURE_FLAG_COUNT = 55; // ────────────────────────────────────────────────────── // Test group 1 — Flag definitions registry diff --git a/tests/unit/follow-up-transcript.test.ts b/tests/unit/follow-up-transcript.test.ts new file mode 100644 index 0000000000..d886d73e25 --- /dev/null +++ b/tests/unit/follow-up-transcript.test.ts @@ -0,0 +1,705 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { Buffer } from "node:buffer"; + +import { + MAX_RESULT_BYTES_PER_TOOL, + MAX_RESULT_BYTES_TOTAL, + serializeBoundedToolResult, + buildFollowUpSourceBody, +} from "../../src/lib/skills/followUpTranscript.ts"; +import type { + ToolCall, + ExecutedToolResult, + BuildFollowUpTranscriptInput, + BoundedToolResult, +} from "../../src/lib/skills/toolLoopTypes.ts"; + +type UnknownRecord = Record; + +const MESSAGES_REQUIRED = "buildFollowUpSourceBody requires sourceBody.messages array"; +const NON_SERIALIZABLE_JSON = '{"error":"Tool result is not JSON-serializable"}'; + +// ─── OpenAI Chat Completions ───────────────────────────────────────────────── + +function openaiInput(overrides: Partial = {}) { + const sourceBody = { + model: "gpt-4o", + messages: [ + { role: "system", content: "You are a helpful assistant." }, + { role: "user", content: "remember foo" }, + ], + tools: [ + { + type: "function", + function: { + name: "memory_search", + description: "search memory", + parameters: { type: "object", properties: { query: { type: "string" } } }, + }, + }, + ], + tool_choice: "auto", + stream: false, + } as UnknownRecord; + const previousResponse = { + id: "chatcmpl-abc", + object: "chat.completion", + choices: [ + { + index: 0, + message: { + role: "assistant", + content: "I will look that up.", + tool_calls: [ + { + id: "call_1", + type: "function", + function: { name: "memory_search", arguments: '{"query":"foo"}' }, + }, + { + id: "call_2", + type: "function", + function: { name: "memory_save", arguments: '{"key":"k","value":"v"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + } as UnknownRecord; + const toolCalls: ToolCall[] = [ + { id: "call_1", name: "memory_search", arguments: { query: "foo" } }, + { id: "call_2", name: "memory_save", arguments: { key: "k", value: "v" } }, + ]; + const results: ExecutedToolResult[] = [ + { id: "call_1", name: "memory_search", result: { hits: ["a", "b"] }, replayed: false }, + { id: "call_2", name: "memory_save", result: { ok: true }, replayed: true }, + ]; + return { + sourceBody, + previousResponse, + toolCalls, + results, + sourceFormat: "openai", + maxResultBytes: MAX_RESULT_BYTES_PER_TOOL, + ...overrides, + } as BuildFollowUpTranscriptInput & { + sourceBody: UnknownRecord; + previousResponse: UnknownRecord; + toolCalls: ToolCall[]; + results: ExecutedToolResult[]; + }; +} + +test("OpenAI: messages array required — Responses-only body throws", () => { + const responsesBody = { + model: "gpt-4o", + input: [{ role: "user", content: "hi" }], + } as UnknownRecord; + assert.throws( + () => buildFollowUpSourceBody(openaiInput({ sourceBody: responsesBody })), + (err: unknown) => err instanceof Error && err.message === MESSAGES_REQUIRED + ); + assert.throws( + () => + buildFollowUpSourceBody( + openaiInput({ sourceBody: { model: "gpt-4o", messages: "not-an-array" } as UnknownRecord }) + ), + (err: unknown) => err instanceof Error && err.message === MESSAGES_REQUIRED + ); +}); + +test("OpenAI: calls/results length, duplicates and ID set mismatch all fail closed", () => { + const missingResult = openaiInput({ + toolCalls: [ + { id: "call_1", name: "memory_search", arguments: { query: "foo" } }, + { id: "call_9", name: "client_tool", arguments: {} }, + ], + }); + assert.throws(() => buildFollowUpSourceBody(missingResult), /matching result/); + + const extraResult = openaiInput({ + results: [ + { id: "call_1", name: "memory_search", result: { hits: ["a"] }, replayed: false }, + { id: "call_2", name: "memory_save", result: { ok: true }, replayed: false }, + { id: "call_9", name: "ghost", result: "no call", replayed: false }, + ], + }); + assert.throws(() => buildFollowUpSourceBody(extraResult), /same length/); + + const duplicateCallId = openaiInput({ + toolCalls: [ + { id: "call_1", name: "memory_search", arguments: {} }, + { id: "call_1", name: "memory_save", arguments: {} }, + ], + }); + assert.throws(() => buildFollowUpSourceBody(duplicateCallId), /unique tool call ids/); + + const duplicateResultId = openaiInput({ + results: [ + { id: "call_1", name: "memory_search", result: "a", replayed: false }, + { id: "call_1", name: "memory_search", result: "b", replayed: false }, + ], + }); + assert.throws(() => buildFollowUpSourceBody(duplicateResultId), /unique tool result ids/); +}); + +test("OpenAI: appends assistant turn and tool messages, preserves tools/tool_choice/model", () => { + const input = openaiInput(); + const sourceBodySnapshot = JSON.parse(JSON.stringify(input.sourceBody)); + const responseSnapshot = JSON.parse(JSON.stringify(input.previousResponse)); + + const out = buildFollowUpSourceBody(input); + + assert.notStrictEqual(out, input.sourceBody); + assert.deepEqual(input.sourceBody, sourceBodySnapshot); + assert.deepEqual(input.previousResponse, responseSnapshot); + + assert.strictEqual(out.model, "gpt-4o"); + assert.deepEqual(out.tools, input.sourceBody.tools); + assert.deepEqual(out.tool_choice, "auto"); + assert.strictEqual(out.stream, false); + + const messages = out.messages as UnknownRecord[]; + assert.strictEqual(messages.length, 5); + assert.deepEqual(messages.slice(0, 2), input.sourceBody.messages); + + const assistant = messages[2]; + assert.strictEqual(assistant.role, "assistant"); + assert.strictEqual(assistant.content, "I will look that up."); + assert.deepEqual(assistant.tool_calls, input.previousResponse.choices[0].message.tool_calls); + + for (let i = 0; i < input.results.length; i++) { + const toolMessage = messages[3 + i]; + assert.strictEqual(toolMessage.role, "tool"); + assert.strictEqual(toolMessage.tool_call_id, input.results[i].id); + assert.strictEqual(toolMessage.content, JSON.stringify(input.results[i].result)); + } +}); + +test("OpenAI: reconstructs assistant tool_calls from parsed calls when response carries none", () => { + const input = openaiInput(); + (input.previousResponse.choices[0].message as UnknownRecord).tool_calls = undefined; + + const out = buildFollowUpSourceBody(input); + + const assistant = (out.messages as UnknownRecord[])[2]; + assert.deepEqual(assistant.tool_calls, [ + { + id: "call_1", + type: "function", + function: { name: "memory_search", arguments: '{"query":"foo"}' }, + }, + { + id: "call_2", + type: "function", + function: { name: "memory_save", arguments: '{"key":"k","value":"v"}' }, + }, + ]); +}); + +test("OpenAI: bounds each tool result to 32768 UTF-8 bytes and total to 65536", () => { + const big = "x".repeat(40_000); + const input = openaiInput({ + results: [ + { id: "call_1", name: "memory_search", result: { payload: big }, replayed: false }, + { id: "call_2", name: "memory_save", result: { payload: big }, replayed: false }, + ], + }); + + const out = buildFollowUpSourceBody(input); + const messages = out.messages as UnknownRecord[]; + const firstContent = String(messages[3].content); + const secondContent = String(messages[4].content); + const firstBytes = Buffer.byteLength(firstContent, "utf8"); + const secondBytes = Buffer.byteLength(secondContent, "utf8"); + + assert.strictEqual(MAX_RESULT_BYTES_PER_TOOL, 32_768); + assert.strictEqual(MAX_RESULT_BYTES_TOTAL, 65_536); + assert.ok(firstBytes <= MAX_RESULT_BYTES_PER_TOOL, `first ${firstBytes}`); + assert.ok(secondBytes <= MAX_RESULT_BYTES_PER_TOOL, `second ${secondBytes}`); + assert.ok(firstBytes + secondBytes <= MAX_RESULT_BYTES_TOTAL, "total byte bound"); + assert.ok(firstContent.includes("[TRUNCATED")); + assert.ok(secondContent.includes("[TRUNCATED")); +}); + +test("OpenAI: maxTotalResultBytes=20 exhausts budget, second tool message is empty and total stays <=20", () => { + const input = openaiInput({ + results: [ + { id: "call_1", name: "memory_search", result: "a".repeat(240), replayed: false }, + { id: "call_2", name: "memory_save", result: "z".repeat(240), replayed: false }, + ], + maxTotalResultBytes: 20, + }); + + const out = buildFollowUpSourceBody(input); + const messages = out.messages as UnknownRecord[]; + const firstBytes = Buffer.byteLength(String(messages[3].content), "utf8"); + const secondContent = messages[4].content; + + assert.ok(firstBytes <= 20, `first ${firstBytes}`); + assert.strictEqual(secondContent, ""); + assert.ok(firstBytes + Buffer.byteLength(String(secondContent), "utf8") <= 20); + assert.ok(String(messages[3].content).includes("[TRUNCATED")); +}); + +// ─── serializeBoundedToolResult ────────────────────────────────────────────── + +function assertValidUtf8(text: string): void { + assert.strictEqual(Buffer.from(text, "utf8").toString("utf8"), text); +} + +function hasLoneSurrogate(text: string): boolean { + for (let i = 0; i < text.length; i++) { + const code = text.charCodeAt(i); + if (code >= 0xd800 && code <= 0xdbff) { + const next = text.charCodeAt(i + 1); + if (next < 0xdc00 || next > 0xdfff) return true; + } else if (code >= 0xdc00 && code <= 0xdfff) { + const prev = text.charCodeAt(i - 1); + if (prev < 0xd800 || prev > 0xdbff) return true; + } + } + return false; +} + +test("serializeBoundedToolResult: CJK + astral cut lands on code-point boundaries", () => { + const raw = "你好🌍".repeat(9); + const serialized = JSON.stringify(raw); + const originalBytes = Buffer.byteLength(serialized, "utf8"); + const maxBytes = 50; + + const bounded = serializeBoundedToolResult(raw, maxBytes); + + assert.strictEqual(bounded.originalBytes, originalBytes); + assert.ok(bounded.truncated); + assert.ok( + Buffer.byteLength(bounded.text, "utf8") <= maxBytes, + `byte bound: ${Buffer.byteLength(bounded.text, "utf8")} <= ${maxBytes}` + ); + assert.ok(bounded.text.includes("[TRUNCATED"), "marker present"); + const match = bounded.text.match(/\[TRUNCATED (\d+) BYTES BY OMNIROUTE\]/); + assert.ok(match, "marker format valid"); + const droppedBytes = Number(match![1]); + assert.ok(droppedBytes > 0, "dropped bytes > 0"); + assert.ok(droppedBytes <= originalBytes, "dropped <= original"); + assertValidUtf8(bounded.text); + assert.ok(!hasLoneSurrogate(bounded.text), "no orphan surrogate"); +}); + +test("serializeBoundedToolResult: marker-only UTF-8-safe prefix when maxBytes <= markerBytes, empty at 0", () => { + const text = "你好世界🌍🌎"; + const markerOnly = serializeBoundedToolResult(text, 20); + assert.ok(markerOnly.truncated); + assert.strictEqual(markerOnly.text, "[TRUNCATED 22 BYTES "); + assert.strictEqual(Buffer.byteLength(markerOnly.text, "utf8"), 20); + assertValidUtf8(markerOnly.text); + + const empty = serializeBoundedToolResult(text, 0); + assert.strictEqual(empty.text, ""); + assertValidUtf8(empty.text); +}); + +test("serializeBoundedToolResult: unfits pass through untruncated with byte counts", () => { + const bounded = serializeBoundedToolResult({ hits: ["a"] }, 1024); + assert.strictEqual(bounded.truncated, false); + assert.strictEqual(bounded.text, '{"hits":["a"]}'); + assert.strictEqual(bounded.originalBytes, Buffer.byteLength('{"hits":["a"]}', "utf8")); +}); + +test("serializeBoundedToolResult: projects undefined/Error/bigint and rejects non-JSON values", () => { + assert.strictEqual(serializeBoundedToolResult(undefined, 1024).text, "null"); + assert.strictEqual(serializeBoundedToolResult(new Error("x"), 1024).text, '{"error":"x"}'); + assert.strictEqual(serializeBoundedToolResult(10n, 1024).text, '"10"'); + + assert.strictEqual(serializeBoundedToolResult(() => undefined, 1024).text, NON_SERIALIZABLE_JSON); + assert.strictEqual(serializeBoundedToolResult(Symbol("s"), 1024).text, NON_SERIALIZABLE_JSON); + + const cyclic: UnknownRecord = {}; + cyclic.self = cyclic; + assert.strictEqual(serializeBoundedToolResult(cyclic, 1024).text, NON_SERIALIZABLE_JSON); + assert.strictEqual(serializeBoundedToolResult({ a: 10n }, 1024).text, NON_SERIALIZABLE_JSON); + + const boundedResult: BoundedToolResult = serializeBoundedToolResult(10n, 1024); + assert.deepEqual(Object.keys(boundedResult).sort(), ["originalBytes", "text", "truncated"]); +}); + +// ─── Task 4: Budget validation (Req 1) ───────────────────────────────────── + +test("serializeBoundedToolResult: NaN budget throws RangeError", () => { + assert.throws(() => serializeBoundedToolResult("hello", NaN), { + name: "RangeError", + message: /maxBytes must be a non-negative finite integer/, + }); +}); + +test("serializeBoundedToolResult: Infinity budget throws RangeError", () => { + assert.throws(() => serializeBoundedToolResult("hello", Infinity), { + name: "RangeError", + message: /maxBytes must be a non-negative finite integer/, + }); +}); + +test("serializeBoundedToolResult: -Infinity budget throws RangeError", () => { + assert.throws(() => serializeBoundedToolResult("hello", -Infinity), { + name: "RangeError", + message: /maxBytes must be a non-negative finite integer/, + }); +}); + +test("serializeBoundedToolResult: negative budget throws RangeError", () => { + assert.throws(() => serializeBoundedToolResult("hello", -1), { + name: "RangeError", + message: /maxBytes must be a non-negative finite integer/, + }); +}); + +test("serializeBoundedToolResult: non-integer budget throws RangeError", () => { + assert.throws(() => serializeBoundedToolResult("hello", 1.5), { + name: "RangeError", + message: /maxBytes must be a non-negative finite integer/, + }); +}); + +test("buildFollowUpSourceBody: NaN maxResultBytes throws RangeError", () => { + const input = openaiInput({ maxResultBytes: NaN }); + assert.throws(() => buildFollowUpSourceBody(input), { + name: "RangeError", + message: /must be a non-negative finite integer/, + }); +}); + +test("buildFollowUpSourceBody: negative maxTotalResultBytes throws RangeError", () => { + const input = openaiInput({ maxTotalResultBytes: -1 }); + assert.throws(() => buildFollowUpSourceBody(input), { + name: "RangeError", + message: /must be a non-negative finite integer/, + }); +}); + +// ─── Task 4: sourceFormat validation (Req 2) ─────────────────────────────── + +test("buildFollowUpSourceBody: invalid sourceFormat 'anthropic' fails closed", () => { + const input = openaiInput({ sourceFormat: "anthropic" } as unknown as { + sourceFormat: "openai" | "claude"; + }); + assert.throws(() => buildFollowUpSourceBody(input), /sourceFormat must be "openai" or "claude"/); +}); + +test("buildFollowUpSourceBody: empty sourceFormat fails closed", () => { + const input = openaiInput({ sourceFormat: "" } as unknown as { + sourceFormat: "openai" | "claude"; + }); + assert.throws(() => buildFollowUpSourceBody(input), /sourceFormat must be "openai" or "claude"/); +}); + +// ─── Task 4: calls/results name match (Req 3) ────────────────────────────── + +test("calls/results name mismatch on same ID fails closed", () => { + const input = openaiInput({ + toolCalls: [ + { id: "call_1", name: "memory_search", arguments: { query: "foo" } }, + { id: "call_2", name: "memory_save", arguments: { key: "k", value: "v" } }, + ], + results: [ + { id: "call_1", name: "WRONG_NAME", result: { hits: ["a"] }, replayed: false }, + { id: "call_2", name: "memory_save", result: { ok: true }, replayed: false }, + ], + }); + assert.throws( + () => buildFollowUpSourceBody(input), + /result name .* does not match tool call name/ + ); +}); + +// ─── Task 4: OpenAI previous response tool_call match (Req 4) ─────────────── + +test("OpenAI: partial tool_calls in previous response fails closed (missing call)", () => { + const input = openaiInput(); + // Only include call_1 in the previous response, omit call_2 + (input.previousResponse.choices[0].message as UnknownRecord).tool_calls = [ + { + id: "call_1", + type: "function", + function: { name: "memory_search", arguments: '{"query":"foo"}' }, + }, + ]; + assert.throws( + () => buildFollowUpSourceBody(input), + /previous response tool_calls must contain exactly one match per call ID/ + ); +}); + +test("OpenAI: duplicate tool_calls for same ID in previous response fails closed", () => { + const input = openaiInput(); + // Duplicate call_1 in the previous response + (input.previousResponse.choices[0].message as UnknownRecord).tool_calls = [ + { + id: "call_1", + type: "function", + function: { name: "memory_search", arguments: '{"query":"foo"}' }, + }, + { + id: "call_1", + type: "function", + function: { name: "memory_search", arguments: '{"query":"foo"}' }, + }, + { + id: "call_2", + type: "function", + function: { name: "memory_save", arguments: '{"key":"k","value":"v"}' }, + }, + ]; + assert.throws( + () => buildFollowUpSourceBody(input), + /previous response tool_calls must contain exactly one match per call ID/ + ); +}); + +test("OpenAI: mismatched name in previous response tool_call fails closed", () => { + const input = openaiInput(); + // call_1 has a different name in the previous response + (input.previousResponse.choices[0].message as UnknownRecord).tool_calls = [ + { + id: "call_1", + type: "function", + function: { name: "WRONG_NAME", arguments: '{"query":"foo"}' }, + }, + { + id: "call_2", + type: "function", + function: { name: "memory_save", arguments: '{"key":"k","value":"v"}' }, + }, + ]; + assert.throws( + () => buildFollowUpSourceBody(input), + /previous response tool_call .* name .* does not match/ + ); +}); + +// ─── Task 4: Claude previousResponse tool_use match (Req 5) ───────────────── + +test("Claude: partial tool_use blocks in previous response fails closed", () => { + const input = claudeInput(); + // Only include toolu_1, omit toolu_2 + (input.previousResponse as UnknownRecord).content = [ + { type: "text", text: "Let me look." }, + { type: "tool_use", id: "toolu_1", name: "memory_search", input: { query: "foo" } }, + ]; + assert.throws( + () => buildFollowUpSourceBody(input), + /previous response content must contain exactly one tool_use match per call ID/ + ); +}); + +test("Claude: duplicate tool_use blocks for same ID in previous response fails closed", () => { + const input = claudeInput(); + // Duplicate toolu_1 + (input.previousResponse as UnknownRecord).content = [ + { type: "text", text: "Let me look." }, + { type: "tool_use", id: "toolu_1", name: "memory_search", input: { query: "foo" } }, + { type: "tool_use", id: "toolu_1", name: "memory_search", input: { query: "foo" } }, + { type: "tool_use", id: "toolu_2", name: "memory_save", input: { key: "k" } }, + ]; + assert.throws( + () => buildFollowUpSourceBody(input), + /previous response content must contain exactly one tool_use match per call ID/ + ); +}); + +test("Claude: mismatched name in previous response tool_use fails closed", () => { + const input = claudeInput(); + // toolu_1 has a different name + (input.previousResponse as UnknownRecord).content = [ + { type: "text", text: "Let me look." }, + { type: "tool_use", id: "toolu_1", name: "WRONG_NAME", input: { query: "foo" } }, + { type: "tool_use", id: "toolu_2", name: "memory_save", input: { key: "k" } }, + ]; + assert.throws( + () => buildFollowUpSourceBody(input), + /previous response tool_use .* name .* does not match/ + ); +}); + +// ─── Claude Messages ───────────────────────────────────────────────────────── + +function claudeInput(overrides: Partial = {}) { + const sourceBody = { + model: "claude-sonnet-4-5", + max_tokens: 1024, + stream: true, + messages: [{ role: "user", content: "remember foo" }], + tools: [ + { + name: "memory_search", + description: "search memory", + input_schema: { type: "object", properties: { query: { type: "string" } } }, + }, + ], + } as UnknownRecord; + const previousResponse = { + id: "msg_1", + type: "message", + role: "assistant", + content: [ + { type: "text", text: "Let me look." }, + { type: "tool_use", id: "toolu_1", name: "memory_search", input: { query: "foo" } }, + { type: "tool_use", id: "toolu_2", name: "memory_save", input: { key: "k" } }, + ], + stop_reason: "tool_use", + } as UnknownRecord; + const toolCalls: ToolCall[] = [ + { id: "toolu_1", name: "memory_search", arguments: { query: "foo" } }, + { id: "toolu_2", name: "memory_save", arguments: { key: "k" } }, + ]; + const results: ExecutedToolResult[] = [ + { id: "toolu_1", name: "memory_search", result: { hits: ["a"] }, replayed: false }, + { id: "toolu_2", name: "memory_save", result: { stored: true }, replayed: true }, + ]; + return { + sourceBody, + previousResponse, + toolCalls, + results, + sourceFormat: "claude", + maxResultBytes: MAX_RESULT_BYTES_PER_TOOL, + ...overrides, + } as BuildFollowUpTranscriptInput & { + sourceBody: UnknownRecord; + previousResponse: UnknownRecord; + toolCalls: ToolCall[]; + results: ExecutedToolResult[]; + }; +} + +test("Claude: appends assistant tool_use turn, then a separate user tool_result message", () => { + const input = claudeInput(); + const sourceBodySnapshot = JSON.parse(JSON.stringify(input.sourceBody)); + + const out = buildFollowUpSourceBody(input); + + assert.notStrictEqual(out, input.sourceBody); + assert.deepEqual(input.sourceBody, sourceBodySnapshot); + assert.strictEqual(out.stream, true, "stream is not forced for Claude bodies"); + + const messages = out.messages as UnknownRecord[]; + assert.strictEqual(messages.length, 3); + assert.deepEqual(messages[0], input.sourceBody.messages[0]); + + const assistant = messages[1]; + assert.strictEqual(assistant.role, "assistant"); + const assistantBlocks = assistant.content as UnknownRecord[]; + assert.strictEqual(assistantBlocks.length, 3, "text + 2 tool_use blocks preserved"); + assert.strictEqual(assistantBlocks[0].type, "text"); + assert.strictEqual(assistantBlocks[0].text, "Let me look."); + assert.ok(assistantBlocks.slice(1).every((block) => block.type === "tool_use")); + assert.deepEqual( + assistantBlocks.slice(1), + input.previousResponse.content.filter((block: UnknownRecord) => block.type === "tool_use") + ); + + const user = messages[2]; + assert.strictEqual(user.role, "user"); + const resultBlocks = user.content as UnknownRecord[]; + assert.strictEqual(resultBlocks.length, 2); + assert.ok(resultBlocks.every((block) => block.type === "tool_result")); + assert.deepEqual( + resultBlocks.map((block) => block.tool_use_id), + ["toolu_1", "toolu_2"] + ); + assert.strictEqual(resultBlocks[0].content, JSON.stringify({ hits: ["a"] })); + assert.strictEqual(resultBlocks[1].content, JSON.stringify({ stored: true })); + + assert.deepEqual(out.tools, input.sourceBody.tools); +}); + +// ─── Task 4 Fix R1: Claude text/thinking block preservation ───────────────── + +test("Claude: assistant content preserves original text blocks in original order", () => { + const input = claudeInput(); + const out = buildFollowUpSourceBody(input); + const messages = out.messages as UnknownRecord[]; + const assistant = messages[1]; + const assistantBlocks = assistant.content as UnknownRecord[]; + // Must include the text block (index 0) AND the two tool_use blocks + assert.strictEqual(assistantBlocks.length, 3, "should have text + 2 tool_use blocks"); + assert.strictEqual(assistantBlocks[0].type, "text"); + assert.strictEqual(assistantBlocks[0].text, "Let me look."); + assert.strictEqual(assistantBlocks[1].type, "tool_use"); + assert.strictEqual(assistantBlocks[1].id, "toolu_1"); + assert.strictEqual(assistantBlocks[2].type, "tool_use"); + assert.strictEqual(assistantBlocks[2].id, "toolu_2"); +}); + +test("Claude: assistant content preserves thinking blocks alongside tool_use", () => { + const input = claudeInput(); + (input.previousResponse as UnknownRecord).content = [ + { type: "thinking", thinking: "Let me reason about this.", signature: "sig_1" }, + { type: "text", text: "I'll search now." }, + { type: "tool_use", id: "toolu_1", name: "memory_search", input: { query: "foo" } }, + { type: "tool_use", id: "toolu_2", name: "memory_save", input: { key: "k" } }, + ]; + const out = buildFollowUpSourceBody(input); + const messages = out.messages as UnknownRecord[]; + const assistantBlocks = messages[1].content as UnknownRecord[]; + assert.strictEqual(assistantBlocks.length, 4, "should have thinking + text + 2 tool_use"); + assert.strictEqual(assistantBlocks[0].type, "thinking"); + assert.strictEqual(assistantBlocks[0].thinking, "Let me reason about this."); + assert.strictEqual(assistantBlocks[1].type, "text"); + assert.strictEqual(assistantBlocks[1].text, "I'll search now."); + assert.strictEqual(assistantBlocks[2].type, "tool_use"); + assert.strictEqual(assistantBlocks[3].type, "tool_use"); +}); + +test("Claude: assistant content filters unmatched tool_use but keeps text", () => { + const input = claudeInput(); + (input.previousResponse as UnknownRecord).content = [ + { type: "text", text: "Looking..." }, + { type: "tool_use", id: "toolu_1", name: "memory_search", input: { query: "foo" } }, + { type: "tool_use", id: "toolu_UNMATCHED", name: "other_tool", input: {} }, + { type: "tool_use", id: "toolu_2", name: "memory_save", input: { key: "k" } }, + ]; + const out = buildFollowUpSourceBody(input); + const messages = out.messages as UnknownRecord[]; + const assistantBlocks = messages[1].content as UnknownRecord[]; + // text preserved, unmatched tool_use filtered out, matched tool_use kept + assert.strictEqual(assistantBlocks.length, 3); + assert.strictEqual(assistantBlocks[0].type, "text"); + assert.strictEqual(assistantBlocks[1].id, "toolu_1"); + assert.strictEqual(assistantBlocks[2].id, "toolu_2"); +}); + +test("Claude: strips text after first tool_use but preserves later thinking blocks", () => { + const input = claudeInput(); + (input.previousResponse as UnknownRecord).content = [ + { type: "text", text: "before" }, + { type: "tool_use", id: "toolu_1", name: "memory_search", input: { query: "foo" } }, + { type: "text", text: "after — must be stripped" }, + { type: "thinking", thinking: "signed thought", signature: "sig_after" }, + { type: "tool_use", id: "toolu_2", name: "memory_save", input: { key: "k" } }, + ]; + + const out = buildFollowUpSourceBody(input); + const assistantBlocks = (out.messages as UnknownRecord[])[1].content as UnknownRecord[]; + assert.deepEqual( + assistantBlocks.map((block) => block.type), + ["text", "tool_use", "thinking", "tool_use"] + ); + assert.equal( + assistantBlocks.some( + (block) => block.type === "text" && block.text === "after — must be stripped" + ), + false + ); +}); + +test("Claude: mismatched IDs fail closed even when response content would filter them", () => { + const input = claudeInput({ + toolCalls: [ + { id: "toolu_1", name: "memory_search", arguments: { query: "foo" } }, + { id: "toolu_9", name: "client_tool", arguments: {} }, + ], + }); + assert.throws(() => buildFollowUpSourceBody(input), /matching result/); +}); diff --git a/tests/unit/non-streaming-client-translate.test.ts b/tests/unit/non-streaming-client-translate.test.ts new file mode 100644 index 0000000000..fb382929fa --- /dev/null +++ b/tests/unit/non-streaming-client-translate.test.ts @@ -0,0 +1,322 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; + +import { + translateNonStreamingClientResponse, + type NonStreamingClientTranslateInput, +} from "../../open-sse/handlers/chatCore/nonStreamingClientTranslate.ts"; +import { FORMATS } from "../../open-sse/translator/formats.ts"; +import { + buildAssistantMessageCacheKey, + clearReasoningCacheAll, + lookupReasoning, +} from "../../open-sse/services/reasoningCache.ts"; +import { invalidateBufferTokensCache } from "../../open-sse/utils/usageTracking.ts"; + +/* ── helpers ─────────────────────────────────────────────────────────────── */ + +function baseInput( + overrides: Partial = {} +): NonStreamingClientTranslateInput { + return { + responseBody: { + id: "chatcmpl-test", + object: "chat.completion", + choices: [ + { + index: 0, + message: { role: "assistant", content: "Hello!" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }, + responsePayloadFormat: "openai", + clientResponseFormat: "openai", + sourceFormat: "openai", + provider: "openai", + model: "gpt-4o", + requestBody: { messages: [{ role: "user", content: "hi" }] }, + responseToolNameMap: null, + requestToolIdentityMap: null, + reasoningCacheScope: null, + clientHeaders: null, + isClaudeCodeCompatible: false, + phase: "final", + ...overrides, + }; +} + +/* ── characterization tests ──────────────────────────────────────────────── */ + +test("basic translate: same-format passthrough returns responseBody", () => { + const input = baseInput(); + const result = translateNonStreamingClientResponse(input); + assert.equal(result.response.choices[0].message.content, "Hello!"); + assert.ok(result.responseForMemoryExtraction); +}); + +test("translate from claude to openai format", () => { + const input = baseInput({ + responseBody: { + id: "msg-123", + content: [{ type: "text", text: "Hi there" }], + stop_reason: "end_turn", + usage: { input_tokens: 10, output_tokens: 5 }, + }, + responsePayloadFormat: "claude", + clientResponseFormat: "openai", + sourceFormat: "claude", + provider: "anthropic", + model: "claude-sonnet-4-20250514", + }); + const result = translateNonStreamingClientResponse(input); + const msg = result.response.choices?.[0]?.message; + assert.ok(msg, "should have choices[0].message"); + assert.equal((msg as { content: string }).content, "Hi there"); +}); + +test("claude source strips markdown code fence", () => { + const input = baseInput({ + responseBody: { + id: "msg-123", + content: [ + { + type: "text", + text: '```json\n{"key": "value"}\n```', + }, + ], + stop_reason: "end_turn", + usage: { input_tokens: 10, output_tokens: 5 }, + }, + responsePayloadFormat: "claude", + clientResponseFormat: "openai", + sourceFormat: "claude", + provider: "anthropic", + model: "claude-sonnet-4-20250514", + }); + const result = translateNonStreamingClientResponse(input); + const content = result.response.choices?.[0]?.message?.content; + assert.ok(typeof content === "string"); + // After stripping, the content should not have the markdown fence wrapper + assert.ok(!content.startsWith("```json"), "markdown fence should be stripped"); +}); + +test("normalizeOpenAIToolFinishReasons: tool_calls present → finish_reason tool_calls", () => { + const input = baseInput({ + responseBody: { + id: "chatcmpl-test", + choices: [ + { + index: 0, + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_1", + type: "function", + function: { name: "get_weather", arguments: "{}" }, + }, + ], + }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }, + }); + const result = translateNonStreamingClientResponse(input); + assert.equal(result.response.choices[0].finish_reason, "tool_calls"); +}); + +test("reasoning replay: no-tool history comes from historyMessages, not requestBody.input", () => { + clearReasoningCacheAll(); + const scope = "api-key:test:s...6a"; + const historyMessages = [{ role: "user", content: "hi from translatedBody" }]; + const assistantMessage = { + role: "assistant", + content: "thinking result", + reasoning_content: "let me think...", + }; + const input = baseInput({ + responseBody: { + id: "chatcmpl-test", + choices: [ + { + index: 0, + message: assistantMessage, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }, + // Responses-shaped finalBody: input, no messages. Parent used translatedBody.messages. + requestBody: { input: [{ role: "user", content: "wrong body" }] }, + historyMessages, + provider: "deepseek", + model: "deepseek-v4-pro", + reasoningCacheScope: scope, + phase: "intermediate", + }); + const result = translateNonStreamingClientResponse(input); + assert.ok(result.response); + const cacheKey = buildAssistantMessageCacheKey( + scope, + [...historyMessages, assistantMessage], + historyMessages.length + ); + assert.equal( + lookupReasoning(cacheKey), + "let me think...", + "must cache against translatedBody.messages, not finalBody.input" + ); +}); + +test("phase=final applies client usage buffer", () => { + // Gemini format skips OpenAI/Responses sanitize, so extra usage fields + // only disappear if applyClientUsageBuffer → filterUsageForFormat runs. + const input = baseInput({ + phase: "final", + clientResponseFormat: FORMATS.GEMINI, + responsePayloadFormat: FORMATS.GEMINI, + sourceFormat: FORMATS.GEMINI, + responseBody: { + id: "chatcmpl-test", + choices: [ + { + index: 0, + message: { role: "assistant", content: "Hello!" }, + finish_reason: "stop", + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + x_provider_extra: 99, + }, + }, + }); + const result = translateNonStreamingClientResponse(input); + const usage = (result.response as { usage: Record }).usage; + assert.equal(usage.x_provider_extra, undefined, "final phase must filter extra usage fields"); + assert.equal( + usage.prompt_tokens, + undefined, + "final Gemini filter must drop OpenAI-shaped prompt_tokens" + ); +}); + +test("phase=intermediate skips applyClientUsageBuffer", () => { + const input = baseInput({ + phase: "intermediate", + clientResponseFormat: FORMATS.GEMINI, + responsePayloadFormat: FORMATS.GEMINI, + sourceFormat: FORMATS.GEMINI, + responseBody: { + id: "chatcmpl-test", + choices: [ + { + index: 0, + message: { role: "assistant", content: "partial" }, + finish_reason: null, + }, + ], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + x_provider_extra: 99, + }, + }, + }); + const result = translateNonStreamingClientResponse(input); + const usage = (result.response as { usage: Record }).usage; + assert.equal( + usage.x_provider_extra, + 99, + "intermediate must keep raw extra usage fields (buffer not applied)" + ); + assert.equal(usage.prompt_tokens, 10); +}); + +test("Responses API format: sanitizeResponsesApiResponse is applied", () => { + const input = baseInput({ + clientResponseFormat: FORMATS.OPENAI_RESPONSES, + responseBody: { + id: "resp_123", + object: "response", + output: [ + { + type: "function_call", + name: "ns__get_weather", + arguments: "{}", + }, + ], + usage: { input_tokens: 10, output_tokens: 5 }, + illegal_top_level: "drop-me", + }, + requestToolIdentityMap: new Map([ + ["ns__get_weather", { namespace: "ns", name: "get_weather" }], + ]), + }); + const result = translateNonStreamingClientResponse(input); + assert.equal(result.response.object, "response"); + assert.equal(result.response.illegal_top_level, undefined, "sanitizer must drop illegal fields"); + const output = result.response.output as Array>; + assert.equal(output[0]?.type, "function_call"); + assert.equal(output[0]?.namespace, "ns", "#7936 restore namespace"); + assert.equal(output[0]?.name, "get_weather", "#7936 restore original name"); +}); + +test("empty content response: passthrough without crash", () => { + const input = baseInput({ + responseBody: {}, + }); + const result = translateNonStreamingClientResponse(input); + assert.ok(result.response); + assert.ok(result.responseForMemoryExtraction); +}); + +test("isClaudeCodeCompatible preserves context budget usage", () => { + const saved = process.env.USAGE_TOKEN_BUFFER; + process.env.USAGE_TOKEN_BUFFER = "2000"; + invalidateBufferTokensCache(); + try { + const input = baseInput({ + isClaudeCodeCompatible: true, + clientResponseFormat: FORMATS.OPENAI, + phase: "final", + responseBody: { + id: "chatcmpl-test", + object: "chat.completion", + choices: [ + { + index: 0, + message: { role: "assistant", content: "test" }, + finish_reason: "stop", + }, + ], + usage: { + prompt_tokens: 100, + completion_tokens: 50, + total_tokens: 150, + }, + }, + }); + const result = translateNonStreamingClientResponse(input); + const usage = (result.response as { usage: Record }).usage; + assert.equal( + usage.prompt_tokens, + 2100, + "Claude Code path must fold context_budget_prompt_tokens (100+2000) into visible prompt_tokens" + ); + assert.equal(usage.total_tokens, 2150); + assert.equal("context_budget_prompt_tokens" in usage, false); + } finally { + if (saved === undefined) delete process.env.USAGE_TOKEN_BUFFER; + else process.env.USAGE_TOKEN_BUFFER = saved; + invalidateBufferTokensCache(); + } +}); diff --git a/tests/unit/non-streaming-finalization.test.ts b/tests/unit/non-streaming-finalization.test.ts new file mode 100644 index 0000000000..1ab2dbfd1c --- /dev/null +++ b/tests/unit/non-streaming-finalization.test.ts @@ -0,0 +1,219 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; + +import type { + ChatCoreErrorResult, + ProviderLegReceipt, + ProviderLegUsage, + ServerOwnedToolLoopResult, +} from "../../src/lib/skills/toolLoopTypes.ts"; +import { + buildNonStreamingFinalizationPlan, + finalizeNonStreamingRequest, + finalizeToolLoopError, + type NonStreamingFinalizationDeps, +} from "../../open-sse/handlers/chatCore/nonStreamingFinalization.ts"; + +function receipt(index: number, overrides: Partial = {}): ProviderLegReceipt { + return { + index, + connectionId: "conn-1", + provider: "openai", + model: "gpt-4o", + startedAt: "2026-01-01T00:00:00.000Z", + endedAt: "2026-01-01T00:00:01.000Z", + latencyMs: 100, + httpStatus: 200, + errorType: null, + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + serviceTier: null, + computedCostUsd: 0.01, + toolCalls: [], + termination: "completed", + clientVisible: true, + ...overrides, + }; +} + +function usage(overrides: Partial = {}): ProviderLegUsage { + return { + prompt_tokens: 100, + completion_tokens: 20, + total_tokens: 120, + cache_read_input_tokens: 10, + reasoning_tokens: 5, + ...overrides, + }; +} + +function errorResult(): ChatCoreErrorResult { + return { + success: false, + status: 429, + response: new Response(JSON.stringify({ error: { message: "rate limited" } }), { + status: 429, + }), + error: "rate limited", + errorCode: "rate_limited", + errorType: "rate_limit_error", + retryAfterMs: 1000, + }; +} + +function spyDeps(): NonStreamingFinalizationDeps & { calls: Record } { + const calls = { + writeUsage: 0, + writeCost: 0, + scheduleQuota: 0, + writeAttempt: 0, + finalizePending: 0, + }; + return { + calls, + writeUsage: () => { + calls.writeUsage += 1; + }, + writeCost: () => { + calls.writeCost += 1; + }, + scheduleQuota: () => { + calls.scheduleQuota += 1; + }, + writeAttempt: () => { + calls.writeAttempt += 1; + }, + finalizePending: () => { + calls.finalizePending += 1; + }, + }; +} + +test("failure plan maps loop aggregate and 429 error", () => { + const loop: ServerOwnedToolLoopResult = { + kind: "error", + errorResult: errorResult(), + cumulativeUsage: usage(), + totalCostUsd: 0.03, + receipts: [ + receipt(0), + receipt(1, { + httpStatus: 429, + errorType: "rate_limit_error", + termination: "provider_error", + computedCostUsd: 0.02, + }), + ], + followUps: 1, + termination: "provider_error", + }; + const plan = buildNonStreamingFinalizationPlan(loop); + assert.equal(plan.kind, "failure"); + if (plan.kind !== "failure") return; + assert.equal(plan.error.status, 429); + assert.equal(plan.error.errorCode, "rate_limited"); + assert.deepEqual(plan.usage, usage()); + assert.equal(plan.totalCostUsd, 0.03); + assert.equal(plan.receiptCount, 2); +}); + +test("success plan maps loop usage and receipt count", () => { + const loop: ServerOwnedToolLoopResult = { + kind: "ok", + response: { choices: [{ message: { content: "done" }, finish_reason: "stop" }] }, + cumulativeUsage: usage({ prompt_tokens: 11, completion_tokens: 7, total_tokens: 18 }), + totalCostUsd: 0.02, + receipts: [receipt(0)], + followUps: 0, + termination: "completed", + }; + const plan = buildNonStreamingFinalizationPlan(loop); + assert.equal(plan.kind, "success"); + if (plan.kind !== "success") return; + assert.equal(plan.usage?.prompt_tokens, 11); + assert.equal(plan.totalCostUsd, 0.02); + assert.equal(plan.receiptCount, 1); +}); + +test("failure finalizer writes usage/cost/attempt/pending once and skips quota", async () => { + const loop: ServerOwnedToolLoopResult = { + kind: "error", + errorResult: errorResult(), + cumulativeUsage: usage(), + totalCostUsd: 0.03, + receipts: [receipt(0), receipt(1, { httpStatus: 429 })], + followUps: 1, + termination: "provider_error", + }; + const deps = spyDeps(); + await finalizeNonStreamingRequest(buildNonStreamingFinalizationPlan(loop), deps); + assert.equal(deps.calls.writeUsage, 1); + assert.equal(deps.calls.writeCost, 1); + assert.equal(deps.calls.scheduleQuota, 0); + assert.equal(deps.calls.writeAttempt, 1); + assert.equal(deps.calls.finalizePending, 1); +}); + +test("success finalizer writes usage/cost/quota/attempt/pending once", async () => { + const loop: ServerOwnedToolLoopResult = { + kind: "ok", + response: { choices: [{ message: { content: "done" }, finish_reason: "stop" }] }, + cumulativeUsage: usage(), + totalCostUsd: 0.02, + receipts: [receipt(0)], + followUps: 0, + termination: "completed", + }; + const deps = spyDeps(); + await finalizeNonStreamingRequest(buildNonStreamingFinalizationPlan(loop), deps); + assert.equal(deps.calls.writeUsage, 1); + assert.equal(deps.calls.writeCost, 1); + assert.equal(deps.calls.scheduleQuota, 1); + assert.equal(deps.calls.writeAttempt, 1); + assert.equal(deps.calls.finalizePending, 1); +}); + +test("finalizeToolLoopError delegates through finalization plan and deps", async () => { + const loop: ServerOwnedToolLoopResult = { + kind: "error", + errorResult: errorResult(), + cumulativeUsage: usage(), + totalCostUsd: 0.05, + receipts: [receipt(0), receipt(1, { httpStatus: 429 })], + followUps: 1, + termination: "provider_error", + }; + let usageSaved = false; + let attemptLogged = false; + let pendingTracked = false; + + const res = await finalizeToolLoopError({ + loop, + model: "gpt-4o", + provider: "openai", + connectionId: "conn-1", + providerRequest: { messages: [] }, + persistFailureUsage: (status, code, u) => { + assert.equal(status, 429); + assert.equal(code, "rate_limited"); + assert.equal(u?.prompt_tokens, 100); + assert.equal(u?.completion_tokens, 20); + assert.equal(u?.cache_read_input_tokens, 10); + assert.equal(u?.reasoning_tokens, 5); + usageSaved = true; + }, + persistAttemptLogs: (params) => { + assert.equal(params.status, 429); + attemptLogged = true; + }, + trackPendingRequest: (m, _p, _conn, pending) => { + assert.equal(m, "gpt-4o"); + assert.equal(pending, false); + pendingTracked = true; + }, + }); + + assert.equal(res.status, 429); + assert.equal(usageSaved, true); + assert.equal(attemptLogged, true); + assert.equal(pendingTracked, true); +}); diff --git a/tests/unit/non-streaming-provider-leg.test.ts b/tests/unit/non-streaming-provider-leg.test.ts new file mode 100644 index 0000000000..5001396f64 --- /dev/null +++ b/tests/unit/non-streaming-provider-leg.test.ts @@ -0,0 +1,1129 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; + +import { + runNonStreamingProviderLeg, + type ChatCoreExecutorResult, + type ProviderLegInput, +} from "../../open-sse/handlers/chatCore/nonStreamingProviderLeg.ts"; +import { + buildAssistantMessageCacheKey, + clearReasoningCacheAll, + lookupReasoning, +} from "../../open-sse/services/reasoningCache.ts"; + +/* -- helpers --------------------------------------------------------------- */ + +function makeResponse( + body: string | object, + status = 200, + headers: Record = {} +): Response { + const text = typeof body === "string" ? body : JSON.stringify(body); + return { + ok: status >= 200 && status < 300, + status, + statusText: status === 200 ? "OK" : "Error", + headers: new Headers({ "content-type": "application/json", ...headers }), + text: async () => text, + clone() { + return { ...this, text: async () => text } as unknown as Response; + }, + body: null, + } as unknown as Response; +} + +function makeExecutorResult( + responseBody: unknown, + status = 200, + headers: Record = {} +): ChatCoreExecutorResult { + return { + response: makeResponse(responseBody, status, headers), + url: "https://api.openai.com/v1/chat/completions", + headers: {}, + transformedBody: responseBody, + }; +} + +function baseInput(overrides: Partial = {}): ProviderLegInput { + return { + phase: "initial", + sourceBody: { + model: "gpt-4o", + messages: [{ role: "user", content: "hi" }], + }, + expectedConnectionId: undefined, + allowAccountRotation: true, + allowModelFallback: true, + executeProviderRequest: async () => + makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "Hello!" }, finish_reason: "stop" }], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }), + setRequestWireState: () => {}, + provider: "openai", + model: "gpt-4o", + connectionId: "conn-test", + ...overrides, + }; +} + +/* -- characterization tests ------------------------------------------------ */ + +test("200 JSON: returns ok with usage and receipt", async () => { + const result = await runNonStreamingProviderLeg(baseInput()); + assert.equal(result.kind, "ok"); + if (result.kind !== "ok") return; + assert.equal(result.response.choices[0].message.content, "Hello!"); + assert.ok(result.usage, "usage should be present"); + assert.equal(result.usage!.prompt_tokens, 10); + assert.equal(result.usage!.completion_tokens, 5); + assert.equal(result.receipt.httpStatus, 200); + assert.equal(result.receipt.termination, "completed"); +}); + +test("runProviderExecution is called once with policy; first send skips executeProviderRequest", async () => { + let pipelineCalls = 0; + let executorCalls = 0; + let seenPolicy: { allowAccountRotation: boolean; allowModelFallback: boolean } | undefined; + const okBody = { + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "from-pipeline" }, finish_reason: "stop" }], + usage: { prompt_tokens: 3, completion_tokens: 2, total_tokens: 5 }, + }; + const input = baseInput({ + executeProviderRequest: async () => { + executorCalls++; + throw new Error("first send must not use executeProviderRequest when pipeline is set"); + }, + runProviderExecution: async ({ policy, model, translatedBody }) => { + pipelineCalls++; + seenPolicy = { + allowAccountRotation: policy.allowAccountRotation, + allowModelFallback: policy.allowModelFallback, + }; + assert.equal(model, "gpt-4o"); + assert.equal((translatedBody as { model?: string }).model, "gpt-4o"); + return { + kind: "response", + response: { + ok: true, + status: 200, + statusText: "OK", + headers: new Headers({ "content-type": "application/json" }), + text: async () => JSON.stringify(okBody), + body: null, + } as unknown as Response, + url: "https://api.openai.com/v1/chat/completions", + headers: {}, + transformedBody: okBody, + model: "gpt-4o", + connectionId: "conn-test", + }; + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(pipelineCalls, 1); + assert.equal(executorCalls, 0); + assert.deepEqual(seenPolicy, { allowAccountRotation: true, allowModelFallback: true }); + assert.equal(result.kind, "ok"); + if (result.kind === "ok") { + assert.equal(result.response.choices[0].message.content, "from-pipeline"); + } +}); + +test("response body is parsed exactly once", async () => { + let parseCount = 0; + const responseBody = { + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "ok" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }; + const input = baseInput({ + executeProviderRequest: async () => makeExecutorResult(responseBody), + }); + const orig = input.executeProviderRequest; + input.executeProviderRequest = async (...args) => { + const result = await orig(...args); + const origText = result.response.text.bind(result.response); + result.response = { + ...result.response, + text: async () => { + parseCount++; + return origText(); + }, + } as unknown as Response; + return result; + }; + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "ok"); + assert.equal(parseCount, 1, "body should be parsed exactly once"); +}); + +test("setRequestWireState is called before executor with current wire state", async () => { + const wireStates: Array<{ translatedBody: unknown; effectiveModel: string }> = []; + const input = baseInput({ + setRequestWireState: (state) => wireStates.push(state), + }); + await runNonStreamingProviderLeg(input); + assert.ok(wireStates.length >= 1, "setRequestWireState should be called"); + assert.equal(wireStates[0].effectiveModel, "gpt-4o"); +}); + +test("buffered SSE response -> JSON body", async () => { + const sseBody = + 'data: {"id":"c1","choices":[{"delta":{"content":"hello"},"index":0,"finish_reason":null}]}\n\n' + + 'data: {"id":"c1","choices":[{"delta":{},"index":0,"finish_reason":"stop"}]}\n\n' + + "data: [DONE]\n\n"; + const sseResponse = { + ok: true, + status: 200, + statusText: "OK", + headers: new Headers({ "content-type": "text/event-stream" }), + text: async () => sseBody, + body: null, + } as unknown as Response; + const input = baseInput({ + executeProviderRequest: async () => ({ + response: sseResponse, + url: "https://api.example.com/v1/chat/completions", + headers: {}, + transformedBody: null, + }), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "ok"); +}); + +test("429 error -> error receipt with usage from body", async () => { + const errorBody = { + error: { message: "Rate limit exceeded", type: "rate_limit_error" }, + usage: { prompt_tokens: 5, completion_tokens: 0, total_tokens: 5 }, + }; + const input = baseInput({ + executeProviderRequest: async () => makeExecutorResult(errorBody, 429), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind !== "error") return; + assert.equal(result.result.status, 429); + assert.ok(result.usage, "error leg should include usage"); + assert.equal(result.receipt.httpStatus, 429); +}); + +test("500 error -> error receipt", async () => { + const input = baseInput({ + executeProviderRequest: async () => + makeExecutorResult({ error: { message: "Internal error" } }, 500), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind !== "error") return; + assert.equal(result.result.status, 500); + assert.equal(result.usage, null, "no usage in 500 error body"); +}); + +test("network throw -> error receipt with status 502", async () => { + const input = baseInput({ + executeProviderRequest: async () => { + throw new TypeError("fetch failed"); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind !== "error") return; + assert.ok(result.result.status >= 500, "should be 5xx"); + assert.equal(result.usage, null); +}); + +/* -- connection mismatch tests --------------------------------------------- */ + +test("expectedConnectionId mismatch -> 409 LEASE_CONNECTION_MISMATCH with receipt", async () => { + let executorCalled = false; + const input = baseInput({ + expectedConnectionId: "conn-abc", + getCurrentConnectionId: () => "conn-xyz", + executeProviderRequest: async () => { + executorCalled = true; + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "ok" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error", "should return error for connection mismatch"); + assert.equal(executorCalled, false, "mismatch at entry must not call the pipeline/executor"); + if (result.kind !== "error") return; + assert.equal(result.result.status, 409, "status must be 409"); + assert.equal( + result.result.errorCode, + "LEASE_CONNECTION_MISMATCH", + "errorCode must be LEASE_CONNECTION_MISMATCH" + ); + assert.equal(result.result.errorType, "lease_error", "errorType must be lease_error"); + assert.equal( + result.receipt.termination, + "connection_mismatch", + "receipt termination must be connection_mismatch" + ); + assert.equal(result.receipt.httpStatus, 409, "receipt httpStatus must be 409"); + assert.equal(result.usage, null, "usage must be null on mismatch"); +}); + +test("no expectedConnectionId -> mismatch check skipped, executor runs", async () => { + let executorCalled = false; + const input = baseInput({ + expectedConnectionId: undefined, + executeProviderRequest: async () => { + executorCalled = true; + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "ok" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "ok"); + assert.ok(executorCalled, "executor should be called when no expectedConnectionId"); +}); + +/* -- rotation guard tests -------------------------------------------------- */ + +test("follow-up Codex 429: rotationPolicy passed with allowAccountRotation=false, resolver call=0", async () => { + let resolverCallCount = 0; + let receivedPolicy: { allowAccountRotation?: boolean } | undefined; + const input = baseInput({ + phase: "follow-up", + allowAccountRotation: false, + provider: "codex", + executeProviderRequest: async (_model, _dedup, policy) => { + receivedPolicy = policy; + return makeExecutorResult( + { error: { message: "rate limited", type: "rate_limit_error" } }, + 429 + ); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + assert.deepEqual( + receivedPolicy, + { allowAccountRotation: false }, + "policy must have allowAccountRotation=false" + ); + assert.equal(resolverCallCount, 0, "resolver must not be called when rotation is disabled"); + if (result.kind === "error") { + assert.equal(result.result.status, 429); + } +}); + +test("initial Codex 429: rotationPolicy passed with allowAccountRotation=true, resolver>=1 and successful retry", async () => { + let receivedPolicy: { allowAccountRotation?: boolean } | undefined; + let executorCallCount = 0; + let resolverCallCount = 0; + const input = baseInput({ + phase: "initial", + allowAccountRotation: true, + provider: "codex", + executeProviderRequest: async (_model, _dedup, policy) => { + receivedPolicy = policy; + executorCallCount++; + if (executorCallCount === 1) { + return makeExecutorResult( + { error: { message: "rate limited", type: "rate_limit_error" } }, + 429 + ); + } + resolverCallCount++; + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "rotated" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.deepEqual( + receivedPolicy, + { allowAccountRotation: true }, + "policy must have allowAccountRotation=true" + ); + // Account rotation lives in the pipeline (6b), not this leg. The leg + // surfaces the 429 and forwards allowAccountRotation so the pipeline can + // retry. Claiming resolver>=1 here without a second execute was a false green. + assert.equal(result.kind, "error"); + assert.equal(executorCallCount, 1, "leg does not rotate; pipeline owns the retry"); + assert.equal(resolverCallCount, 0); + if (result.kind === "error") { + assert.equal(result.result.status, 429); + } +}); + +test("follow-up Antigravity 422: rotationPolicy passed with allowAccountRotation=false, resolver call=0", async () => { + let resolverCallCount = 0; + let receivedPolicy: { allowAccountRotation?: boolean } | undefined; + const input = baseInput({ + phase: "follow-up", + allowAccountRotation: false, + provider: "antigravity", + executeProviderRequest: async (_model, _dedup, policy) => { + receivedPolicy = policy; + return makeExecutorResult( + { error: { message: "gcp_project_required", type: "invalid_request" } }, + 422 + ); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + assert.deepEqual( + receivedPolicy, + { allowAccountRotation: false }, + "policy must have allowAccountRotation=false" + ); + assert.equal(resolverCallCount, 0, "resolver must not be called when rotation is disabled"); + if (result.kind === "error") { + assert.equal(result.result.status, 422); + } +}); + +test("initial Antigravity 422: rotationPolicy passed with allowAccountRotation=true", async () => { + let receivedPolicy: { allowAccountRotation?: boolean } | undefined; + const input = baseInput({ + phase: "initial", + allowAccountRotation: true, + provider: "antigravity", + executeProviderRequest: async (_model, _dedup, policy) => { + receivedPolicy = policy; + return makeExecutorResult( + { error: { message: "gcp_project_required", type: "invalid_request" } }, + 422 + ); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.deepEqual( + receivedPolicy, + { allowAccountRotation: true }, + "policy must have allowAccountRotation=true" + ); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 422); + } +}); + +/* -- side-effect tests ----------------------------------------------------- */ + +test("intermediate phase: leg does not apply client usage buffer", async () => { + const input = baseInput({ + phase: "intermediate", + executeProviderRequest: async () => + makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "partial" }, finish_reason: null }], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "ok"); + if (result.kind !== "ok") return; + // Intermediate phase should have raw usage (not buffered) + assert.ok(result.usage, "usage should be present"); + assert.equal(result.usage!.prompt_tokens, 10); +}); + +test("setRequestWireState receives updated body on fallback", async () => { + const wireStates: Array<{ translatedBody: Record; effectiveModel: string }> = []; + let executorCallCount = 0; + const input = baseInput({ + allowModelFallback: true, + provider: "gemini", + model: "gemini-3-pro", + setRequestWireState: (state) => + wireStates.push(state as { translatedBody: Record; effectiveModel: string }), + executeProviderRequest: async (_modelToCall) => { + executorCallCount++; + if (executorCallCount === 1) { + // First call: empty content -> triggers fallback + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 0, total_tokens: 1 }, + }); + } + // Fallback call + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "fallback" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const _result = await runNonStreamingProviderLeg(input); + // Fallback should have been attempted + assert.ok(executorCallCount >= 2, "executor should be called for fallback"); + // Wire state should have been updated with fallback model + const fallbackState = wireStates.find((s) => s.effectiveModel !== "gemini-3-pro"); + assert.ok(fallbackState, "setRequestWireState should have been called with fallback model"); +}); + +/* -- ClinePass tests ------------------------------------------------------- */ + +test("ClinePass retry on empty envelope: injected sleep, two calls, retry content/usage/finalBody/headers", async () => { + let retryCount = 0; + let sleepCalledWith: number[] = []; + const input = baseInput({ + provider: "clinepass", + sleep: async (ms) => { + sleepCalledWith.push(ms); + }, + executeProviderRequest: async () => { + retryCount++; + if (retryCount === 1) { + return makeExecutorResult({ success: false, error: "empty content" }); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "retried" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(retryCount, 2, "executor should be called twice (initial + retry)"); + assert.deepEqual(sleepCalledWith, [2000], "sleep should be called with 2000ms"); + assert.equal(result.kind, "ok", "retry should succeed"); + if (result.kind === "ok") { + assert.equal(result.response.choices[0].message.content, "retried"); + assert.ok(result.usage, "usage should be present from retry"); + assert.equal(result.usage!.prompt_tokens, 1); + assert.equal(result.usage!.completion_tokens, 1); + assert.equal(result.providerBody.choices[0].message.content, "retried"); + assert.ok(result.headers, "headers should be present"); + } +}); + +test("ClinePass retry: connection check before and after retry executor", async () => { + let connectionChecks: string[] = []; + let retryCount = 0; + const input = baseInput({ + phase: "follow-up", + expectedConnectionId: "conn-abc", + provider: "clinepass", + sleep: async () => {}, + getCurrentConnectionId: () => { + connectionChecks.push("check"); + return "conn-abc"; + }, + executeProviderRequest: async () => { + retryCount++; + if (retryCount === 1) { + return makeExecutorResult({ success: false, error: "empty content" }); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "ok" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "ok"); + // connection check happens: before initial executor + after initial executor + // + before retry + after retry = 4 checks at minimum + assert.ok( + connectionChecks.length >= 4, + `should have >=4 connection checks, got ${connectionChecks.length}` + ); +}); + +test("ClinePass retry: connection mismatch during retry -> 409", async () => { + let retryCount = 0; + const input = baseInput({ + phase: "follow-up", + expectedConnectionId: "conn-abc", + provider: "clinepass", + sleep: async () => {}, + getCurrentConnectionId: () => (retryCount <= 1 ? "conn-abc" : "conn-xyz"), + executeProviderRequest: async () => { + retryCount++; + if (retryCount === 1) { + return makeExecutorResult({ success: false, error: "empty content" }); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "ok" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 409); + assert.equal(result.result.errorCode, "LEASE_CONNECTION_MISMATCH"); + } +}); + +test("ClinePass non-empty error envelope returns error without retry", async () => { + let retryCount = 0; + const input = baseInput({ + provider: "clinepass", + executeProviderRequest: async () => { + retryCount++; + return makeExecutorResult({ success: false, error: "quota exceeded" }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(retryCount, 1, "executor should be called once (no retry for non-empty error)"); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 502); + assert.equal(result.receipt.termination, "provider_error"); + } +}); + +/* -- fallback tests -------------------------------------------------------- */ + +test("empty content initial with fallback enabled -> retries with next model", async () => { + let executorCallCount = 0; + const executedModels: string[] = []; + const input = baseInput({ + allowModelFallback: true, + provider: "gemini", + model: "gemini-3-pro", + executeProviderRequest: async (modelToCall) => { + executorCallCount++; + executedModels.push(modelToCall); + if (executorCallCount === 1) { + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 0, total_tokens: 1 }, + }); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "fallback" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const _result = await runNonStreamingProviderLeg(input); + assert.ok(executorCallCount >= 2, "should attempt fallback"); + assert.notEqual(executedModels[0], executedModels[1], "fallback should use different model"); +}); + +test("empty content follow-up with allowModelFallback=false -> error, no fallback", async () => { + let executorCallCount = 0; + const input = baseInput({ + phase: "follow-up", + allowModelFallback: false, + executeProviderRequest: async () => { + executorCallCount++; + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 0, total_tokens: 1 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(executorCallCount, 1, "executor should be called exactly once (no fallback)"); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 502); + assert.equal(result.usage, null); + } +}); + +test("model-unavailable follow-up with allowModelFallback=false -> error, no fallback", async () => { + let executorCallCount = 0; + const input = baseInput({ + phase: "follow-up", + allowModelFallback: false, + executeProviderRequest: async () => { + executorCallCount++; + return makeExecutorResult( + { error: { message: "model_not_found", type: "invalid_request_error" } }, + 404 + ); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(executorCallCount, 1, "executor should be called exactly once (no fallback)"); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 404); + } +}); + +test("context-overflow follow-up with allowModelFallback=false -> error, no fallback", async () => { + let executorCallCount = 0; + const input = baseInput({ + phase: "follow-up", + allowModelFallback: false, + executeProviderRequest: async () => { + executorCallCount++; + return makeExecutorResult({ error: { message: "maximum context length exceeded" } }, 400); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(executorCallCount, 1, "executor should be called exactly once (no fallback)"); + assert.equal(result.kind, "error"); +}); + +test("initial winning result preserves winning model/connection in receipt", async () => { + const input = baseInput({ + phase: "initial", + connectionId: "conn-win", + executeProviderRequest: async () => + makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "winner" }, finish_reason: "stop" }], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "ok"); + if (result.kind === "ok") { + assert.equal(result.connectionId, "conn-win", "receipt must carry winning connectionId"); + assert.equal(result.receipt.model, "gpt-4o", "receipt must carry winning model"); + assert.equal(result.receipt.termination, "completed"); + } +}); + +/* -- Retry-After header parsing ------------------------------------------- */ + +test("Retry-After header is parsed from upstream response into retryAfterMs", async () => { + const input = baseInput({ + executeProviderRequest: async () => + makeExecutorResult({ error: { message: "Rate limited", type: "rate_limit_error" } }, 429, { + "Retry-After": "30", + }), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 429); + assert.equal( + result.result.retryAfterMs, + 30_000, + "retryAfterMs should be parsed from Retry-After header" + ); + } +}); + +test("upstream error type is preserved in ChatCoreErrorResult", async () => { + const input = baseInput({ + executeProviderRequest: async () => + makeExecutorResult( + { + error: { + message: "Invalid request", + code: "invalid_request", + type: "invalid_request_error", + }, + }, + 400 + ), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 400); + assert.equal( + result.result.errorCode, + "invalid_request", + "upstream errorCode should be preserved" + ); + assert.equal( + result.result.errorType, + "invalid_request_error", + "upstream errorType should be preserved" + ); + } +}); + +/* -- connection mismatch before/after fallback ---------------------------- */ + +test("connection mismatch before fallback -> 409, no fallback executor call", async () => { + let executorCallCount = 0; + const input = baseInput({ + allowModelFallback: true, + provider: "gemini", + model: "gemini-3-pro", + expectedConnectionId: "conn-abc", + getCurrentConnectionId: () => (executorCallCount <= 1 ? "conn-abc" : "conn-xyz"), + executeProviderRequest: async (_modelToCall) => { + executorCallCount++; + if (executorCallCount === 1) { + return makeExecutorResult( + { error: { message: "model_not_found", type: "invalid_request_error" } }, + 404 + ); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "fallback" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 409); + assert.equal(result.result.errorCode, "LEASE_CONNECTION_MISMATCH"); + } +}); + +test("connection mismatch after fallback -> 409, fallback result discarded", async () => { + let executorCallCount = 0; + const input = baseInput({ + allowModelFallback: true, + provider: "gemini", + model: "gemini-3-pro", + expectedConnectionId: "conn-abc", + getCurrentConnectionId: () => (executorCallCount <= 1 ? "conn-abc" : "conn-xyz"), + executeProviderRequest: async (_modelToCall) => { + executorCallCount++; + if (executorCallCount === 1) { + return makeExecutorResult( + { error: { message: "model_not_found", type: "invalid_request_error" } }, + 404 + ); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "fallback" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 409); + assert.equal(result.result.errorCode, "LEASE_CONNECTION_MISMATCH"); + } +}); + +/* -- dynamic connection ID ------------------------------------------------ */ + +test("getCurrentConnectionId is read after initial executor, expected ID enforced", async () => { + let connectionChecks = 0; + const input = baseInput({ + expectedConnectionId: "conn-abc", + getCurrentConnectionId: () => { + connectionChecks++; + return "conn-abc"; + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "ok"); + assert.ok( + connectionChecks >= 1, + "getCurrentConnectionId should be called at least once after executor" + ); +}); + +test("dynamic connection: ID changes between initial and retry -> 409 on retry path", async () => { + let executorCallCount = 0; + const input = baseInput({ + phase: "follow-up", + expectedConnectionId: "conn-abc", + provider: "clinepass", + sleep: async () => {}, + getCurrentConnectionId: () => (executorCallCount < 1 ? "conn-abc" : "conn-xyz"), + executeProviderRequest: async () => { + executorCallCount++; + if (executorCallCount === 1) { + return makeExecutorResult({ success: false, error: "empty content" }); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "ok" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind === "error") { + assert.equal(result.result.status, 409); + assert.equal(result.result.errorCode, "LEASE_CONNECTION_MISMATCH"); + } + assert.equal( + executorCallCount, + 1, + "retry executor must not run after the lease already moved" + ); +}); + +/* -- fallback with real parsed response ----------------------------------- */ + +test("model-unavailable fallback returns real parsed response, usage, headers, winning connection/model", async () => { + let executorCallCount = 0; + const input = baseInput({ + allowModelFallback: true, + provider: "gemini", + model: "gemini-3-pro", + connectionId: "conn-initial", + getCurrentConnectionId: () => "conn-initial", + executeProviderRequest: async (_modelToCall) => { + executorCallCount++; + if (executorCallCount === 1) { + return makeExecutorResult( + { error: { message: "model_not_found", type: "invalid_request_error" } }, + 404 + ); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [ + { message: { role: "assistant", content: "fallback response" }, finish_reason: "stop" }, + ], + usage: { prompt_tokens: 20, completion_tokens: 10, total_tokens: 30 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.ok(executorCallCount >= 2, "should attempt fallback"); + assert.equal(result.kind, "ok", "fallback should succeed"); + if (result.kind === "ok") { + assert.equal(result.response.choices[0].message.content, "fallback response"); + assert.ok(result.usage, "usage should be present from fallback"); + assert.equal(result.usage!.prompt_tokens, 20); + assert.equal(result.usage!.completion_tokens, 10); + assert.equal(result.providerBody.choices[0].message.content, "fallback response"); + assert.ok(result.headers, "headers should be present"); + assert.equal(result.connectionId, "conn-initial", "winning connection should be in receipt"); + // The actual fallback model from getNextFamilyFallback + assert.equal( + result.receipt.model, + "gemini-3.1-pro-preview", + "winning model should be in receipt" + ); + } +}); + +test("empty content fallback returns real parsed response, usage, headers", async () => { + let executorCallCount = 0; + const input = baseInput({ + allowModelFallback: true, + provider: "gemini", + model: "gemini-3-pro", + executeProviderRequest: async (_modelToCall) => { + executorCallCount++; + if (executorCallCount === 1) { + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 0, total_tokens: 1 }, + }); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [ + { message: { role: "assistant", content: "fallback from empty" }, finish_reason: "stop" }, + ], + usage: { prompt_tokens: 15, completion_tokens: 8, total_tokens: 23 }, + }); + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.ok(executorCallCount >= 2, "should attempt fallback"); + assert.equal(result.kind, "ok", "fallback should succeed"); + if (result.kind === "ok") { + assert.equal(result.response.choices[0].message.content, "fallback from empty"); + assert.ok(result.usage, "usage should be present from fallback"); + assert.equal(result.usage!.prompt_tokens, 15); + assert.equal(result.usage!.completion_tokens, 8); + assert.equal(result.providerBody.choices[0].message.content, "fallback from empty"); + } +}); + +/* -- setRequestWireState before every executor ---------------------------- */ + +test("setRequestWireState is called before initial, retry, and fallback executors", async () => { + const wireStates: string[] = []; + let executorCallCount = 0; + const input = baseInput({ + allowModelFallback: true, + provider: "gemini", + model: "gemini-3-pro", + setRequestWireState: (state) => wireStates.push(state.effectiveModel), + executeProviderRequest: async (_modelToCall) => { + executorCallCount++; + if (executorCallCount === 1) { + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 0, total_tokens: 1 }, + }); + } + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "ok" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }); + }, + }); + await runNonStreamingProviderLeg(input); + // wire state should be set: initial + fallback = at least 2 + assert.ok( + wireStates.length >= 2, + `setRequestWireState should be called >=2 times, got ${wireStates.length}` + ); + assert.equal(wireStates[0], "gemini-3-pro", "first call should be initial model"); + assert.notEqual(wireStates[1], "gemini-3-pro", "second call should be fallback model"); +}); + +test("invalid SSE payload returns errorCode invalid_sse_payload, not upstream_error", async () => { + const sseBody = 'data: {"error":{"message":"Devin CLI not found"}}\n\n'; + const input = baseInput({ + executeProviderRequest: async () => ({ + response: { + ok: true, + status: 200, + statusText: "OK", + headers: new Headers({ "content-type": "text/event-stream" }), + text: async () => sseBody, + body: null, + } as unknown as Response, + url: "https://api.example.com/v1/chat/completions", + headers: {}, + transformedBody: null, + }), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind !== "error") return; + assert.equal(result.result.status, 502); + assert.equal(result.result.errorCode, "invalid_sse_payload"); + assert.equal(result.result.errorType, "invalid_sse_payload"); + assert.notEqual(result.result.errorCode, "upstream_error"); +}); + +test("invalid JSON payload returns errorCode invalid_json_payload, not upstream_error", async () => { + const input = baseInput({ + executeProviderRequest: async () => ({ + response: { + ok: true, + status: 200, + statusText: "OK", + headers: new Headers({ "content-type": "application/json" }), + text: async () => "not-json{{{", + body: null, + } as unknown as Response, + url: "https://api.example.com/v1/chat/completions", + headers: {}, + transformedBody: null, + }), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "error"); + if (result.kind !== "error") return; + assert.equal(result.result.status, 502); + assert.equal(result.result.errorCode, "invalid_json_payload"); + assert.equal(result.result.errorType, "invalid_json_payload"); + assert.notEqual(result.result.errorCode, "upstream_error"); +}); + +test("empty-content fallback with invalid SSE body is 502, not 200 empty", async () => { + const sseBody = 'data: {"error":{"message":"Devin CLI not found"}}\n\n'; + let executorCallCount = 0; + const input = baseInput({ + allowModelFallback: true, + provider: "gemini", + model: "gemini-3-pro", + executeProviderRequest: async () => { + executorCallCount++; + if (executorCallCount === 1) { + return makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ message: { role: "assistant", content: "" }, finish_reason: "stop" }], + usage: { prompt_tokens: 1, completion_tokens: 0, total_tokens: 1 }, + }); + } + return { + response: { + ok: true, + status: 200, + statusText: "OK", + headers: new Headers({ "content-type": "text/event-stream" }), + text: async () => sseBody, + body: null, + } as unknown as Response, + url: "https://api.example.com/v1/chat/completions", + headers: {}, + transformedBody: null, + }; + }, + }); + const result = await runNonStreamingProviderLeg(input); + assert.ok(executorCallCount >= 2, "should attempt fallback"); + assert.equal(result.kind, "error", "invalid SSE on fallback must not finishOk the empty original"); + if (result.kind !== "error") return; + assert.equal(result.result.status, 502); + assert.equal(result.result.errorCode, "invalid_sse_payload"); +}); + +test("finishOk caches reasoning against translatedBody.messages, not Responses input", async () => { + clearReasoningCacheAll(); + const scope = "api-key:leg-history"; + const historyMessages = [{ role: "user", content: "hi from translatedBody" }]; + const assistantMessage = { + role: "assistant", + content: "thinking result", + reasoning_content: "let me think...", + }; + const input = baseInput({ + provider: "deepseek", + model: "deepseek-v4-pro", + reasoningCacheScope: scope, + translatedBody: { messages: historyMessages }, + sourceBody: { input: [{ role: "user", content: "wrong body" }] }, + executeProviderRequest: async () => + makeExecutorResult({ + id: "chatcmpl-test", + choices: [{ index: 0, message: assistantMessage, finish_reason: "stop" }], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }), + }); + const result = await runNonStreamingProviderLeg(input); + assert.equal(result.kind, "ok"); + const cacheKey = buildAssistantMessageCacheKey( + scope, + [...historyMessages, assistantMessage], + historyMessages.length + ); + assert.equal( + lookupReasoning(cacheKey), + "let me think...", + "finishOk must pass historyMessages from translatedBody.messages" + ); +}); + +test("semaphore capacity errors rethrow so chatCore can map them to 429", async () => { + const timeout = Object.assign(new Error("Semaphore timeout"), { code: "SEMAPHORE_TIMEOUT" }); + const input = baseInput({ + executeProviderRequest: async () => { + throw timeout; + }, + }); + await assert.rejects( + () => runNonStreamingProviderLeg(input), + (err: unknown) => { + assert.equal((err as { code?: string }).code, "SEMAPHORE_TIMEOUT"); + return true; + } + ); +}); diff --git a/tests/unit/provider-execution-pipeline.test.ts b/tests/unit/provider-execution-pipeline.test.ts new file mode 100644 index 0000000000..94b60c0156 --- /dev/null +++ b/tests/unit/provider-execution-pipeline.test.ts @@ -0,0 +1,632 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import type { + ChatCoreExecutorResult, + PipelineConnectionContext, + PipelineStateHooks, + PipelineTargetContext, + PipelineWireState, + ProviderExecutionPipelineInput, + ProviderExecutionPolicy, +} from "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts"; + +test("runProviderExecutionPipeline is importable", async () => { + const mod = await import("../../open-sse/handlers/chatCore/providerExecutionPipeline.ts"); + assert.equal(typeof mod.runProviderExecutionPipeline, "function"); +}); + +function jsonResponse(body: unknown, status: number, extraHeaders: Record = {}) { + return new Response(JSON.stringify(body), { + status, + headers: { "content-type": "application/json", ...extraHeaders }, + }); +} + +function makeAttempt( + body: unknown, + status: number, + extra: Partial = {} +): ChatCoreExecutorResult { + const response = jsonResponse(body, status, extra.headers as Record | undefined); + return { + response, + url: extra.url ?? "https://upstream.test/v1/chat/completions", + headers: extra.headers ?? { "content-type": "application/json" }, + transformedBody: extra.transformedBody ?? { model: "gpt-5" }, + ...extra, + }; +} + +function noopState(): PipelineStateHooks { + return { + updatePendingStage: () => {}, + recordRateLimitHeaders: () => {}, + recordRateLimitBody: () => {}, + writeTerminalStatus: async () => {}, + persistConnectionPatch: () => {}, + setConnectionRateLimitedUntil: () => {}, + lockModel: () => {}, + recordAntigravityQuotaState: async () => {}, + markAccountSemaphoreBlocked: () => {}, + isolateProbeFailures: () => false, + }; +} + +function makeInput(opts: { + policy: ProviderExecutionPolicy; + provider: string; + model?: string; + stream?: boolean; + connectionId?: string; + send: (model: string, allowDedup: boolean) => Promise; + getProviderCredentials?: PipelineConnectionContext["getProviderCredentials"]; + replaceCredentials?: PipelineConnectionContext["replaceCredentials"]; + getCurrentConnectionId?: () => string | undefined; + refreshCredentials?: PipelineConnectionContext["refreshCredentials"]; + onCredentialsRefreshed?: PipelineConnectionContext["onCredentialsRefreshed"]; + getNextFamilyFallback?: ProviderExecutionPipelineInput["getNextFamilyFallback"]; + state?: Partial; +}): ProviderExecutionPipelineInput { + const model = opts.model ?? "gpt-5"; + const connectionId = opts.connectionId ?? "conn-a"; + let currentId: string | undefined = connectionId; + let credentials: Record = { connectionId }; + const target: PipelineTargetContext = { + provider: opts.provider, + requestedModel: model, + sourceFormat: "openai", + targetFormat: "openai", + stream: opts.stream ?? false, + }; + const wire: PipelineWireState = { + body: { model, messages: [{ role: "user", content: "hi" }] }, + currentModel: model, + triedModels: new Set([model]), + setBodyAndModel: (body, nextModel) => { + wire.body = body; + wire.currentModel = nextModel; + wire.triedModels.add(nextModel); + }, + }; + const connection: PipelineConnectionContext = { + initialConnectionId: connectionId, + getCurrentConnectionId: opts.getCurrentConnectionId ?? (() => currentId), + getCredentials: () => credentials, + replaceCredentials: + opts.replaceCredentials ?? + ((next) => { + credentials = next; + currentId = typeof next.connectionId === "string" ? next.connectionId : currentId; + }), + onCredentialsRefreshed: opts.onCredentialsRefreshed ?? (() => {}), + assertManagedLeaseFence: () => {}, + refreshCredentials: opts.refreshCredentials, + getProviderCredentials: + opts.getProviderCredentials ?? + (async () => { + throw new Error("getProviderCredentials must not be called in this fixture"); + }), + }; + return { + policy: opts.policy, + target, + connection, + wire, + state: { ...noopState(), ...(opts.state || {}) }, + sendProviderAttempt: opts.send, + getNextFamilyFallback: opts.getNextFamilyFallback, + }; +} + +test("initial Codex 429: rotation resolver>=1 and successful retry", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + let sendCount = 0; + let resolverCallCount = 0; + const input = makeInput({ + policy: { allowAccountRotation: true, allowModelFallback: true, expectedConnectionId: undefined }, + provider: "codex", + connectionId: "conn-a", + send: async () => { + sendCount += 1; + if (sendCount === 1) { + return makeAttempt({ error: { message: "rate limited", type: "rate_limit_error" } }, 429, { + headers: { "retry-after": "1" }, + }); + } + return makeAttempt({ + id: "chatcmpl-ok", + choices: [{ message: { role: "assistant", content: "rotated" }, finish_reason: "stop" }], + }, 200); + }, + getProviderCredentials: (async () => { + resolverCallCount += 1; + return { connectionId: "conn-b", allRateLimited: false }; + }) as PipelineConnectionContext["getProviderCredentials"], + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(resolverCallCount >= 1, true, "resolver must run on initial Codex 429"); + assert.equal(sendCount, 2, "second send after rotation"); + assert.equal(outcome.kind, "response"); + if (outcome.kind === "response") { + assert.equal(outcome.connectionId, "conn-b"); + assert.equal(outcome.response.status, 200); + } +}); + +test("initial Antigravity 422 gcp_project_required: rotation resolver>=1 and successful retry", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + let sendCount = 0; + let resolverCallCount = 0; + const input = makeInput({ + policy: { allowAccountRotation: true, allowModelFallback: true }, + provider: "antigravity", + connectionId: "agy-a", + send: async () => { + sendCount += 1; + if (sendCount === 1) { + return makeAttempt({ error: { message: "gcp_project_required", type: "invalid_request" } }, 422); + } + return makeAttempt( + { + id: "chatcmpl-ok", + choices: [{ message: { role: "assistant", content: "rotated" }, finish_reason: "stop" }], + }, + 200 + ); + }, + getProviderCredentials: (async () => { + resolverCallCount += 1; + return { connectionId: "agy-b", allRateLimited: false }; + }) as PipelineConnectionContext["getProviderCredentials"], + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(resolverCallCount >= 1, true, "resolver must run on initial Antigravity BYOP 422"); + assert.equal(sendCount, 2, "second send after BYOP rotation"); + assert.equal(outcome.kind, "response"); + if (outcome.kind === "response") { + assert.equal(outcome.connectionId, "agy-b"); + assert.equal(outcome.response.status, 200); + } +}); + +test("follow-up rotation blocks resolver on Antigravity 422", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + let sendCount = 0; + let resolverCallCount = 0; + const input = makeInput({ + policy: { + allowAccountRotation: false, + allowModelFallback: false, + expectedConnectionId: "agy-a", + }, + provider: "antigravity", + connectionId: "agy-a", + send: async () => { + sendCount += 1; + return makeAttempt({ error: { message: "gcp_project_required", type: "invalid_request" } }, 422); + }, + getProviderCredentials: (async () => { + resolverCallCount += 1; + return { connectionId: "agy-b", allRateLimited: false }; + }) as PipelineConnectionContext["getProviderCredentials"], + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(resolverCallCount, 0, "follow-up must not call credentials resolver"); + assert.equal(sendCount, 1, "follow-up sends once"); + assert.equal(outcome.kind, "error"); + if (outcome.kind === "error") { + assert.equal(outcome.result.status, 422); + assert.equal(outcome.connectionId, "agy-a"); + } +}); + +test("follow-up rotation blocks resolver on Codex 429", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + let sendCount = 0; + let resolverCallCount = 0; + const input = makeInput({ + policy: { + allowAccountRotation: false, + allowModelFallback: false, + expectedConnectionId: "conn-a", + }, + provider: "codex", + connectionId: "conn-a", + send: async () => { + sendCount += 1; + return makeAttempt({ error: { message: "rate limited", type: "rate_limit_error" } }, 429); + }, + getProviderCredentials: (async () => { + resolverCallCount += 1; + return { connectionId: "conn-b", allRateLimited: false }; + }) as PipelineConnectionContext["getProviderCredentials"], + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(resolverCallCount, 0, "follow-up must not call credentials resolver"); + assert.equal(sendCount, 1, "follow-up sends once"); + assert.equal(outcome.kind, "error"); + if (outcome.kind === "error") { + assert.equal(outcome.result.status, 429); + assert.equal(outcome.connectionId, "conn-a"); + } +}); + +test("401 refresh succeeds then retries once on same connection", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + let sendCount = 0; + let refreshCount = 0; + let persistCount = 0; + let resolverCallCount = 0; + const input = makeInput({ + policy: { allowAccountRotation: true, allowModelFallback: true }, + provider: "openai", + connectionId: "conn-a", + send: async () => { + sendCount += 1; + if (sendCount === 1) { + return makeAttempt({ error: { message: "invalid_api_key", type: "authentication_error" } }, 401); + } + return makeAttempt( + { + id: "chatcmpl-ok", + choices: [{ message: { role: "assistant", content: "refreshed" }, finish_reason: "stop" }], + }, + 200 + ); + }, + refreshCredentials: async (creds) => { + refreshCount += 1; + return { ...creds, accessToken: "new-token" }; + }, + onCredentialsRefreshed: async () => { + persistCount += 1; + }, + getProviderCredentials: (async () => { + resolverCallCount += 1; + return { connectionId: "conn-b", allRateLimited: false }; + }) as PipelineConnectionContext["getProviderCredentials"], + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(refreshCount, 1, "refresh once"); + assert.equal(persistCount, 1, "onCredentialsRefreshed once"); + assert.equal(resolverCallCount, 0, "401 refresh must not rotate accounts"); + assert.equal(sendCount, 2, "retry once after refresh"); + assert.equal(outcome.kind, "response"); + if (outcome.kind === "response") { + assert.equal(outcome.connectionId, "conn-a"); + assert.equal(outcome.response.status, 200); + } +}); + +test("status restatement rewrites agentrouter 403 quota exhaustion to 429 before classification", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + let sendCount = 0; + const input = makeInput({ + policy: { allowAccountRotation: true, allowModelFallback: true }, + provider: "agentrouter", + connectionId: "ar-a", + send: async () => { + sendCount += 1; + return makeAttempt({ error: { message: "用户额度不足", type: "forbidden" } }, 403); + }, + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(sendCount, 1); + assert.equal(outcome.kind, "error"); + if (outcome.kind === "error") { + assert.equal(outcome.result.status, 429, "restated before classification"); + assert.equal(outcome.connectionId, "ar-a"); + } +}); + +test("thinking-signature recovery returns winning response", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + let sendCount = 0; + const input = makeInput({ + policy: { allowAccountRotation: true, allowModelFallback: true }, + provider: "claude", + connectionId: "cl-a", + send: async () => { + sendCount += 1; + if (sendCount === 1) { + return makeAttempt( + { error: { message: "invalid signature in thinking block", type: "invalid_request_error" } }, + 400 + ); + } + return makeAttempt( + { + id: "msg-ok", + type: "message", + role: "assistant", + content: [{ type: "text", text: "recovered" }], + }, + 200 + ); + }, + }); + input.wire.body = { + model: "gpt-5", + messages: [ + { role: "user", content: "q1" }, + { + role: "assistant", + content: [ + { type: "thinking", thinking: "old" }, + { type: "text", text: "a1" }, + ], + }, + { role: "user", content: "q2" }, + ], + }; + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(sendCount, 2, "one recovery send after signature error"); + assert.equal(outcome.kind, "response"); + if (outcome.kind === "response") { + assert.equal(outcome.response.status, 200); + assert.equal(outcome.connectionId, "cl-a"); + } +}); + +test("initial model-unavailable falls back to sibling model", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + let sendCount = 0; + let fallbackLookupCount = 0; + const sentModels: string[] = []; + const input = makeInput({ + policy: { allowAccountRotation: true, allowModelFallback: true }, + provider: "openai", + model: "gpt-5", + connectionId: "conn-a", + send: async (model) => { + sendCount += 1; + sentModels.push(model); + if (model === "gpt-5") { + return makeAttempt( + { error: { message: "model is not available", type: "invalid_request_error" } }, + 404 + ); + } + return makeAttempt( + { + id: "chatcmpl-ok", + choices: [{ message: { role: "assistant", content: "fallback" }, finish_reason: "stop" }], + }, + 200 + ); + }, + getNextFamilyFallback: (current) => { + fallbackLookupCount += 1; + return current === "gpt-5" ? "gpt-5-mini" : null; + }, + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(fallbackLookupCount >= 1, true, "family fallback consulted"); + assert.deepEqual(sentModels, ["gpt-5", "gpt-5-mini"]); + assert.equal(sendCount, 2); + assert.equal(outcome.kind, "response"); + if (outcome.kind === "response") { + assert.equal(outcome.model, "gpt-5-mini"); + assert.equal(outcome.response.status, 200); + } +}); + +test("follow-up allowModelFallback=false blocks model-unavailable fallback", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + let sendCount = 0; + let fallbackLookupCount = 0; + const input = makeInput({ + policy: { + allowAccountRotation: false, + allowModelFallback: false, + expectedConnectionId: "conn-a", + }, + provider: "openai", + model: "gpt-5", + connectionId: "conn-a", + send: async () => { + sendCount += 1; + return makeAttempt( + { error: { message: "model is not available", type: "invalid_request_error" } }, + 404 + ); + }, + getNextFamilyFallback: () => { + fallbackLookupCount += 1; + return "gpt-5-mini"; + }, + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(fallbackLookupCount, 0, "follow-up must not consult family fallback"); + assert.equal(sendCount, 1); + assert.equal(outcome.kind, "error"); + if (outcome.kind === "error") { + assert.equal(outcome.result.status, 404); + assert.equal(outcome.model, "gpt-5"); + } +}); + +test("Codex 429 rotation calls scope-rate-limit, affinity-clear, and audit hooks", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + const rateLimited: Array> = []; + const affinityCleared: string[] = []; + const audits: Array> = []; + let sendCount = 0; + const input = makeInput({ + policy: { allowAccountRotation: true, allowModelFallback: true }, + provider: "codex", + connectionId: "conn-a", + send: async () => { + sendCount += 1; + if (sendCount === 1) { + return makeAttempt({ error: { message: "rate limited", type: "rate_limit_error" } }, 429, { + headers: { "retry-after": "2" }, + }); + } + return makeAttempt( + { + id: "chatcmpl-ok", + choices: [{ message: { role: "assistant", content: "rotated" }, finish_reason: "stop" }], + }, + 200 + ); + }, + getProviderCredentials: (async () => ({ + connectionId: "conn-b", + allRateLimited: false, + })) as PipelineConnectionContext["getProviderCredentials"], + state: { + onCodexScopeRateLimited: (params) => { + rateLimited.push(params as unknown as Record); + }, + onClearSessionAffinity: (params) => { + affinityCleared.push(params.failedConnectionId); + }, + onAuditAccountRotation: (params) => { + audits.push(params as unknown as Record); + }, + }, + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(outcome.kind, "response"); + assert.equal(rateLimited.length, 1, "must persist Codex model-scope cooldown"); + assert.equal(rateLimited[0]?.failedConnectionId, "conn-a"); + assert.deepEqual(affinityCleared, ["conn-a"]); + assert.equal(audits.length, 1); + assert.equal(audits[0]?.action, "codex.account_rotation"); + assert.equal(audits[0]?.failedConnectionId, "conn-a"); + assert.equal(audits[0]?.newConnectionId, "conn-b"); +}); + +test("Codex 429 cooldown reads Retry-After from the response, not request headers", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + const rateLimited: Array> = []; + let sendCount = 0; + const input = makeInput({ + policy: { allowAccountRotation: true, allowModelFallback: true }, + provider: "codex", + connectionId: "conn-a", + send: async () => { + sendCount += 1; + if (sendCount === 1) { + // BaseExecutor puts REQUEST headers on attempt.headers (Authorization). + // Upstream Retry-After lives on the Response. Mixing the two bags is the + // extract regression: cooldown silently falls back to 60s. + return { + response: jsonResponse( + { error: { message: "rate limited", type: "rate_limit_error" } }, + 429, + { "Retry-After": "5" } + ), + url: "https://upstream.test/v1/chat/completions", + headers: { Authorization: "Bearer request-token", "content-type": "application/json" }, + transformedBody: { model: "gpt-5" }, + }; + } + return makeAttempt( + { + id: "chatcmpl-ok", + choices: [{ message: { role: "assistant", content: "rotated" }, finish_reason: "stop" }], + }, + 200 + ); + }, + getProviderCredentials: (async () => ({ + connectionId: "conn-b", + allRateLimited: false, + })) as PipelineConnectionContext["getProviderCredentials"], + state: { + onCodexScopeRateLimited: (params) => { + rateLimited.push(params as unknown as Record); + }, + }, + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(outcome.kind, "response"); + assert.equal(rateLimited.length, 1, "must persist Codex model-scope cooldown"); + const until = new Date(String(rateLimited[0]?.rateLimitedUntil)).getTime(); + const delta = until - Date.now(); + assert.ok( + delta > 4_000 && delta < 8_000, + `Retry-After: 5 must yield ~5s cooldown, got ${delta}ms (60s = still reading request headers)` + ); +}); + +test("Antigravity BYOP 422 rotation persists cooldown via setConnectionRateLimitedUntil", async () => { + const { runProviderExecutionPipeline } = await import( + "../../open-sse/handlers/chatCore/providerExecutionPipeline.ts" + ); + const cooldowns: Array<{ id: string; untilMs: number | null }> = []; + let sendCount = 0; + const input = makeInput({ + policy: { allowAccountRotation: true, allowModelFallback: true }, + provider: "antigravity", + connectionId: "agy-a", + send: async () => { + sendCount += 1; + if (sendCount === 1) { + return makeAttempt( + { error: { message: "gcp_project_required", type: "invalid_request" } }, + 422 + ); + } + return makeAttempt( + { + id: "chatcmpl-ok", + choices: [{ message: { role: "assistant", content: "rotated" }, finish_reason: "stop" }], + }, + 200 + ); + }, + getProviderCredentials: (async () => ({ + connectionId: "agy-b", + allRateLimited: false, + })) as PipelineConnectionContext["getProviderCredentials"], + state: { + setConnectionRateLimitedUntil: (id, untilMs) => { + cooldowns.push({ id, untilMs }); + }, + }, + }); + + const outcome = await runProviderExecutionPipeline(input); + assert.equal(outcome.kind, "response"); + assert.equal(sendCount, 2); + assert.equal(cooldowns.length, 1, "BYOP rotate must persist cooldown before picking sibling"); + assert.equal(cooldowns[0]?.id, "agy-a"); + assert.equal(typeof cooldowns[0]?.untilMs, "number"); + assert.equal((cooldowns[0]?.untilMs ?? 0) > Date.now(), true); +}); diff --git a/tests/unit/request-logger-endpoints.test.ts b/tests/unit/request-logger-endpoints.test.ts index 66823c927c..968542c6d6 100644 --- a/tests/unit/request-logger-endpoints.test.ts +++ b/tests/unit/request-logger-endpoints.test.ts @@ -738,3 +738,52 @@ test("generic client snapshots exclude hard-lease control headers", async () => assert.equal(out.headers["x-omniroute-lease-generation"], undefined); assert.equal(out.headers["x-session-id"], "independent-routing-session"); }); + +function syntheticReceipt(index: number) { + return { + index, + connectionId: "conn-1", + provider: "openai", + model: "gpt-4o", + startedAt: "2026-01-01T00:00:00.000Z", + endedAt: "2026-01-01T00:00:01.000Z", + latencyMs: 10 + index, + httpStatus: 200, + errorType: null, + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + serviceTier: null, + computedCostUsd: 0.001, + toolCalls: [{ id: `call-${index}`, name: "memory_search" }], + termination: "completed", + clientVisible: true, + }; +} + +test("logToolLoopReceipt keeps first 4 receipts and clones them", async () => { + const { createRequestLogger } = await import("../../open-sse/utils/requestLogger.ts"); + const logger = await createRequestLogger("openai", "openai", "gpt-4o", { + enabled: true, + captureStreamChunks: false, + }); + for (let i = 0; i < 5; i++) { + logger.logToolLoopReceipt(syntheticReceipt(i)); + } + const payloads = logger.getPipelinePayloads(); + assert.ok(payloads?.toolLoop); + assert.equal(payloads.toolLoop.legs.length, 4); + assert.deepEqual( + payloads.toolLoop.legs.map((leg) => (leg as { index: number }).index), + [0, 1, 2, 3] + ); + assert.equal("arguments" in (payloads.toolLoop.legs[0] as object), false); +}); + +test("logToolLoopReceipt is a no-op when logger is disabled", async () => { + const { createRequestLogger } = await import("../../open-sse/utils/requestLogger.ts"); + const logger = await createRequestLogger("openai", "openai", "gpt-4o", { + enabled: false, + captureStreamChunks: false, + }); + logger.logToolLoopReceipt(syntheticReceipt(0)); + assert.equal(logger.getPipelinePayloads(), null); +}); diff --git a/tests/unit/server-owned-tool-loop-flag.test.ts b/tests/unit/server-owned-tool-loop-flag.test.ts new file mode 100644 index 0000000000..2f6a383044 --- /dev/null +++ b/tests/unit/server-owned-tool-loop-flag.test.ts @@ -0,0 +1,140 @@ +import { describe, it, before, after, beforeEach } from "node:test"; +import assert from "node:assert/strict"; +import os from "node:os"; +import path from "node:path"; +import fs from "node:fs"; +import { fileURLToPath } from "node:url"; + +const __dirname = path.dirname(fileURLToPath(import.meta.url)); + +const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-test-flag-loop-")); +process.env.DATA_DIR = tmpDir; + +const { FEATURE_FLAG_DEFINITIONS } = + await import("../../src/shared/constants/featureFlagDefinitions.ts"); +const { setFeatureFlagOverride, clearAllFeatureFlagOverrides } = + await import("../../src/lib/db/featureFlags.ts"); +const { isServerOwnedToolLoopEnabled } = await import("../../src/shared/utils/featureFlags.ts"); + +describe("SERVER_OWNED_TOOL_LOOP_ENABLED flag definition", () => { + it("exists in FEATURE_FLAG_DEFINITIONS", () => { + const def = FEATURE_FLAG_DEFINITIONS.find((d) => d.key === "SERVER_OWNED_TOOL_LOOP_ENABLED"); + assert.ok(def, "SERVER_OWNED_TOOL_LOOP_ENABLED should exist"); + assert.equal(def.category, "runtime"); + assert.equal(def.defaultValue, "false"); + assert.equal(def.requiresRestart, false); + assert.equal(def.descriptionI18nKey, "featureFlagServerOwnedToolLoopDescription"); + }); +}); + +describe("isServerOwnedToolLoopEnabled wrapper", () => { + beforeEach(() => { + clearAllFeatureFlagOverrides(); + }); + + it("returns false when no override is set (default)", () => { + assert.equal(isServerOwnedToolLoopEnabled(), false); + }); + + it("returns true when DB override is set to true", () => { + setFeatureFlagOverride("SERVER_OWNED_TOOL_LOOP_ENABLED", "true"); + assert.equal(isServerOwnedToolLoopEnabled(), true); + }); + + it("returns false and logs when injected reader throws", () => { + const logs: unknown[] = []; + const origError = console.error; + console.error = (...args: unknown[]) => { + logs.push(args); + }; + try { + const throwingReader = () => { + throw new Error("flag read failed"); + }; + const result = isServerOwnedToolLoopEnabled(throwingReader); + assert.equal(result, false); + assert.ok(logs.length >= 1, "console.error should be called at least once"); + assert.ok( + logs.some((args) => + String(args).includes("Failed to resolve SERVER_OWNED_TOOL_LOOP_ENABLED") + ), + "error log should mention the flag key" + ); + } finally { + console.error = origError; + } + }); +}); + +describe("feature-flags-settings count update", () => { + it("flag count matches updated expected value", () => { + assert.equal(FEATURE_FLAG_DEFINITIONS.length, 55); + }); +}); + +describe("i18n key parity for SERVER_OWNED_TOOL_LOOP_ENABLED", () => { + let enMessages: Record; + let ptBrMessages: Record; + + before(async () => { + const enRaw = fs.readFileSync( + path.resolve(__dirname, "../../src/i18n/messages/en.json"), + "utf8" + ); + enMessages = JSON.parse(enRaw); + const ptBrRaw = fs.readFileSync( + path.resolve(__dirname, "../../src/i18n/messages/pt-BR.json"), + "utf8" + ); + ptBrMessages = JSON.parse(ptBrRaw); + }); + + it("en.json has nested featureFlags.definitions.SERVER_OWNED_TOOL_LOOP_ENABLED.label", () => { + const defs = enMessages.featureFlags as Record | undefined; + assert.ok(defs, "en.json should have featureFlags section"); + const definitions = (defs as Record).definitions as + Record | undefined; + assert.ok(definitions, "en.json featureFlags should have definitions"); + const flagDef = definitions.SERVER_OWNED_TOOL_LOOP_ENABLED as + Record | undefined; + assert.ok(flagDef, "definitions should contain SERVER_OWNED_TOOL_LOOP_ENABLED"); + assert.equal(flagDef.label, "Server-Owned Tool Loop"); + assert.equal( + flagDef.description, + "Continue non-streaming server-owned tool calls until the model returns a client-usable response." + ); + }); + + it("pt-BR.json has nested featureFlags.definitions.SERVER_OWNED_TOOL_LOOP_ENABLED.label", () => { + const defs = ptBrMessages.featureFlags as Record | undefined; + assert.ok(defs, "pt-BR.json should have featureFlags section"); + const definitions = (defs as Record).definitions as + Record | undefined; + assert.ok(definitions, "pt-BR.json featureFlags should have definitions"); + const flagDef = definitions.SERVER_OWNED_TOOL_LOOP_ENABLED as + Record | undefined; + assert.ok(flagDef, "definitions should contain SERVER_OWNED_TOOL_LOOP_ENABLED"); + assert.equal(flagDef.label, "Server-Owned Tool Loop"); + assert.equal(typeof flagDef.description, "string"); + assert.ok( + ((flagDef.description as string) || "").length > 0, + "description should be non-empty" + ); + }); + + it("en.json does NOT have stale top-level featureFlagServerOwnedToolLoopDescription", () => { + assert.equal( + (enMessages as Record).featureFlagServerOwnedToolLoopDescription, + undefined, + "top-level key should be removed" + ); + }); +}); + +after(() => { + try { + fs.rmSync(tmpDir, { recursive: true, force: true }); + } catch { + // ignore + } +}); diff --git a/tests/unit/server-owned-tool-loop-gate.test.ts b/tests/unit/server-owned-tool-loop-gate.test.ts new file mode 100644 index 0000000000..3501014c40 --- /dev/null +++ b/tests/unit/server-owned-tool-loop-gate.test.ts @@ -0,0 +1,82 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { shouldRunServerOwnedToolLoop } from "../../open-sse/handlers/chatCore/serverOwnedToolLoopGate.ts"; +import { FORMATS } from "../../open-sse/translator/formats.ts"; + +test("flag-off never runs the loop", () => { + assert.equal( + shouldRunServerOwnedToolLoop({ + enabled: false, + stream: false, + isResponsesEndpoint: false, + sourceFormat: FORMATS.OPENAI, + }), + false + ); +}); + +test("streaming never runs the loop", () => { + assert.equal( + shouldRunServerOwnedToolLoop({ + enabled: true, + stream: true, + isResponsesEndpoint: false, + sourceFormat: FORMATS.OPENAI, + }), + false + ); +}); + +test("Responses endpoint and format keep the old path", () => { + assert.equal( + shouldRunServerOwnedToolLoop({ + enabled: true, + stream: false, + isResponsesEndpoint: true, + sourceFormat: FORMATS.OPENAI, + }), + false + ); + assert.equal( + shouldRunServerOwnedToolLoop({ + enabled: true, + stream: false, + isResponsesEndpoint: false, + sourceFormat: FORMATS.OPENAI_RESPONSES, + }), + false + ); +}); + +test("non-streaming Chat and Claude run the loop when enabled", () => { + assert.equal( + shouldRunServerOwnedToolLoop({ + enabled: true, + stream: false, + isResponsesEndpoint: false, + sourceFormat: FORMATS.OPENAI, + }), + true + ); + assert.equal( + shouldRunServerOwnedToolLoop({ + enabled: true, + stream: false, + isResponsesEndpoint: false, + sourceFormat: FORMATS.CLAUDE, + }), + true + ); +}); + +test("gemini and other source formats keep the old path", () => { + assert.equal( + shouldRunServerOwnedToolLoop({ + enabled: true, + stream: false, + isResponsesEndpoint: false, + sourceFormat: FORMATS.GEMINI, + }), + false + ); +}); diff --git a/tests/unit/server-owned-tool-loop-wire.test.ts b/tests/unit/server-owned-tool-loop-wire.test.ts new file mode 100644 index 0000000000..c403dfc40a --- /dev/null +++ b/tests/unit/server-owned-tool-loop-wire.test.ts @@ -0,0 +1,246 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; + +import { + applyServerOwnedToolLoopIfNeeded, + derivePostInjectionRequestIdentity, +} from "../../open-sse/handlers/chatCore/serverOwnedToolLoopWire.ts"; +import type { + NonStreamingProviderLegResult, + ProviderLegReceipt, +} from "../../src/lib/skills/toolLoopTypes.ts"; +import { FORMATS } from "../../open-sse/translator/formats.ts"; + +function receipt(index: number): ProviderLegReceipt { + return { + index, + connectionId: "conn-1", + provider: "openai", + model: "gpt-4o", + startedAt: "2026-01-01T00:00:00.000Z", + endedAt: "2026-01-01T00:00:01.000Z", + latencyMs: 10, + httpStatus: 200, + errorType: null, + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + serviceTier: null, + computedCostUsd: 0.01, + toolCalls: [], + termination: "completed", + clientVisible: true, + }; +} + +function okLeg(overrides: Partial = {}) { + return { + kind: "ok" as const, + response: { + choices: [ + { + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_1", + type: "function", + function: { name: "memory_search", arguments: '{"q":"x"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }, + responseForMemoryExtraction: {}, + providerBody: {}, + providerRequest: {}, + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + responsePayloadFormat: "openai", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: receipt(0), + ...overrides, + }; +} + +test("flag-off skips the loop and does not resume", async () => { + let followUps = 0; + const result = await applyServerOwnedToolLoopIfNeeded({ + enabled: false, + stream: false, + isResponsesEndpoint: false, + sourceFormat: FORMATS.OPENAI, + initialLeg: okLeg(), + sourceBody: { model: "gpt-4o", messages: [{ role: "user", content: "hi" }] }, + skillsModelId: "openai", + executionContext: { + apiKeyId: "k", + sessionId: "s", + requestId: "r", + builtinToolNames: ["memory_search"], + }, + expectedConnectionId: "conn-1", + followUpLeg: async () => { + followUps += 1; + throw new Error("must not resume"); + }, + logReceipt: () => {}, + }); + assert.equal(result.kind, "skip"); + assert.equal(followUps, 0); +}); + +test("Responses format skips even when enabled", async () => { + const result = await applyServerOwnedToolLoopIfNeeded({ + enabled: true, + stream: false, + isResponsesEndpoint: true, + sourceFormat: FORMATS.OPENAI, + initialLeg: okLeg(), + sourceBody: { model: "gpt-4o", messages: [] }, + skillsModelId: "openai", + executionContext: { apiKeyId: "k", sessionId: "s", requestId: "r" }, + expectedConnectionId: "conn-1", + followUpLeg: async () => { + throw new Error("must not resume"); + }, + logReceipt: () => {}, + }); + assert.equal(result.kind, "skip"); +}); + +test("enabled Chat loop resumes once and logs receipts", async () => { + const logged: number[] = []; + let followUps = 0; + const result = await applyServerOwnedToolLoopIfNeeded({ + enabled: true, + stream: false, + isResponsesEndpoint: false, + sourceFormat: FORMATS.OPENAI, + initialLeg: okLeg(), + sourceBody: { + model: "gpt-4o", + messages: [{ role: "user", content: "hi" }], + }, + skillsModelId: "openai", + executionContext: { + apiKeyId: "k", + sessionId: "s", + requestId: "r", + builtinToolNames: ["memory_search"], + }, + expectedConnectionId: "conn-1", + followUpLeg: async () => { + followUps += 1; + return { + kind: "ok", + response: { + choices: [{ message: { role: "assistant", content: "done" }, finish_reason: "stop" }], + }, + responseForMemoryExtraction: { text: "done" }, + providerBody: { id: "2" }, + providerRequest: { messages: [] }, + usage: { prompt_tokens: 4, completion_tokens: 2, total_tokens: 6 }, + responsePayloadFormat: "openai", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: receipt(1), + }; + }, + logReceipt: (r) => logged.push(r.index), + executeServerOwned: async (calls) => + calls.map((c) => ({ id: c.id, name: c.name, result: { hits: [] }, replayed: false })), + }); + assert.equal(result.kind, "ok"); + if (result.kind !== "ok") return; + assert.equal(followUps, 1); + assert.deepEqual(logged, [0, 1]); + assert.equal( + (result.leg.response as { choices: Array<{ message: { content: string } }> }).choices[0].message + .content, + "done" + ); + assert.equal(result.usage?.prompt_tokens, 14); +}); + +test("provider error from follow-up is returned as error", async () => { + const result = await applyServerOwnedToolLoopIfNeeded({ + enabled: true, + stream: false, + isResponsesEndpoint: false, + sourceFormat: FORMATS.OPENAI, + initialLeg: okLeg(), + sourceBody: { model: "gpt-4o", messages: [{ role: "user", content: "hi" }] }, + skillsModelId: "openai", + executionContext: { + apiKeyId: "k", + sessionId: "s", + requestId: "r", + builtinToolNames: ["memory_search"], + }, + expectedConnectionId: "conn-1", + followUpLeg: async () => ({ + kind: "error", + result: { + success: false, + status: 429, + response: new Response(null, { status: 429 }), + error: "rate limited", + errorCode: "rate_limited", + }, + receipt: receipt(1), + usage: { prompt_tokens: 1, completion_tokens: 0, total_tokens: 1 }, + }), + logReceipt: () => {}, + executeServerOwned: async (calls) => + calls.map((c) => ({ id: c.id, name: c.name, result: { hits: [] }, replayed: false })), + }); + assert.equal(result.kind, "error"); + if (result.kind !== "error") return; + assert.equal(result.loop.errorResult?.status, 429); +}); + +test("derivePostInjectionRequestIdentity uses the client idempotency header", () => { + const a = derivePostInjectionRequestIdentity({ + apiKeyId: "k", + headers: { "idempotency-key": "tool-loop-retry-1" }, + skillRequestId: "internal-1", + postInjectionBody: { messages: [{ role: "user", content: "hi" }] }, + }); + const b = derivePostInjectionRequestIdentity({ + apiKeyId: "k", + headers: { "idempotency-key": "tool-loop-retry-1" }, + skillRequestId: "internal-2", + postInjectionBody: { messages: [{ role: "user", content: "hi" }] }, + }); + const c = derivePostInjectionRequestIdentity({ + apiKeyId: "k", + headers: { "idempotency-key": "other" }, + skillRequestId: "internal-1", + postInjectionBody: { messages: [{ role: "user", content: "hi" }] }, + }); + assert.equal(a, b); + assert.notEqual(a, c); +}); + +test("applyServerOwnedToolLoopIfNeeded accepts undefined expectedConnectionId for unmanaged leases", async () => { + const result = await applyServerOwnedToolLoopIfNeeded({ + enabled: false, + stream: false, + isResponsesEndpoint: false, + sourceFormat: FORMATS.OPENAI, + initialLeg: okLeg(), + sourceBody: { model: "gpt-4o", messages: [] }, + skillsModelId: "openai", + executionContext: { apiKeyId: "k", sessionId: "s", requestId: "r" }, + expectedConnectionId: undefined, + followUpLeg: async () => { + throw new Error("must not resume"); + }, + logReceipt: () => {}, + }); + assert.equal(result.kind, "skip"); +}); diff --git a/tests/unit/server-owned-tool-loop.test.ts b/tests/unit/server-owned-tool-loop.test.ts new file mode 100644 index 0000000000..6c5bc931d2 --- /dev/null +++ b/tests/unit/server-owned-tool-loop.test.ts @@ -0,0 +1,1104 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + runServerOwnedToolLoop, + MAX_FOLLOW_UPS, + LOOP_BUDGET_MS, + MIN_REMAINING_FOR_FOLLOW_UP_MS, +} from "../../src/lib/skills/serverOwnedToolLoop.ts"; +import type { + ServerOwnedToolLoopOptions, + NonStreamingProviderLegResult, + ProviderLegReceipt, + ChatCoreErrorResult, + ToolCall, +} from "../../src/lib/skills/toolLoopTypes.ts"; +import { ServerOwnedExecutionError, extractToolCalls } from "../../src/lib/skills/interception.ts"; +import { buildFollowUpSourceBody } from "../../src/lib/skills/followUpTranscript.ts"; + +// ─── Fix 6: serializedResultTextById verbatim use ──────────────────────────── + +test("buildFollowUpSourceBody uses serializedResultTextById verbatim when provided", () => { + const sentinel = '{"custom":"SENTINEL_12345"}'; + const toolCalls = [{ id: "tc1", name: "memory_search", arguments: { query: "x" } }]; + const results = [{ id: "tc1", name: "memory_search", result: { hits: ["a"] }, replayed: false }]; + const previousResponse = { + choices: [ + { + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: "tc1", + type: "function", + function: { name: "memory_search", arguments: '{"query":"x"}' }, + }, + ], + }, + }, + ], + }; + const sourceBody = { model: "gpt-4o", messages: [{ role: "user", content: "hi" }] }; + + // With serializedResultTextById → uses sentinel verbatim + const serMap = new Map([["tc1", sentinel]]); + const withMap = buildFollowUpSourceBody({ + sourceBody, + previousResponse, + toolCalls, + results, + sourceFormat: "openai", + maxResultBytes: 32_768, + maxTotalResultBytes: 65_536, + serializedResultTextById: serMap, + }); + + const toolMsg = (withMap.messages as Record[]).find( + (m: Record) => m.role === "tool" && m.tool_call_id === "tc1" + ); + assert.strictEqual(toolMsg!.content, sentinel, "must use pre-serialized text verbatim"); + + // Without serializedResultTextById → serializer would produce JSON of { hits: ["a"] } + const withoutMap = buildFollowUpSourceBody({ + sourceBody, + previousResponse, + toolCalls, + results, + sourceFormat: "openai", + maxResultBytes: 32_768, + maxTotalResultBytes: 65_536, + }); + + const toolMsgNoMap = (withoutMap.messages as Record[]).find( + (m: Record) => m.role === "tool" && m.tool_call_id === "tc1" + ); + const defaultSerialized = JSON.stringify({ hits: ["a"] }); + assert.strictEqual(toolMsgNoMap!.content, defaultSerialized, "without map, uses JSON.stringify"); + assert.notStrictEqual( + toolMsgNoMap!.content, + sentinel, + "without map, content differs from sentinel" + ); +}); + +// ─── Helpers ────────────────────────────────────────────────────────────────── + +type UnknownRecord = Record; + +function makeReceipt(overrides: Partial = {}): ProviderLegReceipt { + return { + index: 0, + connectionId: "conn-1", + provider: "openai", + model: "gpt-4o", + startedAt: new Date().toISOString(), + endedAt: new Date().toISOString(), + latencyMs: 100, + httpStatus: 200, + errorType: null, + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + serviceTier: null, + computedCostUsd: 0.001, + toolCalls: [], + termination: "completed", + clientVisible: true, + ...overrides, + }; +} + +function makeOkLeg( + overrides: Partial = {} +): NonStreamingProviderLegResult & { kind: "ok" } { + return { + kind: "ok", + response: { + id: "chatcmpl-abc", + choices: [ + { + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_1", + type: "function", + function: { name: "memory_search", arguments: '{"query":"foo"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }, + responseForMemoryExtraction: { + choices: [{ message: { role: "assistant", content: null } }], + }, + providerBody: {}, + providerRequest: {}, + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + responsePayloadFormat: "openai", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: makeReceipt(), + ...overrides, + }; +} + +function makeServerOwnedCallResponse( + _callId: string, + _name: string +): NonStreamingProviderLegResult & { kind: "ok" } { + return { + kind: "ok", + response: { + id: "chatcmpl-def", + choices: [ + { + message: { + role: "assistant", + content: "Here is what I found about foo.", + tool_calls: undefined, + }, + finish_reason: "stop", + }, + ], + }, + responseForMemoryExtraction: { + choices: [{ message: { role: "assistant", content: "Here is what I found about foo." } }], + }, + providerBody: {}, + providerRequest: {}, + usage: { prompt_tokens: 150, completion_tokens: 80, total_tokens: 230 }, + responsePayloadFormat: "openai", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: makeReceipt({ + index: 1, + usage: { prompt_tokens: 150, completion_tokens: 80, total_tokens: 230 }, + computedCostUsd: 0.002, + }), + }; +} + +function makeErrorResult( + status: number, + message: string, + code?: string, + errorType?: string +): ChatCoreErrorResult { + return { + success: false, + status, + response: new Response(JSON.stringify({ error: message }), { status }), + error: message, + errorCode: code, + errorType, + }; +} + +function makeDefaultOptions( + overrides: Partial = {} +): ServerOwnedToolLoopOptions { + const initialLeg = makeOkLeg(); + return { + initialLeg, + sourceBody: { + model: "gpt-4o", + messages: [ + { role: "system", content: "You are helpful." }, + { role: "user", content: "Search for foo" }, + ], + tools: [ + { + type: "function", + function: { + name: "memory_search", + description: "search memory", + parameters: { type: "object", properties: { query: { type: "string" } } }, + }, + }, + ], + }, + sourceFormat: "openai", + skillsModelId: "gpt-4o", + executionContext: { + apiKeyId: "key-1", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["memory_search"], + }, + executeServerOwned: async (calls: ToolCall[]) => { + return calls.map((call) => ({ + id: call.id, + name: call.name, + result: { hits: ["a", "b"] }, + replayed: false, + })); + }, + resumeUpstream: async () => makeServerOwnedCallResponse("call_1", "memory_search"), + deadlineAtMs: 120_000, + ...overrides, + }; +} + +// ─── Interface contract tests (no casts) ───────────────────────────────────── + +test("ServerOwnedToolLoopOptions accepts abortSignal and now fields at type level", () => { + // This test proves the interface has the fields; a cast would hide missing fields. + const ac = new AbortController(); + let customNow = 5000; + const opts: ServerOwnedToolLoopOptions = { + initialLeg: makeOkLeg(), + sourceBody: { model: "gpt-4o", messages: [] }, + sourceFormat: "openai", + skillsModelId: "gpt-4o", + executionContext: { apiKeyId: "k", sessionId: "s", requestId: "r" }, + executeServerOwned: async () => [], + resumeUpstream: async () => makeOkLeg(), + deadlineAtMs: 120_000, + abortSignal: ac.signal, + now: () => customNow, + }; + // Verify the fields are present and accessible without cast + assert.ok(opts.abortSignal, "abortSignal must be accessible"); + assert.strictEqual(typeof opts.now, "function", "now must be a function"); + assert.strictEqual(opts.now!(), 5000, "now() returns the injected value"); +}); + +test("NonStreamingProviderLegResult error arm carries usage field", () => { + const errLeg: NonStreamingProviderLegResult = { + kind: "error", + result: { + success: false, + status: 500, + response: new Response(), + error: "fail", + }, + receipt: makeReceipt({ + usage: { prompt_tokens: 50, completion_tokens: 10, total_tokens: 60 }, + }), + usage: { prompt_tokens: 50, completion_tokens: 10, total_tokens: 60 }, + }; + assert.ok(errLeg.kind === "error"); + assert.ok(errLeg.usage, "error leg must carry usage"); + assert.strictEqual(errLeg.usage!.prompt_tokens, 50); +}); + +// ─── Fix 3: sourceFormat-based extraction regression ───────────────────────── + +test("extractToolCalls with opaque model alias + Claude shape returns 1 tool call via sourceFormat", () => { + const claudeResponse = { + content: [{ type: "tool_use", id: "tu_1", name: "memory_search", input: { query: "test" } }], + stop_reason: "tool_use", + }; + + // Opaque model alias that detectProvider maps to "openai" → returns 0 + const viaModelId = extractToolCalls(claudeResponse, "official-fable"); + assert.strictEqual(viaModelId.length, 0, "opaque alias without sourceFormat returns 0"); + + // Explicit sourceFormat "claude" → returns 1 + const viaSourceFormat = extractToolCalls(claudeResponse, "claude"); + assert.strictEqual(viaSourceFormat.length, 1, "sourceFormat=claude returns 1 tool call"); + assert.strictEqual(viaSourceFormat[0].name, "memory_search"); +}); + +test("extractToolCalls with opaque model alias + Claude shape actually executes in loop", async () => { + const claudeResponse: NonStreamingProviderLegResult & { kind: "ok" } = { + kind: "ok", + response: { + content: [{ type: "tool_use", id: "tu_exec", name: "memory_search", input: { query: "x" } }], + stop_reason: "tool_use", + }, + responseForMemoryExtraction: {}, + providerBody: {}, + providerRequest: {}, + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + responsePayloadFormat: "claude", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: makeReceipt({ index: 0 }), + }; + + let executeCalled = false; + const opts: ServerOwnedToolLoopOptions = { + initialLeg: claudeResponse, + sourceBody: { model: "official-fable", messages: [{ role: "user", content: "hi" }] }, + sourceFormat: "claude", + skillsModelId: "official-fable", + executionContext: { + apiKeyId: "k", + sessionId: "s", + requestId: "r", + builtinToolNames: ["memory_search"], + }, + executeServerOwned: async (calls) => { + executeCalled = true; + return calls.map((c) => ({ id: c.id, name: c.name, result: { ok: true }, replayed: false })); + }, + resumeUpstream: async () => ({ + kind: "ok", + response: { + content: [{ type: "text", text: "Done" }], + stop_reason: "end_turn", + }, + responseForMemoryExtraction: {}, + providerBody: {}, + providerRequest: {}, + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + responsePayloadFormat: "claude", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: makeReceipt({ index: 1 }), + }), + deadlineAtMs: 120_000, + }; + + const result = await runServerOwnedToolLoop(opts); + assert.ok(executeCalled, "server-owned call must execute even with opaque model alias"); + assert.strictEqual(result.termination, "completed"); + assert.strictEqual(result.followUps, 1); +}); + +// ─── §5.5 Constants ─────────────────────────────────────────────────────────── + +test("MAX_FOLLOW_UPS is 3", () => { + assert.strictEqual(MAX_FOLLOW_UPS, 3); +}); + +test("LOOP_BUDGET_MS is 120000", () => { + assert.strictEqual(LOOP_BUDGET_MS, 120_000); +}); + +test("MIN_REMAINING_FOR_FOLLOW_UP_MS is 10000", () => { + assert.strictEqual(MIN_REMAINING_FOR_FOLLOW_UP_MS, 10_000); +}); + +// ─── Happy-path ────────────────────────────────────────────────────────────── + +test("happy-path: server-owned call executed, resume produces text, followUps=1, termination=completed", async () => { + const opts = makeDefaultOptions(); + const result = await runServerOwnedToolLoop(opts); + + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.followUps, 1); + assert.strictEqual(result.receipts.length, 2, "initial leg + 1 follow-up"); + assert.strictEqual(result.termination, "completed"); + assert.ok(result.response, "response must be non-empty"); + assert.ok(result.cumulativeUsage, "cumulativeUsage must be non-null"); + + const response = result.response as UnknownRecord; + const choices = response.choices as Array; + const message = choices[0].message as UnknownRecord; + assert.ok( + typeof message.content === "string" && message.content.length > 0, + "final content must be non-empty string" + ); +}); + +// ─── Step 4: Termination branches ───────────────────────────────────────────── + +test("no server-owned calls: completed with no follow-ups", async () => { + const initialLeg = makeOkLeg({ + response: { + id: "chatcmpl-ok", + choices: [ + { + message: { + role: "assistant", + content: "Plain text response", + tool_calls: undefined, + }, + finish_reason: "stop", + }, + ], + }, + }); + const opts = makeDefaultOptions({ + initialLeg, + executeServerOwned: async () => [], + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.followUps, 0); + assert.strictEqual(result.termination, "completed"); + assert.strictEqual(result.receipts.length, 1); +}); + +test("client_tools: no follow-up, client native calls only", async () => { + const initialLeg = makeOkLeg({ + response: { + id: "chatcmpl-ct", + choices: [ + { + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: "cli1", + type: "function", + function: { name: "Bash", arguments: '{"cmd":"ls"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }, + }); + const opts = makeDefaultOptions({ initialLeg }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.followUps, 0); + assert.strictEqual(result.termination, "client_tools"); +}); + +test("mixed_tools: server call executed, no follow-up, results appended to content", async () => { + const initialLeg = makeOkLeg({ + response: { + id: "chatcmpl-mix", + choices: [ + { + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: "srv1", + type: "function", + function: { name: "memory_search", arguments: '{"query":"foo"}' }, + }, + { + id: "cli1", + type: "function", + function: { name: "Bash", arguments: '{"cmd":"ls"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }, + }); + + let executeCount = 0; + const opts = makeDefaultOptions({ + initialLeg, + executeServerOwned: async (calls) => { + executeCount++; + return calls.map((c) => ({ + id: c.id, + name: c.name, + result: { hits: ["x"] }, + replayed: false, + })); + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.followUps, 0); + assert.strictEqual(result.termination, "mixed_tools"); + assert.strictEqual(executeCount, 1, "server call executed once"); +}); + +test("max_followups: 3 follow-ups, 4th leg server call still executed but no 5th leg", async () => { + let legCount = 0; + let executeCount = 0; + + const opts = makeDefaultOptions({ + executeServerOwned: async (calls) => { + executeCount++; + return calls.map((c) => ({ + id: c.id, + name: c.name, + result: { ok: true }, + replayed: false, + })); + }, + resumeUpstream: async () => { + legCount++; + return { + kind: "ok", + response: { + id: `chatcmpl-leg${legCount}`, + choices: [ + { + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: `call_${legCount + 1}`, + type: "function", + function: { name: "memory_search", arguments: '{"query":"q"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }, + responseForMemoryExtraction: {}, + providerBody: {}, + providerRequest: {}, + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + responsePayloadFormat: "openai", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: makeReceipt({ + index: legCount + 1, + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + }), + }; + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.followUps, 3); + assert.strictEqual(result.termination, "max_followups"); + assert.strictEqual(executeCount, 4, "server calls executed for initial + 3 follow-ups"); + assert.strictEqual(result.receipts.length, 4, "exactly 4 receipts: initial + 3 follow-ups"); +}); + +test("deadline: not enough remaining time → termination=deadline, no follow-up", async () => { + const startMs = 1000; + let now = startMs; + const opts = makeDefaultOptions({ + now: () => now, + deadlineAtMs: startMs + MIN_REMAINING_FOR_FOLLOW_UP_MS - 1, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.followUps, 0); + assert.strictEqual(result.termination, "deadline"); +}); + +test("client_abort: abortSignal.aborted → termination=client_abort, no formatter/resume", async () => { + const ac = new AbortController(); + ac.abort(); + + const opts = makeDefaultOptions({ abortSignal: ac.signal }); + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "client_abort"); + assert.ok(result.errorResult, "errorResult must be present"); + assert.strictEqual(result.errorResult!.status, 499); + assert.strictEqual(result.errorResult!.errorCode, "client_closed_request"); +}); + +test("provider_error: resumeUpstream returns error → termination=provider_error, no formatter", async () => { + const errorResult = makeErrorResult( + 500, + "Internal Server Error", + "internal_error", + "server_error" + ); + const opts = makeDefaultOptions({ + resumeUpstream: async () => ({ + kind: "error", + result: errorResult, + receipt: makeReceipt({ + index: 1, + httpStatus: 500, + errorType: "server_error", + usage: { prompt_tokens: 50, completion_tokens: 10, total_tokens: 60 }, + computedCostUsd: 0.003, + }), + usage: { prompt_tokens: 50, completion_tokens: 10, total_tokens: 60 }, + }), + }); + + const origExec = opts.executeServerOwned; + opts.executeServerOwned = async (calls, context) => { + const results = await origExec(calls, context); + return results; + }; + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "provider_error"); + assert.ok(result.errorResult, "errorResult must be present"); + assert.strictEqual(result.errorResult!.status, 500); + assert.strictEqual(result.receipts.length, 2, "initial + failed follow-up receipts"); + assert.ok(result.receipts[1].usage, "failed leg receipt must carry usage"); + assert.strictEqual(result.receipts[1].usage!.prompt_tokens, 50); + assert.strictEqual(result.totalCostUsd, 0.004, "cost includes failed leg computedCostUsd"); +}); + +test("connection_mismatch: follow-up connectionId differs → termination=connection_mismatch", async () => { + const opts = makeDefaultOptions({ + resumeUpstream: async () => ({ + kind: "ok", + response: { + id: "chatcmpl-mismatch", + choices: [ + { + message: { + role: "assistant", + content: "Response", + tool_calls: undefined, + }, + finish_reason: "stop", + }, + ], + }, + responseForMemoryExtraction: {}, + providerBody: {}, + providerRequest: {}, + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + responsePayloadFormat: "openai", + looksLikeSSE: false, + connectionId: "conn-CHANGED", + headers: new Headers(), + receipt: makeReceipt({ + index: 1, + connectionId: "conn-CHANGED", + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + }), + }), + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "connection_mismatch"); + assert.ok(result.errorResult, "errorResult must be present"); + assert.strictEqual(result.errorResult!.status, 409); + assert.strictEqual(result.errorResult!.errorCode, "LEASE_CONNECTION_MISMATCH"); +}); + +test("ServerOwnedExecutionError: execution_error → termination=execution_error, no formatter", async () => { + const opts = makeDefaultOptions({ + executeServerOwned: async () => { + throw new ServerOwnedExecutionError("handler crashed", "TOOL_EXECUTION_ERROR", 500); + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "execution_error"); + assert.ok(result.errorResult, "errorResult must be present"); + assert.strictEqual(result.errorResult!.status, 500); + assert.strictEqual(result.errorResult!.errorCode, "TOOL_EXECUTION_ERROR"); +}); + +test("ServerOwnedExecutionError: execution_in_progress → termination=execution_in_progress", async () => { + const opts = makeDefaultOptions({ + executeServerOwned: async () => { + throw new ServerOwnedExecutionError("Tool execution in progress", "TOOL_IN_PROGRESS", 409); + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "execution_in_progress"); + assert.ok(result.errorResult); + assert.strictEqual(result.errorResult!.status, 409); + assert.strictEqual(result.errorResult!.errorCode, "TOOL_IN_PROGRESS"); +}); + +test("ServerOwnedExecutionError: execution_unknown → termination=execution_unknown", async () => { + const opts = makeDefaultOptions({ + executeServerOwned: async () => { + throw new ServerOwnedExecutionError( + "Tool execution state unknown", + "TOOL_STATE_UNKNOWN", + 500 + ); + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "execution_unknown"); + assert.ok(result.errorResult); + assert.strictEqual(result.errorResult!.status, 500); + assert.strictEqual(result.errorResult!.errorCode, "TOOL_STATE_UNKNOWN"); +}); + +test("ServerOwnedExecutionError: execution_identity_conflict → termination=execution_identity_conflict", async () => { + const opts = makeDefaultOptions({ + executeServerOwned: async () => { + throw new ServerOwnedExecutionError( + "Tool execution identity conflict", + "IDENTITY_CONFLICT", + 409 + ); + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "execution_identity_conflict"); + assert.ok(result.errorResult); + assert.strictEqual(result.errorResult!.status, 409); + assert.strictEqual(result.errorResult!.errorCode, "IDENTITY_CONFLICT"); +}); + +test("ServerOwnedExecutionError: execution_timeout → termination=execution_timeout", async () => { + const opts = makeDefaultOptions({ + executeServerOwned: async () => { + throw new ServerOwnedExecutionError( + "Tool execution timed out", + "TOOL_EXECUTION_TIMEOUT", + 504 + ); + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "execution_timeout"); + assert.ok(result.errorResult); + assert.strictEqual(result.errorResult!.status, 504); + assert.strictEqual(result.errorResult!.errorCode, "TOOL_EXECUTION_TIMEOUT"); +}); + +// ─── Provider error identity ───────────────────────────────────────────────── + +test("provider_error: error result is same object identity as resumeUpstream return, formatter count=0", async () => { + const errorResult = makeErrorResult(502, "Bad Gateway"); + let resumeCallCount = 0; + + const opts = makeDefaultOptions({ + resumeUpstream: async () => { + resumeCallCount++; + return { + kind: "error", + result: errorResult, + receipt: makeReceipt({ index: 1, httpStatus: 502 }), + usage: null, + }; + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "provider_error"); + assert.strictEqual(result.errorResult, errorResult, "errorResult must be same object identity"); + assert.strictEqual(resumeCallCount, 1); +}); + +// ─── Cumulative UTF-8 budget ───────────────────────────────────────────────── + +test("tool_output_budget: cumulative bytes exhausted → termination=tool_output_budget", async () => { + const bigResult = { data: "x".repeat(70_000) }; + const opts = makeDefaultOptions({ + maxTotalResultBytes: 100, + executeServerOwned: async (calls) => + calls.map((c) => ({ + id: c.id, + name: c.name, + result: bigResult, + replayed: false, + })), + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.termination, "tool_output_budget"); +}); + +test("tool_output_budget: truncated=true on any result must terminate, no resume", async () => { + const bigResult = { data: "y".repeat(50_000) }; + let resumeCount = 0; + const opts = makeDefaultOptions({ + maxResultBytes: 100, + maxTotalResultBytes: 200, + executeServerOwned: async (calls) => + calls.map((c) => ({ + id: c.id, + name: c.name, + result: bigResult, + replayed: false, + })), + resumeUpstream: async () => { + resumeCount++; + return makeServerOwnedCallResponse("call_next", "memory_search"); + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.termination, "tool_output_budget", "truncated result must terminate"); + assert.strictEqual(resumeCount, 0, "resume must not be called when result is truncated"); + // Formatter output must exist and be byte-bounded + assert.ok(result.response, "response must be present"); + const responseStr = JSON.stringify(result.response); + const responseBytes = Buffer.byteLength(responseStr, "utf8"); + assert.ok(responseBytes > 0, "formatter output must be non-empty"); +}); + +test("mixed_tools: abort before execute must be checked", async () => { + const ac = new AbortController(); + ac.abort(); + + const initialLeg = makeOkLeg({ + response: { + id: "chatcmpl-mix-abort", + choices: [ + { + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: "srv1", + type: "function", + function: { name: "memory_search", arguments: '{"query":"x"}' }, + }, + { + id: "cli1", + type: "function", + function: { name: "Bash", arguments: '{"cmd":"ls"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }, + }); + + let executeCount = 0; + const opts = makeDefaultOptions({ + initialLeg, + abortSignal: ac.signal, + executeServerOwned: async (calls) => { + executeCount++; + return calls.map((c) => ({ id: c.id, name: c.name, result: { ok: true }, replayed: false })); + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual( + result.termination, + "client_abort", + "mixed with abort must terminate as client_abort" + ); + assert.strictEqual(executeCount, 0, "execute must not run when abort is signaled"); +}); + +// ─── All-null usage ────────────────────────────────────────────────────────── + +test("all-null usage legs → cumulativeUsage is null", async () => { + const opts = makeDefaultOptions({ + initialLeg: makeOkLeg({ usage: null }), + resumeUpstream: async () => ({ + kind: "ok", + response: { + id: "chatcmpl-null", + choices: [ + { + message: { role: "assistant", content: "Done", tool_calls: undefined }, + finish_reason: "stop", + }, + ], + }, + responseForMemoryExtraction: {}, + providerBody: {}, + providerRequest: {}, + usage: null, + responsePayloadFormat: "openai", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: makeReceipt({ index: 1, usage: null }), + }), + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.termination, "completed"); + assert.strictEqual(result.cumulativeUsage, null, "all-null usage → null cumulative"); +}); + +test("provider error with all-null usage → cumulativeUsage is null", async () => { + const opts = makeDefaultOptions({ + initialLeg: makeOkLeg({ usage: null }), + resumeUpstream: async () => ({ + kind: "error", + result: makeErrorResult(500, "Internal Server Error"), + receipt: makeReceipt({ index: 1, usage: null, httpStatus: 500 }), + usage: null, + }), + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "provider_error"); + assert.strictEqual(result.cumulativeUsage, null, "all-null usage on error → null cumulative"); +}); + +test("input objects are not mutated", async () => { + const sourceBody = { + model: "gpt-4o", + messages: [{ role: "user", content: "hi" }], + tools: [{ type: "function", function: { name: "memory_search" } }], + }; + const sourceBodySnapshot = JSON.parse(JSON.stringify(sourceBody)); + const opts = makeDefaultOptions({ sourceBody }); + + await runServerOwnedToolLoop(opts); + assert.deepStrictEqual(sourceBody, sourceBodySnapshot); +}); + +// ─── Accumulated transcript across rounds ──────────────────────────────────── + +test("two-round server calls: second resumeUpstream receives sourceBody with first round messages", async () => { + const sourceBodiesReceived: Record[] = []; + let resumeCount = 0; + + const opts = makeDefaultOptions({ + resumeUpstream: async (nextSourceBody) => { + sourceBodiesReceived.push(JSON.parse(JSON.stringify(nextSourceBody))); + resumeCount++; + if (resumeCount === 1) { + return { + kind: "ok", + response: { + id: "chatcmpl-round2", + choices: [ + { + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call_2", + type: "function", + function: { name: "memory_search", arguments: '{"query":"bar"}' }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }, + responseForMemoryExtraction: {}, + providerBody: {}, + providerRequest: {}, + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + responsePayloadFormat: "openai", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: makeReceipt({ + index: 1, + usage: { prompt_tokens: 100, completion_tokens: 20, total_tokens: 120 }, + }), + }; + } + return { + kind: "ok", + response: { + id: "chatcmpl-final", + choices: [ + { + message: { + role: "assistant", + content: "Final answer", + tool_calls: undefined, + }, + finish_reason: "stop", + }, + ], + }, + responseForMemoryExtraction: {}, + providerBody: {}, + providerRequest: {}, + usage: { prompt_tokens: 150, completion_tokens: 30, total_tokens: 180 }, + responsePayloadFormat: "openai", + looksLikeSSE: false, + connectionId: "conn-1", + headers: new Headers(), + receipt: makeReceipt({ + index: 2, + usage: { prompt_tokens: 150, completion_tokens: 30, total_tokens: 180 }, + }), + }; + }, + executeServerOwned: async (calls) => + calls.map((c) => ({ + id: c.id, + name: c.name, + result: { ok: true }, + replayed: false, + })), + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "ok"); + assert.strictEqual(result.followUps, 2); + assert.strictEqual(result.termination, "completed"); + + assert.strictEqual(sourceBodiesReceived.length, 2, "resumeUpstream called twice"); + const firstBody = sourceBodiesReceived[0].messages as UnknownRecord[]; + // 2 original + assistant turn + 1 tool message = 4 + assert.strictEqual(firstBody.length, 4, "2 original messages + assistant turn + tool result"); + const secondBody = sourceBodiesReceived[1].messages as UnknownRecord[]; + // 4 (from first resume) + assistant turn + 1 tool message = 6 + assert.strictEqual(secondBody.length, 6, "4 from first resume + round2 assistant/tool turn"); +}); + +// ─── Abort before resume ───────────────────────────────────────────────────── + +test("abort after execute but before resume → client_abort", async () => { + const ac = new AbortController(); + const opts = makeDefaultOptions({ + abortSignal: ac.signal, + executeServerOwned: async (calls) => { + ac.abort(); + return calls.map((c) => ({ + id: c.id, + name: c.name, + result: { ok: true }, + replayed: false, + })); + }, + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "client_abort"); + assert.ok(result.errorResult); + assert.strictEqual(result.errorResult!.status, 499); +}); + +// ─── Resumed leg receipt always enters receipts/usage/cost ─────────────────── + +test("failed leg receipt still enters receipts array with usage and cost", async () => { + const errorResult = makeErrorResult(429, "Rate limited"); + const opts = makeDefaultOptions({ + resumeUpstream: async () => ({ + kind: "error", + result: errorResult, + receipt: makeReceipt({ + index: 1, + httpStatus: 429, + usage: { prompt_tokens: 50, completion_tokens: 0, total_tokens: 50 }, + computedCostUsd: 0.0005, + }), + usage: { prompt_tokens: 50, completion_tokens: 0, total_tokens: 50 }, + }), + }); + + const result = await runServerOwnedToolLoop(opts); + assert.strictEqual(result.kind, "error"); + assert.strictEqual(result.termination, "provider_error"); + assert.strictEqual(result.receipts.length, 2, "initial + failed leg"); + assert.ok(result.receipts[1].usage, "failed leg receipt must have usage"); + assert.strictEqual(result.receipts[1].httpStatus, 429); + assert.strictEqual(result.totalCostUsd, 0.0015, "cost includes failed leg"); +}); diff --git a/tests/unit/skill-execution-fence.test.ts b/tests/unit/skill-execution-fence.test.ts new file mode 100644 index 0000000000..9dd07d501c --- /dev/null +++ b/tests/unit/skill-execution-fence.test.ts @@ -0,0 +1,786 @@ +process.env.DISABLE_SQLITE_AUTO_BACKUP = "true"; + +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import path from "node:path"; +import os from "node:os"; +import { createRequire } from "node:module"; + +const require_ = createRequire(import.meta.url); +const BetterSqlite3 = require_("better-sqlite3") as typeof import("better-sqlite3"); + +import { createBetterSqliteAdapter } from "../../src/lib/db/adapters/betterSqliteAdapter"; +import type { SqliteAdapter } from "../../src/lib/db/adapters/types"; + +import { + claimServerToolExecution, + finalizeServerToolExecution, +} from "../../src/lib/db/skillExecutionFence"; +import { runWithServerToolFence } from "../../src/lib/skills/toolExecutionFence"; + +// Minimal fixture schema — only what tests need; no SCHEMA_SQL import from core.ts +const FIXTURE_SCHEMA = ` + CREATE TABLE IF NOT EXISTS server_tool_executions ( + id TEXT PRIMARY KEY, + api_key_id TEXT NOT NULL, + request_identity TEXT NOT NULL, + tool_call_id TEXT NOT NULL, + tool_name TEXT NOT NULL, + input_digest TEXT NOT NULL, + output TEXT, + status TEXT NOT NULL CHECK(status IN ('running', 'success', 'error', 'timeout')), + error_message TEXT, + duration_ms INTEGER, + claim_expires_at TEXT NOT NULL, + completed_at TEXT, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + UNIQUE(api_key_id, request_identity, tool_call_id) + ); + CREATE INDEX IF NOT EXISTS idx_server_tool_executions_status_expiry + ON server_tool_executions(status, claim_expires_at); + CREATE INDEX IF NOT EXISTS idx_server_tool_executions_created + ON server_tool_executions(created_at); +`; + +function makeTempDb(): { + adapter: SqliteAdapter; + dir: string; + raw: import("better-sqlite3").Database; +} { + const dir = fs.mkdtempSync(path.join(os.tmpdir(), "fence-test-")); + const dbPath = path.join(dir, "test.db"); + const raw = new BetterSqlite3(dbPath); + raw.pragma("journal_mode = WAL"); + raw.pragma("busy_timeout = 2000"); + raw.exec(FIXTURE_SCHEMA); + const adapter = createBetterSqliteAdapter(raw); + return { adapter, dir, raw }; +} + +function makeSecondAdapter(dir: string): { + adapter: SqliteAdapter; + raw: import("better-sqlite3").Database; +} { + const dbPath = path.join(dir, "test.db"); + const raw = new BetterSqlite3(dbPath); + raw.pragma("journal_mode = WAL"); + raw.pragma("busy_timeout = 2000"); + const adapter = createBetterSqliteAdapter(raw); + return { adapter, raw }; +} + +function cleanup(raw: import("better-sqlite3").Database, dir: string) { + raw.close(); + fs.rmSync(dir, { recursive: true, force: true }); +} + +const BASE_INPUT = { + apiKeyId: "key-1", + requestIdentity: "key-1:req-id:body-hash", + toolCallId: "call-1", + toolName: "memory_search", + inputDigest: "abc123", + leaseExpiresAt: new Date(Date.now() + 60_000).toISOString(), +}; + +const BASE_FINALIZE = { + executionId: "", + status: "success" as const, + output: { result: "ok" }, + errorMessage: null, + durationMs: 100, +}; + +// ── Defect 1: Lease expiry compares STORED claim_expires_at, not input ── + +test("claim: stored expired lease → unknown even if retry supplies future lease", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const pastLease = new Date(Date.now() - 60_000).toISOString(); + const claim1 = claimServerToolExecution( + { ...BASE_INPUT, leaseExpiresAt: pastLease }, + adapter, + Date.now() + ); + assert.equal(claim1.kind, "claimed"); + const futureLease = new Date(Date.now() + 300_000).toISOString(); + const claim2 = claimServerToolExecution( + { ...BASE_INPUT, leaseExpiresAt: futureLease }, + adapter, + Date.now() + ); + assert.equal(claim2.kind, "unknown", "must compare stored lease, not input lease"); + } finally { + cleanup(raw, dir); + } +}); + +test("claim: stored future lease with expired retry → in_progress (not unknown)", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const futureLease = new Date(Date.now() + 300_000).toISOString(); + const claim1 = claimServerToolExecution( + { ...BASE_INPUT, leaseExpiresAt: futureLease }, + adapter, + Date.now() + ); + assert.equal(claim1.kind, "claimed"); + const pastLease = new Date(Date.now() - 10_000).toISOString(); + const claim2 = claimServerToolExecution( + { ...BASE_INPUT, leaseExpiresAt: pastLease }, + adapter, + Date.now() + ); + assert.equal(claim2.kind, "in_progress", "stored lease is future, so should be in_progress"); + } finally { + cleanup(raw, dir); + } +}); + +// ── Basic claim/replay tests ── + +test("claim: first claim returns claimed", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const result = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(result.kind, "claimed"); + assert.ok(typeof result.executionId === "string" && result.executionId.length > 0); + } finally { + cleanup(raw, dir); + } +}); + +test("claim: terminal row (success) with same identity returns replay", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const claim1 = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(claim1.kind, "claimed"); + finalizeServerToolExecution({ ...BASE_FINALIZE, executionId: claim1.executionId }, adapter); + const claim2 = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(claim2.kind, "replay"); + if (claim2.kind === "replay") { + assert.equal(claim2.status, "success"); + assert.deepEqual(claim2.output, { result: "ok" }); + } + } finally { + cleanup(raw, dir); + } +}); + +test("claim: terminal row (error) replay preserves error status and message", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const claim1 = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(claim1.kind, "claimed"); + finalizeServerToolExecution( + { + ...BASE_FINALIZE, + executionId: claim1.executionId, + status: "error", + output: null, + errorMessage: "tool failed", + durationMs: 42, + }, + adapter + ); + const claim2 = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(claim2.kind, "replay"); + if (claim2.kind === "replay") { + assert.equal(claim2.status, "error"); + assert.equal(claim2.errorMessage, "tool failed"); + } + } finally { + cleanup(raw, dir); + } +}); + +test("claim: same key but different name returns identity_conflict", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + claimServerToolExecution(BASE_INPUT, adapter); + const claim2 = claimServerToolExecution({ ...BASE_INPUT, toolName: "different_tool" }, adapter); + assert.equal(claim2.kind, "identity_conflict"); + } finally { + cleanup(raw, dir); + } +}); + +test("claim: same key but different inputDigest returns identity_conflict", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + claimServerToolExecution(BASE_INPUT, adapter); + const claim2 = claimServerToolExecution( + { ...BASE_INPUT, inputDigest: "different_digest" }, + adapter + ); + assert.equal(claim2.kind, "identity_conflict"); + } finally { + cleanup(raw, dir); + } +}); + +// ── Defect 2: Cross-handle poll test with gate-based concurrency ── + +test("fence: cross-handle poll — B sees in_progress, then A finalizes, then B replays", async () => { + const { adapter: adapterA, dir, raw: rawA } = makeTempDb(); + const { adapter: adapterB, raw: rawB } = makeSecondAdapter(dir); + let fakeTime = 1_000_000; + const now = () => fakeTime; + const sleep = async (ms: number) => { + fakeTime += ms; + await new Promise((r) => setTimeout(r, 0)); + }; + + // Gate: A signals when it starts, B signals when it observes in_progress + let aStartedResolve!: () => void; + const aStarted = new Promise((r) => { + aStartedResolve = r; + }); + let bObservedResolve!: () => void; + const bObserved = new Promise((r) => { + bObservedResolve = r; + }); + + let handlerCallCount = 0; + const handler = async () => { + handlerCallCount++; + aStartedResolve(); + // Wait until B has observed in_progress before finalizing + await bObserved; + return "handler-result"; + }; + + const sharedArgs = { q: "test-query" }; + + // Handle A claims via fence wrapper + const pA = runWithServerToolFence({ + apiKeyId: "key-1", + requestIdentity: "key-1:req-id:body-hash", + toolCallId: "call-1", + toolName: "memory_search", + arguments: sharedArgs, + leaseDurationMs: 60_000, + execute: handler, + now, + sleep, + db: adapterA, + }); + + // Wait for A to start executing + await aStarted; + + // Handle B calls the fence wrapper — should find running row, observe in_progress + const pB = runWithServerToolFence({ + apiKeyId: "key-1", + requestIdentity: "key-1:req-id:body-hash", + toolCallId: "call-1", + toolName: "memory_search", + arguments: sharedArgs, + leaseDurationMs: 60_000, + execute: async () => "should-not-run", + now, + sleep, + db: adapterB, + }); + + // Injected sleep proves B entered polling before A may finalize. + await new Promise((resolve) => setImmediate(resolve)); + assert.ok(fakeTime > 1_000_000, "B must poll before A finalizes"); + bObservedResolve(); + + const [, fenceResult] = await Promise.all([pA, pB]); + assert.equal(handlerCallCount, 1, "handler executes exactly once"); + assert.equal(fenceResult.kind, "replayed"); + if (fenceResult.kind === "replayed") { + assert.equal(fenceResult.status, "success"); + } + rawA.close(); + rawB.close(); + fs.rmSync(dir, { recursive: true, force: true }); +}); + +test("fence: in_progress poll returns in_progress when not finalized within deadline", async () => { + const { adapter, dir, raw } = makeTempDb(); + let fakeTime = 1_000_000; + const now = () => fakeTime; + const sleep = async (ms: number) => { + fakeTime += ms; + await new Promise((r) => setTimeout(r, 0)); + }; + + const sharedArgs = { q: "test" }; + const digest = require_("crypto") + .createHash("sha256") + .update(JSON.stringify(sharedArgs)) + .digest("hex"); + + // Insert a running row directly (simulates another process that claimed but hasn't finalized) + raw.exec(` + INSERT INTO server_tool_executions + (id, api_key_id, request_identity, tool_call_id, tool_name, input_digest, status, claim_expires_at) + VALUES ('ext-id', 'key-1', 'key-1:req:body', 'call-1', 'tool_a', '${digest}', 'running', datetime('now', '+300 seconds')) + `); + + const result = await runWithServerToolFence({ + apiKeyId: "key-1", + requestIdentity: "key-1:req:body", + toolCallId: "call-1", + toolName: "tool_a", + arguments: sharedArgs, + leaseDurationMs: 60_000, + execute: async () => "should-not-run", + now, + sleep, + db: adapter, + }); + + assert.equal(result.kind, "in_progress"); + assert.ok(fakeTime >= 1_000_000 + 2_000, "time should have advanced by poll budget"); + cleanup(raw, dir); +}); + +// ── Defect 3: Duration test — drives runWithServerToolFence, advances injected clock ── + +test("finalize: duration_ms persisted exactly via fence wrapper (success)", async () => { + const { adapter, dir, raw } = makeTempDb(); + let fakeTime = 10_000; + const now = () => fakeTime; + const sleep = async (ms: number) => { + fakeTime += ms; + }; + + const result = await runWithServerToolFence({ + apiKeyId: "key-1", + requestIdentity: "key-1:req:body", + toolCallId: "call-1", + toolName: "tool_a", + arguments: { x: 1 }, + leaseDurationMs: 60_000, + execute: async (_execId) => { + fakeTime += 42; + return "done"; + }, + now, + sleep, + db: adapter, + }); + + assert.equal(result.kind, "executed"); + // Query duration_ms from DB — wrapper must have calculated it, not hardcoded 0 + const row = raw + .prepare("SELECT duration_ms, status FROM server_tool_executions WHERE status = 'success'") + .get() as { duration_ms: number | null; status: string }; + assert.equal(row.duration_ms, 42, "duration_ms must be exactly 42 via injected clock"); + assert.equal(row.status, "success"); + cleanup(raw, dir); +}); + +test("finalize: duration_ms persisted exactly via fence wrapper (error)", async () => { + const { adapter, dir, raw } = makeTempDb(); + let fakeTime = 10_000; + const now = () => fakeTime; + const sleep = async (ms: number) => { + fakeTime += ms; + }; + + await assert.rejects( + runWithServerToolFence({ + apiKeyId: "key-1", + requestIdentity: "key-1:req:body", + toolCallId: "call-1", + toolName: "tool_a", + arguments: { x: 1 }, + leaseDurationMs: 60_000, + execute: async () => { + fakeTime += 17; + throw new Error("boom"); + }, + now, + sleep, + db: adapter, + }), + /boom/, + "handler error must propagate" + ); + + const row = raw + .prepare("SELECT duration_ms, status FROM server_tool_executions WHERE status = 'error'") + .get() as { duration_ms: number | null; status: string }; + assert.equal(row.duration_ms, 17, "duration_ms must be exactly 17 for error via injected clock"); + assert.equal(row.status, "error"); + cleanup(raw, dir); +}); + +// ── Defect 4: Rejected in-process active Promise → joiner gets replayed ── + +test("fence: rejected in-process promise — joiner gets replayed with error status", async () => { + const { adapter, dir, raw } = makeTempDb(); + let fakeTime = 1_000_000; + const now = () => fakeTime; + const sleep = async (ms: number) => { + fakeTime += ms; + await new Promise((r) => setTimeout(r, 0)); + }; + + // Gate for concurrency: A signals started, B signals it observed in_progress + let aStartedResolve!: () => void; + const aStarted = new Promise((r) => { + aStartedResolve = r; + }); + let bObservedResolve!: () => void; + const bObserved = new Promise((r) => { + bObservedResolve = r; + }); + + let handlerCalled = false; + const failingHandler = async () => { + handlerCalled = true; + aStartedResolve(); + await bObserved; + throw new Error("handler crashed"); + }; + + // Handle A claims via fence — handler will throw, wrapper resolves with error status + const pA = runWithServerToolFence({ + apiKeyId: "key-2", + requestIdentity: "key-2:req:body", + toolCallId: "call-x", + toolName: "tool_a", + arguments: {}, + leaseDurationMs: 60_000, + execute: failingHandler, + now, + sleep, + db: adapter, + }); + + await aStarted; + + // Handle B immediately tries same key — gets in_progress, polls the active promise + const pB = runWithServerToolFence({ + apiKeyId: "key-2", + requestIdentity: "key-2:req:body", + toolCallId: "call-x", + toolName: "tool_a", + arguments: {}, + leaseDurationMs: 60_000, + execute: async () => { + throw new Error("should-not-run"); + }, + now, + sleep, + db: adapter, + }); + + await new Promise((r) => setTimeout(r, 5)); + bObservedResolve(); + + const [rA, rB] = await Promise.allSettled([pA, pB]); + assert.ok(handlerCalled, "handler should have been called"); + assert.equal(rA.status, "rejected", "handle A should reject because handler threw"); + assert.equal(rB.status, "fulfilled"); + if (rB.status === "fulfilled") { + assert.equal(rB.value.kind, "replayed", "joiner must get replayed for terminal error"); + if (rB.value.kind === "replayed") { + assert.equal(rB.value.status, "error", "replayed status must be error"); + assert.equal(rB.value.errorMessage, "handler crashed"); + } + } + cleanup(raw, dir); +}); + +// ── Defect 6: isUniqueConstraintError rejects non-UNIQUE constraints ── + +test("isUniqueConstraintError: NOT NULL error propagates through claimServerToolExecution", () => { + const { dir, raw } = makeTempDb(); + try { + // Build a fake adapter that wraps the real one but makes INSERT throw NOT NULL + const realDir = fs.mkdtempSync(path.join(os.tmpdir(), "fence-fake-")); + const fakeRaw = new BetterSqlite3(path.join(realDir, "test.db")); + fakeRaw.pragma("journal_mode = WAL"); + fakeRaw.exec(FIXTURE_SCHEMA); + const fake = createBetterSqliteAdapter(fakeRaw); + + // Proxy: intercept INSERT INTO server_tool_executions and inject NOT NULL error + const proxyAdapter = new Proxy(fake, { + get(target, prop) { + if (prop === "prepare") { + return (sql: string) => { + const stmt = (target.prepare as Function)(sql); + if (/INSERT INTO server_tool_executions/i.test(sql)) { + return { + ...stmt, + run: (..._args: unknown[]) => { + const err = new Error("NOT NULL constraint failed") as Error & { code: string }; + err.code = "SQLITE_CONSTRAINT_NOTNULL"; + throw err; + }, + }; + } + return stmt; + }; + } + return (target as Record)[prop as string]; + }, + }) as SqliteAdapter; + + assert.throws( + () => { + claimServerToolExecution(BASE_INPUT, proxyAdapter); + }, + (err: unknown) => { + return ( + err instanceof Error && + err.message.includes("NOT NULL") && + (err as { code?: string }).code === "SQLITE_CONSTRAINT_NOTNULL" + ); + }, + "NOT NULL constraint must propagate, not be swallowed by isUniqueConstraintError" + ); + fakeRaw.close(); + fs.rmSync(realDir, { recursive: true, force: true }); + } finally { + cleanup(raw, dir); + } +}); + +test("isUniqueConstraintError: CHECK error propagates through claimServerToolExecution", () => { + const { dir, raw } = makeTempDb(); + try { + const realDir = fs.mkdtempSync(path.join(os.tmpdir(), "fence-fake-")); + const fakeRaw = new BetterSqlite3(path.join(realDir, "test.db")); + fakeRaw.pragma("journal_mode = WAL"); + fakeRaw.exec(FIXTURE_SCHEMA); + const fake = createBetterSqliteAdapter(fakeRaw); + + const proxyAdapter = new Proxy(fake, { + get(target, prop) { + if (prop === "prepare") { + return (sql: string) => { + const stmt = (target.prepare as Function)(sql); + if (/INSERT INTO server_tool_executions/i.test(sql)) { + return { + ...stmt, + run: (..._args: unknown[]) => { + const err = new Error("CHECK constraint failed") as Error & { code: string }; + err.code = "SQLITE_CONSTRAINT_CHECK"; + throw err; + }, + }; + } + return stmt; + }; + } + return (target as Record)[prop as string]; + }, + }) as SqliteAdapter; + + assert.throws( + () => { + claimServerToolExecution(BASE_INPUT, proxyAdapter); + }, + (err: unknown) => { + return ( + err instanceof Error && + err.message.includes("CHECK") && + (err as { code?: string }).code === "SQLITE_CONSTRAINT_CHECK" + ); + }, + "CHECK constraint must propagate, not be swallowed by isUniqueConstraintError" + ); + fakeRaw.close(); + fs.rmSync(realDir, { recursive: true, force: true }); + } finally { + cleanup(raw, dir); + } +}); + +// ── Existing tests (cleaned up, no done callbacks) ── + +test("finalize: only updates running rows, second finalize returns false", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const claim = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(claim.kind, "claimed"); + const firstFinalize = finalizeServerToolExecution( + { ...BASE_FINALIZE, executionId: claim.executionId }, + adapter + ); + assert.equal(firstFinalize, true); + const secondFinalize = finalizeServerToolExecution( + { ...BASE_FINALIZE, executionId: claim.executionId }, + adapter + ); + assert.equal(secondFinalize, false); + } finally { + cleanup(raw, dir); + } +}); + +test("finalize: output sanitized — nested credentials stripped, valid JSON preserved", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const claim = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(claim.kind, "claimed"); + const sensitiveOutput = { + token: "sk-live-secret123", + nested: { bearer: "Bearer abcdefghijklmnop", deep: { key: "sk-test-abcdefghijklmnop" } }, + data: "normal text", + }; + finalizeServerToolExecution( + { ...BASE_FINALIZE, executionId: claim.executionId, output: sensitiveOutput }, + adapter + ); + const row = raw + .prepare("SELECT output FROM server_tool_executions WHERE id = ?") + .get(claim.executionId) as { output: string | null }; + assert.ok(row.output, "output should be stored"); + assert.ok( + !row.output.includes("sk-live-secret123"), + "output must not contain raw sk- credential" + ); + assert.ok( + !row.output.includes("Bearer abcdefghijklmnop"), + "output must not contain raw Bearer token" + ); + const parsed = JSON.parse(row.output); + assert.ok(typeof parsed === "object", "sanitized output must be parseable JSON"); + } finally { + cleanup(raw, dir); + } +}); + +test("finalize: error message sanitized — no raw stack in error_message", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const claim = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(claim.kind, "claimed"); + finalizeServerToolExecution( + { + executionId: claim.executionId, + status: "error", + output: null, + errorMessage: "Something failed\n at /internal/path.js:42\n at processTicksAndRejections", + durationMs: 50, + }, + adapter + ); + const row = raw + .prepare("SELECT error_message FROM server_tool_executions WHERE id = ?") + .get(claim.executionId) as { error_message: string | null }; + assert.ok(row.error_message, "error_message should be stored"); + assert.ok(!row.error_message.includes("at /internal/"), "error must not contain raw stack"); + } finally { + cleanup(raw, dir); + } +}); + +test("finalize: output truncated to 32KB with valid JSON envelope", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const claim = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(claim.kind, "claimed"); + // Build a nested object that serializes to >32KB after sanitization + // (sanitizeUpstreamDetails truncates arrays to 32 elements, and strips 'key'-like names) + const bigObj: Record = {}; + for (let i = 0; i < 2000; i++) { + bigObj[`field_${i}_padding`] = "x".repeat(20); + } + finalizeServerToolExecution( + { ...BASE_FINALIZE, executionId: claim.executionId, output: bigObj }, + adapter + ); + const row = raw + .prepare("SELECT output FROM server_tool_executions WHERE id = ?") + .get(claim.executionId) as { output: string | null }; + assert.ok(row.output, "output must be stored"); + assert.ok(row.output.length <= 32768, `output length ${row.output.length} must be <=32768`); + const parsed = JSON.parse(row.output); + assert.ok(parsed.truncated === true, "truncated envelope must have truncated:true"); + assert.ok(typeof parsed.preview === "string", "truncated envelope must have preview string"); + } finally { + cleanup(raw, dir); + } +}); + +test("finalize: raw tool arguments are never stored in the execution row", () => { + const { adapter, dir, raw } = makeTempDb(); + try { + const claim = claimServerToolExecution(BASE_INPUT, adapter); + assert.equal(claim.kind, "claimed"); + finalizeServerToolExecution( + { ...BASE_FINALIZE, executionId: claim.executionId, output: { key: "value" } }, + adapter + ); + const row = raw + .prepare("SELECT output, input_digest FROM server_tool_executions WHERE id = ?") + .get(claim.executionId) as { output: string | null; input_digest: string }; + assert.ok(row.output, "output should be stored"); + assert.ok(row.input_digest, "input_digest must be digest only, not raw args"); + } finally { + cleanup(raw, dir); + } +}); + +test("contention: two independent SQLite handles on same file — loser UNIQUE caught cleanly", () => { + const { adapter: adapter1, dir, raw: raw1 } = makeTempDb(); + const { adapter: adapter2, raw: raw2 } = makeSecondAdapter(dir); + try { + const claim1 = claimServerToolExecution(BASE_INPUT, adapter1); + assert.equal(claim1.kind, "claimed"); + const claim2 = claimServerToolExecution(BASE_INPUT, adapter2); + assert.ok( + claim2.kind === "replay" || claim2.kind === "in_progress", + `Expected replay or in_progress, got ${claim2.kind}` + ); + } finally { + raw1.close(); + raw2.close(); + fs.rmSync(dir, { recursive: true, force: true }); + } +}); + +test("fence: in-process promise joiner replays success result", async () => { + const { adapter, dir, raw } = makeTempDb(); + let handlerCallCount = 0; + const handler = async () => { + handlerCallCount++; + return "handler-result"; + }; + + const p1 = runWithServerToolFence( + { + apiKeyId: "key-1", + requestIdentity: "key-1:req:body", + toolCallId: "call-1", + toolName: "tool_a", + arguments: { q: "test" }, + leaseDurationMs: 60_000, + execute: handler, + }, + adapter + ); + const p2 = runWithServerToolFence( + { + apiKeyId: "key-1", + requestIdentity: "key-1:req:body", + toolCallId: "call-1", + toolName: "tool_a", + arguments: { q: "test" }, + leaseDurationMs: 60_000, + execute: handler, + }, + adapter + ); + + const [r1, r2] = await Promise.all([p1, p2]); + assert.equal(handlerCallCount, 1, "handler should execute exactly once"); + assert.equal(r1.kind, "executed"); + assert.equal(r2.kind, "replayed"); + if (r2.kind === "replayed") { + assert.equal(r2.status, "success"); + assert.equal(r2.value, "handler-result"); + } + cleanup(raw, dir); +}); diff --git a/tests/unit/skills-executor.test.ts b/tests/unit/skills-executor.test.ts index 17349e84b2..7a95793745 100644 --- a/tests/unit/skills-executor.test.ts +++ b/tests/unit/skills-executor.test.ts @@ -293,3 +293,44 @@ test("skillExecutor turns handler errors and timeouts into error executions", as assert.equal(timedOut.output, null); assert.match(timedOut.errorMessage, /timed out/i); }); + +// ─── Task 3: executeClaimed separation from execute RED tests ───────────────── + +test("executeClaimed executes registered handler and returns SkillExecution without writing skill_executions row", async () => { + await registerEchoSkill(); + + skillExecutor.registerHandler("echo-handler", async (input, context) => ({ + echoed: `${input.value}:${context.apiKeyId}`, + })); + + const execution = await skillExecutor.executeClaimed( + "echo@1.0.0", + { value: "claimed" }, + { apiKeyId: "key-a", sessionId: "session-claimed" }, + "test-execution-id" + ); + + assert.equal(execution.status, "success"); + assert.deepEqual(execution.output, { echoed: "claimed:key-a" }); + + // Must NOT write to skill_executions (only execute() does). + const count = skillExecutor.countExecutions("key-a"); + assert.equal(count, 0, "executeClaimed must not write skill_executions row"); +}); + +test("execute still writes exactly 1 skill_executions row (existing contract preserved)", async () => { + await registerEchoSkill(); + + skillExecutor.registerHandler("echo-handler", async (input) => ({ + echoed: input.value, + })); + + await skillExecutor.execute( + "echo@1.0.0", + { value: "persist" }, + { apiKeyId: "key-a", sessionId: "session-persist" } + ); + + const count = skillExecutor.countExecutions("key-a"); + assert.equal(count, 1, "execute must write exactly 1 skill_executions row"); +}); diff --git a/tests/unit/skills-interception-server-owned.test.ts b/tests/unit/skills-interception-server-owned.test.ts new file mode 100644 index 0000000000..f177878327 --- /dev/null +++ b/tests/unit/skills-interception-server-owned.test.ts @@ -0,0 +1,795 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +// Local shape for formatter return values — avoids explicit any in assertions. +type FormattedResponse = { + choices?: Array<{ + message: { + content: string | null; + tool_calls?: Array<{ id: string; function: { name: string; arguments: string } }>; + }; + finish_reason: string; + }>; + content?: Array<{ type: string; id?: string; text?: string }>; + stop_reason?: string; + stop_sequence?: string | null; + output?: Array<{ type: string; call_id?: string; name?: string; arguments?: string }>; + response?: { + output?: Array<{ type: string; call_id?: string; name?: string; arguments?: string }>; + }; +}; + +const TEST_ROOT = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-skills-interception-owned-")); +const TEST_DATA_DIR = path.join(TEST_ROOT, "data"); +const TEST_PLUGINS_DIR = path.join(TEST_ROOT, "plugins"); +const ORIGINAL_DATA_DIR = process.env.DATA_DIR; +const ORIGINAL_PLUGINS_DIR = process.env.OMNIROUTE_PLUGINS_DIR; +fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +fs.mkdirSync(TEST_PLUGINS_DIR, { recursive: true }); +process.env.DATA_DIR = TEST_DATA_DIR; +process.env.OMNIROUTE_PLUGINS_DIR = TEST_PLUGINS_DIR; + +const coreDb = await import("../../src/lib/db/core.ts"); +const { skillRegistry } = await import("../../src/lib/skills/registry.ts"); +const { skillExecutor } = await import("../../src/lib/skills/executor.ts"); +const { + classifyServerOwnedCalls, + formatEscapeHatchResponse, + executeServerOwned, + ServerOwnedExecutionError, +} = await import("../../src/lib/skills/interception.ts"); + +function resetRuntime() { + skillRegistry["registeredSkills"].clear(); + skillRegistry["versionCache"].clear(); + skillExecutor["handlers"].clear(); + skillExecutor.setTimeout(50); +} + +async function resetStorage() { + resetRuntime(); + coreDb.resetDbInstance(); + fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); + fs.mkdirSync(TEST_DATA_DIR, { recursive: true }); +} + +async function registerRuntimeSkills() { + await skillRegistry.register({ + name: "lookup", + version: "1.0.0", + description: "lookup records", + schema: { input: { id: "string" }, output: { record: "string" } }, + handler: "lookup-handler", + enabled: true, + apiKeyId: "key-a", + }); + await skillRegistry.register({ + name: "broken", + version: "1.0.0", + description: "always fails", + schema: { input: {}, output: {} }, + handler: "broken-handler", + enabled: true, + apiKeyId: "key-a", + }); + + skillExecutor.registerHandler("lookup-handler", async (input) => ({ + record: `resolved:${input.id}`, + })); + skillExecutor.registerHandler("broken-handler", async () => { + throw new Error("skill failure"); + }); +} + +test.beforeEach(async () => { + process.env.DATA_DIR = TEST_DATA_DIR; + process.env.OMNIROUTE_PLUGINS_DIR = TEST_PLUGINS_DIR; + await resetStorage(); + await registerRuntimeSkills(); +}); + +test.after(() => { + resetRuntime(); + coreDb.resetDbInstance(); + if (ORIGINAL_DATA_DIR === undefined) delete process.env.DATA_DIR; + else process.env.DATA_DIR = ORIGINAL_DATA_DIR; + if (ORIGINAL_PLUGINS_DIR === undefined) delete process.env.OMNIROUTE_PLUGINS_DIR; + else process.env.OMNIROUTE_PLUGINS_DIR = ORIGINAL_PLUGINS_DIR; + fs.rmSync(TEST_ROOT, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); +}); + +// ─── Task 3: classifyServerOwnedCalls + formatEscapeHatchResponse RED tests ── + +test("classifyServerOwnedCalls: owner set builtin/custom → serverOwned; registry-registered but not in owner set → clientNative", async () => { + const calls = [ + { id: "c1", name: "http_request", arguments: {} }, + { id: "c2", name: "lookup@1.0.0", arguments: {} }, + { id: "c3", name: "Bash", arguments: {} }, + ]; + + const result = await classifyServerOwnedCalls(calls, { + apiKeyId: "key-a", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + injectedCustomSkillNames: ["lookup@1.0.0"], + customSkillExecutionEnabled: true, + }); + + assert.equal(result.serverOwned.length, 2); + assert.equal(result.serverOwned[0].id, "c1"); + assert.equal(result.serverOwned[1].id, "c2"); + assert.equal(result.clientNative.length, 1); + assert.equal(result.clientNative[0].id, "c3"); +}); + +test("classifyServerOwnedCalls: client same-name memory_search → not server-owned", async () => { + const calls = [ + { id: "c1", name: "memory_search", arguments: {} }, + { id: "c2", name: "http_request", arguments: {} }, + ]; + + const result = await classifyServerOwnedCalls(calls, { + apiKeyId: "key-a", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + // memory_search NOT in builtinToolNames (client owns it) + injectedCustomSkillNames: [], + customSkillExecutionEnabled: true, + }); + + // memory_search is client-native because it's not in any owner set. + assert.equal(result.clientNative.length, 1); + assert.equal(result.clientNative[0].id, "c1"); + assert.equal(result.serverOwned.length, 1); + assert.equal(result.serverOwned[0].id, "c2"); +}); + +test("classifyServerOwnedCalls: registered skill with client same-name → client-native (not server-owned)", async () => { + // Register a skill that the client also declares with the same encoded name. + await skillRegistry.register({ + name: "collision-check", + version: "1.0.0", + description: "collision test", + schema: { input: {}, output: {} }, + handler: "collision-handler", + enabled: true, + apiKeyId: "key-a", + mode: "on", + }); + + const encodedName = (await import("../../src/lib/skills/injection.ts")).encodeSkillToolName( + "collision-check", + "1.0.0" + ); + + // Client declares a tool with the same encoded name. + const calls = [ + { id: "c1", name: encodedName, arguments: {} }, + { id: "c2", name: "http_request", arguments: {} }, + ]; + + const result = await classifyServerOwnedCalls(calls, { + apiKeyId: "key-a", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + injectedCustomSkillNames: [], // empty: client collision prevented injection + customSkillExecutionEnabled: true, + }); + + // The registered skill with client same-name must be client-native. + assert.equal(result.clientNative.length, 1); + assert.equal(result.clientNative[0].id, "c1"); + assert.equal(result.serverOwned.length, 1); + assert.equal(result.serverOwned[0].id, "c2"); + + skillRegistry["registeredSkills"].clear(); + skillRegistry["versionCache"].clear(); +}); + +test("formatEscapeHatchResponse: mixed OpenAI — strip server calls, append results to content, keep client calls, finish_reason:tool_calls", async () => { + const response = { + choices: [ + { + message: { + content: null, + tool_calls: [ + { id: "srv1", function: { name: "http_request", arguments: "{}" } }, + { id: "cli1", function: { name: "Bash", arguments: '{"cmd":"ls"}' } }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }; + + const serverCalls = [{ id: "srv1", name: "http_request", arguments: {} }]; + const clientCalls = [{ id: "cli1", name: "Bash", arguments: { cmd: "ls" } }]; + const results = [{ id: "srv1", name: "http_request", result: { status: 200 }, replayed: false }]; + + const formatted = formatEscapeHatchResponse( + response, + serverCalls, + results, + clientCalls, + "openai" + ); + + const choice = (formatted as FormattedResponse).choices[0]; + // Server call stripped from tool_calls, client call kept. + assert.equal(choice.message.tool_calls.length, 1); + assert.equal(choice.message.tool_calls[0].id, "cli1"); + // Server result appended to content. + assert.ok(typeof choice.message.content === "string"); + assert.ok(choice.message.content.includes("200")); + // finish_reason stays tool_calls (mixed). + assert.equal(choice.finish_reason, "tool_calls"); +}); + +test("formatEscapeHatchResponse: all-server OpenAI — strip tool calls, append results, finish_reason:stop", async () => { + const response = { + choices: [ + { + message: { + content: null, + tool_calls: [{ id: "srv1", function: { name: "http_request", arguments: "{}" } }], + }, + finish_reason: "tool_calls", + }, + ], + }; + + const serverCalls = [{ id: "srv1", name: "http_request", arguments: {} }]; + const results = [{ id: "srv1", name: "http_request", result: { ok: true }, replayed: false }]; + + const formatted = formatEscapeHatchResponse(response, serverCalls, results, [], "openai"); + + const choice = (formatted as FormattedResponse).choices[0]; + assert.ok( + !choice.message.tool_calls || choice.message.tool_calls.length === 0, + "all-server must have no remaining tool_calls" + ); + assert.ok(typeof choice.message.content === "string"); + assert.ok(choice.message.content.includes("ok")); + assert.equal(choice.finish_reason, "stop"); +}); + +test("formatEscapeHatchResponse: Claude mixed — strip server tool_use, keep client tool_use, end_turn stays", async () => { + const response = { + content: [ + { type: "tool_use", id: "srv1", name: "http_request", input: {} }, + { type: "tool_use", id: "cli1", name: "Bash", input: { cmd: "ls" } }, + ], + stop_reason: "tool_use", + }; + + const serverCalls = [{ id: "srv1", name: "http_request", arguments: {} }]; + const clientCalls = [{ id: "cli1", name: "Bash", arguments: { cmd: "ls" } }]; + const results = [{ id: "srv1", name: "http_request", result: { ok: true }, replayed: false }]; + + const formatted = formatEscapeHatchResponse( + response, + serverCalls, + results, + clientCalls, + "claude" + ); + + const content = (formatted as FormattedResponse).content as Array<{ type: string; id?: string }>; + // Server tool_use removed, client tool_use kept. + const toolUses = content.filter((b) => b.type === "tool_use"); + assert.equal(toolUses.length, 1); + assert.equal(toolUses[0].id, "cli1"); + // stop_reason stays tool_use (mixed). + assert.equal((formatted as FormattedResponse).stop_reason, "tool_use"); +}); + +test("formatEscapeHatchResponse: Claude all-server — strip tool_use, append text, end_turn", async () => { + const response = { + content: [{ type: "tool_use", id: "srv1", name: "http_request", input: {} }], + stop_reason: "tool_use", + }; + + const serverCalls = [{ id: "srv1", name: "http_request", arguments: {} }]; + const results = [{ id: "srv1", name: "http_request", result: { ok: true }, replayed: false }]; + + const formatted = formatEscapeHatchResponse(response, serverCalls, results, [], "claude"); + + const content = (formatted as FormattedResponse).content as Array<{ type: string }>; + const toolUses = content.filter((b) => b.type === "tool_use"); + assert.equal(toolUses.length, 0); + const textBlocks = content.filter((b) => b.type === "text"); + assert.ok(textBlocks.length > 0); + assert.equal((formatted as FormattedResponse).stop_reason, "end_turn"); +}); + +test("formatEscapeHatchResponse: formatter does not call interceptToolCalls or any handler (purity)", async () => { + // Purity proof: the formatter is synchronous and its source must not contain + // calls to interceptToolCalls, skillExecutor, or handler invocations. + const fnSource = formatEscapeHatchResponse.toString(); + assert.ok( + !fnSource.includes("interceptToolCalls"), + "formatter source must not reference interceptToolCalls" + ); + assert.ok( + !fnSource.includes("skillExecutor"), + "formatter source must not reference skillExecutor" + ); + assert.ok(!fnSource.includes("await"), "formatter must be synchronous (no await)"); + + // Also verify it returns immediately without side effects. + const response = { + choices: [ + { + message: { + content: null, + tool_calls: [{ id: "srv1", function: { name: "http_request", arguments: "{}" } }], + }, + finish_reason: "tool_calls", + }, + ], + }; + + const result = formatEscapeHatchResponse( + response, + [{ id: "srv1", name: "http_request", arguments: {} }], + [{ id: "srv1", name: "http_request", result: { ok: true }, replayed: false }], + [], + "openai" + ); + + assert.ok(result, "formatter returns a result"); + assert.ok(result.choices[0].message.content, "formatter populates content"); +}); + +test("formatEscapeHatchResponse: Responses wrapper is byte-identical for function_call_output", async () => { + const response = { + object: "response", + output: [{ type: "function_call", call_id: "call1", name: "lookup@1.0.0", arguments: "{}" }], + }; + + const serverCalls = [{ id: "call1", name: "lookup@1.0.0", arguments: {} }]; + const results = [ + { id: "call1", name: "lookup@1.0.0", result: { record: "42" }, replayed: false }, + ]; + + const formatted = formatEscapeHatchResponse(response, serverCalls, results, [], "openai"); + + // Responses format: original output + function_call_output appended. + const output = (formatted as FormattedResponse).output; + assert.equal(output.length, 2); + assert.equal(output[0].type, "function_call"); + assert.equal(output[1].type, "function_call_output"); + assert.equal(output[1].call_id, "call1"); +}); + +// ─── F3: nested Responses output formatter ────────────────────────────────── + +test("formatEscapeHatchResponse: nested {response:{output}} appends function_call_output to nested output, not top-level", async () => { + const nestedResponse = { + object: "response", + response: { + output: [{ type: "function_call", call_id: "nc1", name: "lookup@1.0.0", arguments: "{}" }], + }, + }; + + const serverCalls = [{ id: "nc1", name: "lookup@1.0.0", arguments: {} }]; + const results = [ + { id: "nc1", name: "lookup@1.0.0", result: { record: "nested-42" }, replayed: false }, + ]; + + const formatted = formatEscapeHatchResponse(nestedResponse, serverCalls, results, [], "openai"); + + // Must append to nested response.output, not top-level output. + const nestedOutput = (formatted as { response?: { output?: unknown[] } }).response?.output; + assert.ok(Array.isArray(nestedOutput), "nested response.output must be an array"); + assert.equal(nestedOutput.length, 2, "nested output must have original + appended"); + assert.equal(nestedOutput[0].type, "function_call"); + assert.equal(nestedOutput[1].type, "function_call_output"); + assert.equal((nestedOutput[1] as { call_id: string }).call_id, "nc1"); + + // Top-level must NOT have an output array. + assert.equal( + Array.isArray((formatted as { output?: unknown[] }).output), + false, + "top-level output must not exist" + ); +}); + +// ─── F1: executeServerOwned RED tests ─────────────────────────────────────── + +test("executeServerOwned: requires requestIdentity when executionFenceEnabled", async () => { + const calls = [{ id: "c1", name: "http_request", arguments: { url: "https://example.com" } }]; + const context = { + apiKeyId: "key-fence", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + executionFenceEnabled: true, + // requestIdentity is deliberately missing + }; + + await assert.rejects( + () => executeServerOwned(calls, context), + /requestIdentity/, + "must require requestIdentity when executionFenceEnabled" + ); +}); + +test("executeServerOwned: dispatches memory builtin and returns ExecutedToolResult", async () => { + const calls = [{ id: "c1", name: "memory_search", arguments: { query: "test" } }]; + const context = { + apiKeyId: "key-mem", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["memory_search"], + executionFenceEnabled: false, + }; + + const results = await executeServerOwned(calls, context); + assert.equal(results.length, 1); + assert.equal(results[0].id, "c1"); + assert.equal(results[0].name, "memory_search"); + assert.equal(results[0].replayed, false); + assert.ok(results[0].result !== undefined, "result must be present"); +}); + +test("executeServerOwned: dispatches ordinary builtin (http_request)", async () => { + const calls = [{ id: "c1", name: "http_request", arguments: { url: "https://example.com" } }]; + const context = { + apiKeyId: "key-builtin", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + executionFenceEnabled: false, + }; + + const results = await executeServerOwned(calls, context); + assert.equal(results.length, 1); + assert.equal(results[0].id, "c1"); + assert.equal(results[0].name, "http_request"); + assert.equal(results[0].replayed, false); +}); + +test("executeServerOwned: surfaces identity_conflict as typed error, never feeds to model", async () => { + // Custom skill call with no matching handler — should surface as error + const calls = [{ id: "c1", name: "missing-skill@1.0.0", arguments: {} }]; + const context = { + apiKeyId: "key-err", + sessionId: "s1", + requestId: "r1", + injectedCustomSkillNames: ["missing-skill@1.0.0"], + customSkillExecutionEnabled: true, + executionFenceEnabled: false, + }; + + const results = await executeServerOwned(calls, context); + assert.equal(results.length, 1); + assert.equal(results[0].id, "c1"); + assert.equal(results[0].replayed, false); + // Result must contain an error indicator + const resultRecord = results[0].result as Record; + assert.ok( + resultRecord && (resultRecord.error || resultRecord.status), + "result must contain error indicator" + ); +}); + +// ─── Fix Round 2: Defect 1 — typed errors for fence control-flow states ──── + +test("ServerOwnedExecutionError is an exported class with code and httpStatus", () => { + const err = new ServerOwnedExecutionError("test", "TEST_CODE", 409); + assert.ok(err instanceof Error); + assert.equal(err.code, "TEST_CODE"); + assert.equal(err.httpStatus, 409); + assert.equal(err.message, "test"); +}); + +test("executeServerOwned: in_progress fence state → throws ServerOwnedExecutionError with TOOL_IN_PROGRESS (409)", async () => { + const { setFenceFnForTesting } = await import("../../src/lib/skills/interception.ts"); + const mockFence = async () => ({ kind: "in_progress" as const }); + setFenceFnForTesting(mockFence); + + try { + const calls = [{ id: "c1", name: "http_request", arguments: {} }]; + const context = { + apiKeyId: "key-fence", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + executionFenceEnabled: true, + requestIdentity: "identity-1", + }; + + try { + await executeServerOwned(calls, context); + assert.fail("must throw"); + } catch (e: unknown) { + assert.ok(e instanceof ServerOwnedExecutionError); + assert.equal(e.code, "TOOL_IN_PROGRESS"); + assert.equal(e.httpStatus, 409); + } + } finally { + setFenceFnForTesting(null); + } +}); + +test("executeServerOwned: unknown fence state → throws ServerOwnedExecutionError with TOOL_STATE_UNKNOWN (500)", async () => { + const { setFenceFnForTesting } = await import("../../src/lib/skills/interception.ts"); + const mockFence = async () => ({ kind: "unknown" as const }); + setFenceFnForTesting(mockFence); + + try { + const calls = [{ id: "c1", name: "http_request", arguments: {} }]; + const context = { + apiKeyId: "key-fence", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + executionFenceEnabled: true, + requestIdentity: "identity-1", + }; + + try { + await executeServerOwned(calls, context); + assert.fail("must throw"); + } catch (e: unknown) { + assert.ok(e instanceof ServerOwnedExecutionError); + assert.equal(e.code, "TOOL_STATE_UNKNOWN"); + assert.equal(e.httpStatus, 500); + } + } finally { + setFenceFnForTesting(null); + } +}); + +test("executeServerOwned: identity_conflict fence state → throws ServerOwnedExecutionError with IDENTITY_CONFLICT (409)", async () => { + const { setFenceFnForTesting } = await import("../../src/lib/skills/interception.ts"); + const mockFence = async () => ({ kind: "identity_conflict" as const }); + setFenceFnForTesting(mockFence); + + try { + const calls = [{ id: "c1", name: "http_request", arguments: {} }]; + const context = { + apiKeyId: "key-fence", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + executionFenceEnabled: true, + requestIdentity: "identity-1", + }; + + try { + await executeServerOwned(calls, context); + assert.fail("must throw"); + } catch (e: unknown) { + assert.ok(e instanceof ServerOwnedExecutionError); + assert.equal(e.code, "IDENTITY_CONFLICT"); + assert.equal(e.httpStatus, 409); + } + } finally { + setFenceFnForTesting(null); + } +}); + +// ─── Fix Round 3: Defect 1 — replay error/timeout detection ─────────────── + +test("executeServerOwned: error replay → throws ServerOwnedExecutionError with TOOL_EXECUTION_ERROR (500)", async () => { + const { setFenceFnForTesting } = await import("../../src/lib/skills/interception.ts"); + const mockFence = async () => ({ + kind: "replayed" as const, + value: null, + status: "error" as const, + errorMessage: "handler crashed", + }); + setFenceFnForTesting(mockFence); + + try { + const calls = [{ id: "c1", name: "http_request", arguments: {} }]; + const context = { + apiKeyId: "key-fence", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + executionFenceEnabled: true, + requestIdentity: "identity-1", + }; + + try { + await executeServerOwned(calls, context); + assert.fail("must throw"); + } catch (e: unknown) { + assert.ok(e instanceof ServerOwnedExecutionError); + assert.equal(e.code, "TOOL_EXECUTION_ERROR"); + assert.equal(e.httpStatus, 500); + assert.equal(e.message, "handler crashed"); + } + } finally { + setFenceFnForTesting(null); + } +}); + +test("executeServerOwned: timeout replay → throws ServerOwnedExecutionError with TOOL_EXECUTION_TIMEOUT (504)", async () => { + const { setFenceFnForTesting } = await import("../../src/lib/skills/interception.ts"); + const mockFence = async () => ({ + kind: "replayed" as const, + value: null, + status: "timeout" as const, + errorMessage: "execution exceeded deadline", + }); + setFenceFnForTesting(mockFence); + + try { + const calls = [{ id: "c1", name: "http_request", arguments: {} }]; + const context = { + apiKeyId: "key-fence", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + executionFenceEnabled: true, + requestIdentity: "identity-1", + }; + + try { + await executeServerOwned(calls, context); + assert.fail("must throw"); + } catch (e: unknown) { + assert.ok(e instanceof ServerOwnedExecutionError); + assert.equal(e.code, "TOOL_EXECUTION_TIMEOUT"); + assert.equal(e.httpStatus, 504); + assert.equal(e.message, "execution exceeded deadline"); + } + } finally { + setFenceFnForTesting(null); + } +}); + +test("executeServerOwned: success replay → returns ExecutedToolResult with replayed:true", async () => { + const { setFenceFnForTesting } = await import("../../src/lib/skills/interception.ts"); + const mockFence = async () => ({ + kind: "replayed" as const, + value: { cached: true }, + status: "success" as const, + errorMessage: null, + }); + setFenceFnForTesting(mockFence); + + try { + const calls = [{ id: "c1", name: "http_request", arguments: {} }]; + const context = { + apiKeyId: "key-fence", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + executionFenceEnabled: true, + requestIdentity: "identity-1", + }; + + const results = await executeServerOwned(calls, context); + assert.equal(results.length, 1); + assert.equal(results[0].id, "c1"); + assert.equal(results[0].replayed, true); + assert.deepEqual(results[0].result, { cached: true }); + } finally { + setFenceFnForTesting(null); + } +}); + +test("executeServerOwned: error replay with null errorMessage → uses default message", async () => { + const { setFenceFnForTesting } = await import("../../src/lib/skills/interception.ts"); + const mockFence = async () => ({ + kind: "replayed" as const, + value: null, + status: "error" as const, + errorMessage: null, + }); + setFenceFnForTesting(mockFence); + + try { + const calls = [{ id: "c1", name: "http_request", arguments: {} }]; + const context = { + apiKeyId: "key-fence", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + executionFenceEnabled: true, + requestIdentity: "identity-1", + }; + + try { + await executeServerOwned(calls, context); + assert.fail("must throw"); + } catch (e: unknown) { + assert.ok(e instanceof ServerOwnedExecutionError); + assert.equal(e.code, "TOOL_EXECUTION_ERROR"); + assert.equal(e.message, "Tool execution failed"); + } + } finally { + setFenceFnForTesting(null); + } +}); + +// ─── Fix Round 2: Defect 2 — executeClaimed non-success throws ────────────── + +test("executeClaimed: custom skill handler returns non-success status → executeClaimed throws safe error", async () => { + // Register a custom skill whose handler returns a failure output + await skillRegistry.register({ + name: "fail-skill", + version: "1.0.0", + description: "always returns failure status", + schema: { input: {}, output: {} }, + handler: "fail-handler", + enabled: true, + apiKeyId: "key-a", + }); + + skillExecutor.registerHandler("fail-handler", async () => ({ + status: "failed", + message: "something went wrong", + })); + + await assert.rejects( + () => + skillExecutor.executeClaimed( + "fail-skill", + {}, + { apiKeyId: "key-a", sessionId: "s1" }, + "exec-fail" + ), + /Skill execution failed/, + "executeClaimed must throw when handler returns non-success status" + ); + + skillRegistry["registeredSkills"].clear(); + skillRegistry["versionCache"].clear(); +}); + +// ─── Fix Round 2: Defect 3 — LEASE_DURATION_MS = 120000 ───────────────────── + +test("LEASE_DURATION_MS is 120000 to match loop wall-clock upper bound", async () => { + const fs = await import("node:fs"); + const sourceCode = fs.readFileSync( + new URL("../../src/lib/skills/interception.ts", import.meta.url), + "utf8" + ); + const match = sourceCode.match(/const\s+LEASE_DURATION_MS\s*=\s*([\d_]+)/); + assert.ok(match, "LEASE_DURATION_MS must be defined in interception.ts"); + const value = Number(match[1].replace(/_/g, "")); + assert.equal(value, 120_000, "LEASE_DURATION_MS must be 120000 (not 30000)"); +}); + +// ─── Fix Round 2: Defect 4 — classifyServerOwnedCalls no DB load ───────────── + +test("classifyServerOwnedCalls does not call skillRegistry.loadFromDatabase (ownership from owner sets only)", async () => { + let loadFromDatabaseCalled = false; + const origLoad = skillRegistry.loadFromDatabase.bind(skillRegistry); + skillRegistry.loadFromDatabase = async (..._args: unknown[]) => { + loadFromDatabaseCalled = true; + return origLoad(...(_args as [string])); + }; + + try { + const calls = [{ id: "c1", name: "http_request", arguments: {} }]; + await classifyServerOwnedCalls(calls, { + apiKeyId: "key-a", + sessionId: "s1", + requestId: "r1", + builtinToolNames: ["http_request"], + injectedCustomSkillNames: [], + customSkillExecutionEnabled: false, + }); + + assert.equal( + loadFromDatabaseCalled, + false, + "classifyServerOwnedCalls must NOT call loadFromDatabase — owner sets are sufficient" + ); + } finally { + skillRegistry.loadFromDatabase = origLoad; + } +}); diff --git a/tests/unit/stable-json.test.ts b/tests/unit/stable-json.test.ts new file mode 100644 index 0000000000..ff811d3058 --- /dev/null +++ b/tests/unit/stable-json.test.ts @@ -0,0 +1,239 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; + +const __dirname = path.dirname(fileURLToPath(import.meta.url)); +import { + canonicalJson, + canonicalJsonSha256, + deriveToolRequestIdentity, +} from "../../src/lib/skills/stableJson.ts"; + +test("canonicalJson sorts nested object keys and preserves array order", () => { + const a = { z: [{ b: 2, a: 1 }], a: true }; + const b = { a: true, z: [{ a: 1, b: 2 }] }; + assert.equal(canonicalJson(a), canonicalJson(b)); + assert.equal(canonicalJsonSha256(a), canonicalJsonSha256(b)); + assert.notEqual(canonicalJsonSha256({ a: [1, 2] }), canonicalJsonSha256({ a: [2, 1] })); + assert.equal( + deriveToolRequestIdentity({ + apiKeyId: "key-a", + stableClientRequestId: "req-a", + skillRequestId: "internal-1", + postInjectionBody: a, + }), + deriveToolRequestIdentity({ + apiKeyId: "key-a", + stableClientRequestId: "req-a", + skillRequestId: "internal-2", + postInjectionBody: b, + }) + ); +}); + +test("canonicalJson rejects values that cannot form an execution identity", () => { + const cyclic: Record = {}; + cyclic.self = cyclic; + for (const value of [undefined, 1n, Number.NaN, Number.POSITIVE_INFINITY, cyclic]) { + assert.throws(() => canonicalJson(value), /canonical JSON/i); + } + assert.throws(() => canonicalJson({ bad: () => 1 }), /canonical JSON/i); + assert.throws(() => canonicalJson({ bad: Symbol("x") }), /canonical JSON/i); +}); + +test("canonicalJson sorts keys by Unicode code point, not UTF-16 code unit", () => { + // \uE000 (BMP private-use, code point 57344) vs "a" (code point 97) + // Code-point sort: a=97 < \uE000=57344 + const withBmp = { "\uE000": 1, a: 2 }; + const withAstral = { "\u{10000}": 1, a: 2 }; + const sorted1 = canonicalJson(withBmp); + const sorted2 = canonicalJson(withAstral); + // a < \uE000 in code-point order + const bmpKey = "\uE000"; + assert.ok( + sorted1.indexOf('"a"') < sorted1.indexOf(bmpKey), + `expected "a" before BMP key in: ${sorted1}` + ); + // a < \u{10000} in code-point order; astral is encoded as surrogate pair + const astralKey = "\u{10000}"; + assert.ok( + sorted2.indexOf('"a"') < sorted2.indexOf(astralKey), + `expected "a" before astral key in: ${sorted2}` + ); +}); + +test("canonicalJson normalizes -0 to 0", () => { + assert.equal(canonicalJson({ val: -0 }), '{"val":0}'); + assert.equal(canonicalJson({ val: 0 }), '{"val":0}'); + assert.equal(canonicalJsonSha256({ val: -0 }), canonicalJsonSha256({ val: 0 })); +}); + +test("canonicalJson rejects sparse arrays", () => { + const sparse = [1, , 3]; + assert.throws(() => canonicalJson(sparse), /canonical JSON/i); +}); + +test("canonicalJson rejects objects with getters/accessors", () => { + const obj: Record = {}; + Object.defineProperty(obj, "hidden", { + get() { + return 42; + }, + enumerable: true, + }); + assert.throws(() => canonicalJson(obj), /canonical JSON/i); +}); + +test("canonicalJson rejects non-plain objects (Date, Map, Set, class instances)", () => { + assert.throws(() => canonicalJson(new Date()), /canonical JSON/i); + assert.throws(() => canonicalJson(new Map()), /canonical JSON/i); + assert.throws(() => canonicalJson(new Set()), /canonical JSON/i); + class Custom {} + assert.throws(() => canonicalJson(new Custom()), /canonical JSON/i); +}); + +test("deriveToolRequestIdentity uses stableClientRequestId when present", () => { + const body = { a: 1, b: 2 }; + const withStable = deriveToolRequestIdentity({ + apiKeyId: "key-1", + stableClientRequestId: "idempotent-req-1", + skillRequestId: "uuid-internal", + postInjectionBody: body, + }); + const withDifferentInternal = deriveToolRequestIdentity({ + apiKeyId: "key-1", + stableClientRequestId: "idempotent-req-1", + skillRequestId: "different-uuid", + postInjectionBody: body, + }); + // stable key present => internal UUID ignored + assert.equal(withStable, withDifferentInternal); +}); + +test("deriveToolRequestIdentity uses skillRequestId when stableClientRequestId is null", () => { + const body = { a: 1 }; + const r1 = deriveToolRequestIdentity({ + apiKeyId: "key-1", + stableClientRequestId: null, + skillRequestId: "uuid-a", + postInjectionBody: body, + }); + const r2 = deriveToolRequestIdentity({ + apiKeyId: "key-1", + stableClientRequestId: null, + skillRequestId: "uuid-b", + postInjectionBody: body, + }); + // Different skillRequestId => different identity + assert.notEqual(r1, r2); +}); + +test("deriveToolRequestIdentity includes apiKeyId and body digest", () => { + const body = { x: 1 }; + const r1 = deriveToolRequestIdentity({ + apiKeyId: "key-1", + stableClientRequestId: "req", + skillRequestId: "sr", + postInjectionBody: body, + }); + const r2 = deriveToolRequestIdentity({ + apiKeyId: "key-2", + stableClientRequestId: "req", + skillRequestId: "sr", + postInjectionBody: body, + }); + assert.notEqual(r1, r2); + // Body change => different identity + const r3 = deriveToolRequestIdentity({ + apiKeyId: "key-1", + stableClientRequestId: "req", + skillRequestId: "sr", + postInjectionBody: { x: 2 }, + }); + assert.notEqual(r1, r3); +}); + +test("canonicalJson sorts BMP PUA \uE000 before astral \u{10000} in same object", () => { + // U+E000 (BMP private-use, code point 57344) vs U+10000 (Linear B, code point 65536) + // code-point order: E000 < 10000 + // UTF-16 code-unit order: \uD800 (surrogate of 10000) < \uE000 — reversed! + // So default .sort() would place \u{10000} before \uE000. + const obj = { "\u{10000}": 1, "\uE000": 2 }; + const serialized = canonicalJson(obj); + const e000Pos = serialized.indexOf("\uE000"); + const astralPos = serialized.indexOf("\u{10000}"); + assert.ok( + e000Pos < astralPos, + `expected \\uE000 (pos ${e000Pos}) before \\u{10000} (pos ${astralPos}) in: ${serialized}` + ); +}); + +test("canonicalJson rejects objects with getters without invoking the getter", () => { + let getterCallCount = 0; + const obj: Record = {}; + Object.defineProperty(obj, "hidden", { + get() { + getterCallCount++; + return 42; + }, + enumerable: true, + }); + assert.throws(() => canonicalJson(obj), /canonical JSON/i); + assert.equal( + getterCallCount, + 0, + "getter must not be invoked when canonicalJson rejects the object" + ); +}); + +test("termination union in toolLoopTypes.ts source matches expected set", () => { + const src = fs.readFileSync( + path.resolve(__dirname, "../../src/lib/skills/toolLoopTypes.ts"), + "utf8" + ); + + // Extract the termination union block from ServerOwnedToolLoopResult + const loopResultStart = src.indexOf("interface ServerOwnedToolLoopResult"); + assert.ok(loopResultStart !== -1, "ServerOwnedToolLoopResult must exist in toolLoopTypes.ts"); + const afterLoopResult = src.slice(loopResultStart); + const termStart = afterLoopResult.indexOf("termination:"); + assert.ok(termStart !== -1, "termination must exist in ServerOwnedToolLoopResult"); + const afterTerm = afterLoopResult.slice(termStart); + // Match until the closing brace of the interface + const closingBrace = afterTerm.indexOf("\n}"); + assert.ok(closingBrace !== -1, "closing brace must follow termination union"); + const unionText = afterTerm.slice(0, closingBrace); + + const actual = new Set(); + for (const m of unionText.matchAll(/"([^"]+)"/g)) { + actual.add(m[1]); + } + + const expected = new Set([ + "completed", + "client_tools", + "mixed_tools", + "max_followups", + "tool_output_budget", + "deadline", + "client_abort", + "provider_error", + "connection_mismatch", + "execution_in_progress", + "execution_unknown", + "execution_identity_conflict", + "execution_error", + "execution_timeout", + ]); + + // Every expected member must appear in source + for (const val of expected) { + assert.ok(actual.has(val), `termination union in source must include "${val}"`); + } + // No extra unexpected members + for (const val of actual) { + assert.ok(expected.has(val), `unexpected termination member "${val}" in source`); + } +}); diff --git a/tests/unit/tool-loop-usage.test.ts b/tests/unit/tool-loop-usage.test.ts new file mode 100644 index 0000000000..ca1237c23a --- /dev/null +++ b/tests/unit/tool-loop-usage.test.ts @@ -0,0 +1,149 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { aggregateProviderLegUsage } from "../../src/lib/skills/serverOwnedToolLoop.ts"; +import type { ProviderLegUsage } from "../../src/lib/skills/toolLoopTypes.ts"; + +// ─── Basic aggregation ──────────────────────────────────────────────────────── + +test("aggregateProviderLegUsage: sums prompt_tokens, completion_tokens, total_tokens", () => { + const a: ProviderLegUsage = { + prompt_tokens: 100, + completion_tokens: 20, + total_tokens: 120, + }; + const b: ProviderLegUsage = { + prompt_tokens: 150, + completion_tokens: 80, + total_tokens: 230, + }; + + const result = aggregateProviderLegUsage([a, b]); + assert.strictEqual(result.prompt_tokens, 250); + assert.strictEqual(result.completion_tokens, 100); + assert.strictEqual(result.total_tokens, 350, "total_tokens rederived as prompt+completion"); +}); + +// ─── Optional fields: present in at least one leg ──────────────────────────── + +test("aggregateProviderLegUsage: optional fields summed when present", () => { + const a: ProviderLegUsage = { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + cached_tokens: 3, + cache_read_input_tokens: 2, + cache_creation_input_tokens: 1, + reasoning_tokens: 4, + cost_in_usd_ticks: 100, + }; + const b: ProviderLegUsage = { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + // no optional fields + }; + + const result = aggregateProviderLegUsage([a, b]); + assert.strictEqual(result.prompt_tokens, 20); + assert.strictEqual(result.completion_tokens, 10); + assert.strictEqual(result.total_tokens, 30); + assert.strictEqual(result.cached_tokens, 3); + assert.strictEqual(result.cache_read_input_tokens, 2); + assert.strictEqual(result.cache_creation_input_tokens, 1); + assert.strictEqual(result.reasoning_tokens, 4); + assert.strictEqual(result.cost_in_usd_ticks, 100); +}); + +// ─── All-null usage → null ─────────────────────────────────────────────────── + +test("aggregateProviderLegUsage: all null usages → returns zero usage (not null)", () => { + const result = aggregateProviderLegUsage([null, null]); + assert.strictEqual(result.prompt_tokens, 0); + assert.strictEqual(result.completion_tokens, 0); + assert.strictEqual(result.total_tokens, 0); +}); + +test("aggregateProviderLegUsage: mix of null and non-null → sums only non-null", () => { + const a: ProviderLegUsage = { + prompt_tokens: 100, + completion_tokens: 20, + total_tokens: 120, + }; + + const result = aggregateProviderLegUsage([a, null]); + assert.strictEqual(result.prompt_tokens, 100); + assert.strictEqual(result.completion_tokens, 20); + assert.strictEqual(result.total_tokens, 120); +}); + +// ─── total_tokens rederived ─────────────────────────────────────────────────── + +test("aggregateProviderLegUsage: total_tokens is always prompt+completion, not sum of raw total_tokens", () => { + const a: ProviderLegUsage = { + prompt_tokens: 100, + completion_tokens: 20, + total_tokens: 999, // wrong raw value + }; + const b: ProviderLegUsage = { + prompt_tokens: 100, + completion_tokens: 20, + total_tokens: 999, // wrong raw value + }; + + const result = aggregateProviderLegUsage([a, b]); + assert.strictEqual(result.total_tokens, 240, "100+20 + 100+20 = 240, not 1998"); +}); + +// ─── Optional fields absent when no leg defines them ────────────────────────── + +test("aggregateProviderLegUsage: optional fields absent when no leg defines them", () => { + const a: ProviderLegUsage = { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + }; + + const result = aggregateProviderLegUsage([a]); + assert.strictEqual(result.cached_tokens, undefined); + assert.strictEqual(result.cache_read_input_tokens, undefined); + assert.strictEqual(result.cache_creation_input_tokens, undefined); + assert.strictEqual(result.reasoning_tokens, undefined); + assert.strictEqual(result.cost_in_usd_ticks, undefined); +}); + +// ─── Empty array ────────────────────────────────────────────────────────────── + +test("aggregateProviderLegUsage: empty array → zeros", () => { + const result = aggregateProviderLegUsage([]); + assert.strictEqual(result.prompt_tokens, 0); + assert.strictEqual(result.completion_tokens, 0); + assert.strictEqual(result.total_tokens, 0); +}); + +// ─── Claude-style alias normalization is caller's responsibility ────────────── + +test("aggregateProviderLegUsage: Claude input_tokens/output_tokens are NOT aliases — caller must normalize before aggregation", () => { + // Claude uses input_tokens/output_tokens, not prompt_tokens/completion_tokens. + // The caller must normalize before passing to aggregateProviderLegUsage. + // This test verifies that aggregateProviderLegUsage uses the standard field names. + const a: ProviderLegUsage = { + prompt_tokens: 100, + completion_tokens: 20, + total_tokens: 120, + }; + + const result = aggregateProviderLegUsage([a]); + assert.strictEqual(result.prompt_tokens, 100); + assert.strictEqual(result.completion_tokens, 20); + // If someone passes input_tokens (Claude alias), it would be ignored + const claudeStyle = { + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + input_tokens: 50, + output_tokens: 10, + } as unknown as ProviderLegUsage; + const result2 = aggregateProviderLegUsage([claudeStyle]); + assert.strictEqual(result2.prompt_tokens, 0, "Claude alias not summed into prompt_tokens"); +});