diff --git a/packages/ai/src/route/client.ts b/packages/ai/src/route/client.ts index 03f43d17a19..f4873f347ef 100644 --- a/packages/ai/src/route/client.ts +++ b/packages/ai/src/route/client.ts @@ -5,7 +5,7 @@ import { Endpoint, type EndpointPatch } from "./endpoint" import { RequestExecutor } from "./executor" import { Framing } from "./framing" import { HttpTransport } from "./transport" -import type { HttpRequestTransform, Transport, TransportRuntime } from "./transport" +import type { HttpMiddleware, Transport, TransportRuntime } from "./transport" import { WebSocketExecutor } from "./transport" import type { Protocol } from "./protocol" import { applyCachePolicy } from "../cache-policy" @@ -155,7 +155,7 @@ export interface Interface { } export interface StreamOptions { - readonly transform?: HttpRequestTransform + readonly http?: HttpMiddleware } export interface StreamMethod { @@ -307,7 +307,7 @@ function makeFromTransport( auth: routeInput.auth ?? Auth.none, encodeBody, headers: routeInput.headers, - transform: options?.transform, + middleware: options?.http, }), streamPrepared: (prepared: Prepared, request: LLMRequest, runtime: TransportRuntime) => { const route = `${request.model.provider}/${request.model.route.id}` diff --git a/packages/ai/src/route/executor.ts b/packages/ai/src/route/executor.ts index f7a0fb465b3..73433d05ed4 100644 --- a/packages/ai/src/route/executor.ts +++ b/packages/ai/src/route/executor.ts @@ -20,9 +20,18 @@ import { classifyProviderFailure } from "../provider-error" export interface Interface { readonly execute: ( request: HttpClientRequest.HttpClientRequest, + middleware?: HttpMiddleware, ) => Effect.Effect } +export type HttpHandler = ( + request: HttpClientRequest.HttpClientRequest, +) => Effect.Effect +export type HttpMiddleware = ( + request: HttpClientRequest.HttpClientRequest, + handler: HttpHandler, +) => Effect.Effect + export class Service extends Context.Service()("@opencode/AI/RequestExecutor") {} const BODY_LIMIT = 16_384 @@ -282,12 +291,20 @@ export const layer: Layer.Layer = Layer.e Service, Effect.gen(function* () { const http = yield* HttpClient.HttpClient - const executeOnce = (request: HttpClientRequest.HttpClientRequest) => + const executeOnce = (request: HttpClientRequest.HttpClientRequest, middleware?: HttpMiddleware) => Effect.gen(function* () { const redactedNames = yield* Headers.CurrentRedactedNames - return yield* http - .execute(request) - .pipe(Effect.mapError(toHttpError(redactedNames)), Effect.flatMap(statusError(request, redactedNames))) + if (!middleware) + return yield* http + .execute(request) + .pipe(Effect.mapError(toHttpError(redactedNames)), Effect.flatMap(statusError(request, redactedNames))) + + const response = yield* middleware(request, (input) => + http + .execute(input) + .pipe(Effect.mapError((cause) => (cause instanceof Error ? cause : new Error(String(cause))))), + ).pipe(Effect.mapError(toHttpError(redactedNames))) + return yield* statusError(response.request, redactedNames)(response) }) return Service.of({ execute: executeOnce, diff --git a/packages/ai/src/route/index.ts b/packages/ai/src/route/index.ts index 8b85cfc3f76..eb9f1759c55 100644 --- a/packages/ai/src/route/index.ts +++ b/packages/ai/src/route/index.ts @@ -23,4 +23,4 @@ export type { ApiKeyMode, AuthOverride, ProviderAuthOption } from "./auth-option export type { Definition as EndpointFn, EndpointInput } from "./endpoint" export type { Definition as FramingDef } from "./framing" export type { Protocol as ProtocolDef } from "./protocol" -export type { HttpRequest, HttpRequestTransform, Transport as TransportDef, TransportRuntime } from "./transport" +export type { HttpHandler, HttpMiddleware, Transport as TransportDef, TransportRuntime } from "./transport" diff --git a/packages/ai/src/route/transport/http.ts b/packages/ai/src/route/transport/http.ts index a1de9ec5735..223afb353c1 100644 --- a/packages/ai/src/route/transport/http.ts +++ b/packages/ai/src/route/transport/http.ts @@ -3,7 +3,7 @@ import { Headers, HttpClientRequest } from "effect/unstable/http" import { Auth } from "../auth" import { render as renderEndpoint } from "../endpoint" import { Framing } from "../framing" -import type { Transport, TransportPrepareInput } from "./index" +import type { HttpMiddleware, Transport, TransportPrepareInput } from "./index" import * as ProviderShared from "../../protocols/shared" import { mergeJsonRecords, type LLMRequest } from "../../schema" @@ -19,6 +19,7 @@ export interface JsonRequestParts { export interface HttpPrepared { readonly request: HttpClientRequest.HttpClientRequest readonly framing: Framing.Definition + readonly middleware?: HttpMiddleware } const applyQuery = (url: string, query: Record | undefined) => { @@ -74,21 +75,21 @@ export const httpJson = (input: HttpJsonInput): HttpJs prepare: (prepareInput) => Effect.gen(function* () { const parts = yield* jsonRequestParts({ ...prepareInput }) - const request = { url: parts.url, method: "POST", headers: { ...parts.headers }, body: parts.bodyText } - yield* (prepareInput.transform?.(request) ?? Effect.void) + const request = ProviderShared.jsonPost({ + url: parts.url, + body: parts.bodyText, + headers: parts.headers, + }) return { - request: ProviderShared.jsonPost({ - url: request.url, - body: request.body ?? "", - headers: Headers.fromInput(request.headers), - }), + request, framing: input.framing, + middleware: prepareInput.middleware, } }), frames: (prepared, request, runtime) => Stream.unwrap( runtime.http - .execute(prepared.request) + .execute(prepared.request, prepared.middleware) .pipe( Effect.map((response) => prepared.framing.frame( diff --git a/packages/ai/src/route/transport/index.ts b/packages/ai/src/route/transport/index.ts index 16624cce563..c74e578e8bb 100644 --- a/packages/ai/src/route/transport/index.ts +++ b/packages/ai/src/route/transport/index.ts @@ -1,7 +1,7 @@ import type { Effect, Stream } from "effect" import { Endpoint } from "../endpoint" import { Auth } from "../auth" -import type { Interface as RequestExecutorInterface } from "../executor" +import type { HttpMiddleware, Interface as RequestExecutorInterface } from "../executor" import type { Interface as WebSocketExecutorInterface } from "./websocket" import type { AIError, LLMRequest } from "../../schema" @@ -10,15 +10,6 @@ export interface TransportRuntime { readonly webSocket?: WebSocketExecutorInterface } -export interface HttpRequest { - url: string - readonly method: string - headers: Record - body: string | undefined -} - -export type HttpRequestTransform = (request: HttpRequest) => Effect.Effect - export interface Transport { readonly id: string readonly prepare: (input: TransportPrepareInput) => Effect.Effect @@ -32,8 +23,9 @@ export interface TransportPrepareInput { readonly auth: Auth.Definition readonly encodeBody: (body: Body) => string readonly headers?: (input: { readonly request: LLMRequest }) => Record - readonly transform?: HttpRequestTransform + readonly middleware?: HttpMiddleware } export * as HttpTransport from "./http" +export type { HttpHandler, HttpMiddleware } from "../executor" export { WebSocketExecutor, WebSocketTransport } from "./websocket" diff --git a/packages/ai/test/compile.test.ts b/packages/ai/test/compile.test.ts index 27971bf5c8e..0f1260873e0 100644 --- a/packages/ai/test/compile.test.ts +++ b/packages/ai/test/compile.test.ts @@ -1,6 +1,6 @@ import { describe, expect, test } from "bun:test" -import { Effect, Schema } from "effect" -import { HttpClientRequest } from "effect/unstable/http" +import { Effect, Ref, Schema } from "effect" +import { HttpClientRequest, HttpClientResponse } from "effect/unstable/http" import { LLM, mergeProviderOptions } from "../src" import { AnthropicMessages, OpenAIChat } from "../src/protocols" import { Auth, LLMClient } from "../src/route" @@ -146,12 +146,16 @@ describe("request option precedence", () => { prompt: "Say hello.", }), { - transform: (request) => - Effect.sync(() => { - expect(request.headers.authorization).toBe("Bearer fresh-key") - request.url = "https://proxy.test/v1/chat/completions" - request.headers["x-plugin"] = "transformed" - request.body = JSON.stringify({ transformed: true }) + http: (request, handler) => + Effect.gen(function* () { + return yield* handler( + request.pipe( + HttpClientRequest.setUrl("https://proxy.test/v1/chat/completions"), + HttpClientRequest.setMethod("PUT"), + HttpClientRequest.setHeader("x-plugin", "transformed"), + HttpClientRequest.bodyText(JSON.stringify({ transformed: true }), "application/custom+json"), + ), + ) }), }, ).pipe( @@ -160,7 +164,9 @@ describe("request option precedence", () => { Effect.gen(function* () { const web = yield* HttpClientRequest.toWeb(input.request).pipe(Effect.orDie) expect(web.url).toBe("https://proxy.test/v1/chat/completions") + expect(web.method).toBe("PUT") expect(web.headers.get("x-plugin")).toBe("transformed") + expect(web.headers.get("content-type")).toBe("application/custom+json") expect(decodeJson(input.text)).toEqual({ transformed: true }) return input.respond(sseEvents(deltaChunk({}, "stop")), { headers: { "content-type": "text/event-stream" }, @@ -171,6 +177,82 @@ describe("request option precedence", () => { ), ) + it.effect("transforms the HTTP response before protocol decoding", () => + Effect.gen(function* () { + const response = yield* LLMClient.generate( + LLM.request({ + model: OpenAIChat.route + .with({ endpoint: { baseURL: "https://api.openai.test/v1/" }, auth: Auth.bearer("test") }) + .model({ id: "gpt-4o-mini" }), + prompt: "Say hello.", + }), + { + http: (request, handler) => + Effect.gen(function* () { + const response = yield* handler(request) + return HttpClientResponse.fromWeb( + response.request, + new Response((yield* response.text).replace("network", "hooked"), { + status: response.status, + headers: response.headers, + }), + ) + }), + }, + ).pipe( + Effect.provide( + dynamicResponse((input) => + Effect.succeed( + input.respond(sseEvents(deltaChunk({ content: "network" }, "stop")), { + headers: { "content-type": "text/event-stream" }, + }), + ), + ), + ), + ) + + expect(response.text).toBe("hooked") + }), + ) + + it.effect("can inspect an error response and retry the native request", () => + Effect.gen(function* () { + const attempts = yield* Ref.make(0) + const response = yield* LLMClient.generate( + LLM.request({ + model: OpenAIChat.route + .with({ endpoint: { baseURL: "https://api.openai.test/v1/" }, auth: Auth.bearer("stale") }) + .model({ id: "gpt-4o-mini" }), + prompt: "Say hello.", + }), + { + http: (request, handler) => + Effect.gen(function* () { + const response = yield* handler(request) + expect(response.status).toBe(401) + return yield* handler(HttpClientRequest.setHeader(request, "authorization", "Bearer refreshed")) + }), + }, + ).pipe( + Effect.provide( + dynamicResponse((input) => + Effect.gen(function* () { + yield* Ref.update(attempts, (value) => value + 1) + if (input.request.headers.authorization !== "Bearer refreshed") + return input.respond("unauthorized", { status: 401 }) + return input.respond(sseEvents(deltaChunk({ content: "retried" }, "stop")), { + headers: { "content-type": "text/event-stream" }, + }) + }), + ), + ), + ) + + expect(response.text).toBe("retried") + expect(yield* Ref.get(attempts)).toBe(2) + }), + ) + it.effect("applies raw body overlays after protocol lowering", () => LLMClient.generate( LLM.request({