diff --git a/packages/ai/src/route/client.ts b/packages/ai/src/route/client.ts index 067292329b..2e38b96a86 100644 --- a/packages/ai/src/route/client.ts +++ b/packages/ai/src/route/client.ts @@ -8,6 +8,7 @@ import { HttpTransport } from "./transport" import type { Transport, TransportRuntime } from "./transport" import { WebSocketExecutor } from "./transport" import type { Protocol } from "./protocol" +import type { CallOptions } from "./request-transform" import { applyCachePolicy } from "../cache-policy" import * as ProviderShared from "../protocols/shared" import type { LLMError, PreparedRequestOf, ProtocolID, ProviderOptions } from "../schema" @@ -47,7 +48,11 @@ export interface Route { readonly body: RouteBody readonly with: (patch: RoutePatch) => Route readonly model: (input: RouteMappedModelInput) => Model - readonly prepareTransport: (body: Body, request: LLMRequest) => Effect.Effect + readonly prepareTransport: ( + body: Body, + request: LLMRequest, + options?: CallOptions, + ) => Effect.Effect readonly streamPrepared: ( prepared: Prepared, request: LLMRequest, @@ -158,11 +163,11 @@ export interface Interface { } export interface StreamMethod { - (request: LLMRequest): Stream.Stream + (request: LLMRequest, options?: CallOptions): Stream.Stream } export interface GenerateMethod { - (request: LLMRequest): Effect.Effect + (request: LLMRequest, options?: CallOptions): Effect.Effect } export class Service extends Context.Service()("@opencode/LLMClient") {} @@ -297,7 +302,7 @@ function makeFromTransport( }) }, model: (input) => makeRouteModel(route, input), - prepareTransport: (body, request) => + prepareTransport: (body, request, options) => routeInput.transport.prepare({ body, request, @@ -305,6 +310,7 @@ function makeFromTransport( auth: routeInput.auth ?? Auth.none, encodeBody, headers: routeInput.headers, + transformRequest: options?.transformRequest, }), streamPrepared: (prepared: Prepared, request: LLMRequest, runtime: TransportRuntime) => { const route = `${request.model.provider}/${request.model.route.id}` @@ -373,14 +379,14 @@ export function make( // `compile` is the important boundary: it turns a common `LLMRequest` into a // validated provider body plus transport-private prepared data, but does not // execute transport. -const compile = Effect.fn("LLM.compile")(function* (request: LLMRequest) { +const compile = Effect.fn("LLM.compile")(function* (request: LLMRequest, options?: CallOptions) { const resolved = applyCachePolicy(resolveRequestOptions(request)) const route = resolved.model.route const body = yield* route.body .from(resolved) .pipe(Effect.flatMap(ProviderShared.validateWith(Schema.decodeUnknownEffect(route.body.schema)))) - const prepared = yield* route.prepareTransport(body, resolved) + const prepared = yield* route.prepareTransport(body, resolved, options) return { request: resolved, @@ -403,17 +409,17 @@ const prepareWith = Effect.fn("LLMClient.prepare")(function* (request: LLMReques }) }) -const streamRequestWith = (runtime: TransportRuntime) => (request: LLMRequest) => +const streamRequestWith = (runtime: TransportRuntime) => (request: LLMRequest, options?: CallOptions) => Stream.unwrap( Effect.gen(function* () { - const compiled = yield* compile(request) + const compiled = yield* compile(request, options) return compiled.route.streamPrepared(compiled.prepared, compiled.request, runtime) }), ) const generateWith = (stream: Interface["stream"]) => - Effect.fn("LLM.generate")(function* (request: LLMRequest) { - const state = yield* stream(request).pipe(Stream.runFold(LLMResponse.empty, LLMResponse.reduce)) + Effect.fn("LLM.generate")(function* (request: LLMRequest, options?: CallOptions) { + const state = yield* stream(request, options).pipe(Stream.runFold(LLMResponse.empty, LLMResponse.reduce)) const response = LLMResponse.complete(state) if (response) return response return yield* ProviderShared.eventError( @@ -425,17 +431,17 @@ const generateWith = (stream: Interface["stream"]) => export const prepare = (request: LLMRequest) => prepareWith(request) as Effect.Effect, LLMError> -export function stream(request: LLMRequest): Stream.Stream { +export function stream(request: LLMRequest, options?: CallOptions): Stream.Stream { return Stream.unwrap( Effect.gen(function* () { - return (yield* Service).stream(request) + return (yield* Service).stream(request, options) }), ) as Stream.Stream } -export function generate(request: LLMRequest): Effect.Effect { +export function generate(request: LLMRequest, options?: CallOptions): Effect.Effect { return Effect.gen(function* () { - return yield* (yield* Service).generate(request) + return yield* (yield* Service).generate(request, options) }) as Effect.Effect } diff --git a/packages/ai/src/route/index.ts b/packages/ai/src/route/index.ts index 70db881ea4..08427d9c76 100644 --- a/packages/ai/src/route/index.ts +++ b/packages/ai/src/route/index.ts @@ -9,6 +9,7 @@ export type { Interface as LLMClientShape, Service as LLMClientService, } from "./client" +export type { CallOptions, RequestData, RequestTransform, RequestValue } from "./request-transform" export * from "./executor" export { Auth } from "./auth" export { AuthOptions } from "./auth-options" diff --git a/packages/ai/src/route/request-transform.ts b/packages/ai/src/route/request-transform.ts new file mode 100644 index 0000000000..e278aa34e0 --- /dev/null +++ b/packages/ai/src/route/request-transform.ts @@ -0,0 +1,14 @@ +import type { Effect } from "effect" + +export type RequestValue = null | boolean | number | string | RequestValue[] | { [key: string]: RequestValue } + +export interface RequestData { + readonly headers: Record + readonly body: Record +} + +export type RequestTransform = (request: RequestData) => Effect.Effect + +export interface CallOptions { + readonly transformRequest?: RequestTransform +} diff --git a/packages/ai/src/route/transport/http.ts b/packages/ai/src/route/transport/http.ts index 785e450a5e..469eaca0cb 100644 --- a/packages/ai/src/route/transport/http.ts +++ b/packages/ai/src/route/transport/http.ts @@ -1,4 +1,4 @@ -import { Effect, Stream } from "effect" +import { Effect, Schema, Stream } from "effect" import { Headers, HttpClientRequest } from "effect/unstable/http" import { Auth } from "../auth" import { render as renderEndpoint } from "../endpoint" @@ -6,6 +6,7 @@ import { Framing } from "../framing" import type { Transport, TransportPrepareInput } from "./index" import * as ProviderShared from "../../protocols/shared" import { mergeJsonRecords, type LLMRequest } from "../../schema" +import type { RequestValue } from "../request-transform" export type JsonRequestInput = TransportPrepareInput @@ -86,24 +87,61 @@ const bodyWithOverlay = (body: Body, request: LLMRequest, encodeBody: (bod return yield* ProviderShared.invalidRequest("http.body can only overlay JSON object request bodies") }) -export const jsonRequestParts = (input: JsonRequestInput) => +const isRequestValue = (value: unknown): value is RequestValue => { + if (value === null || ["boolean", "number", "string"].includes(typeof value)) return true + if (Array.isArray(value)) return value.every(isRequestValue) + return ProviderShared.isRecord(value) && Object.values(value).every(isRequestValue) +} + +export const isRequestBody = (value: unknown): value is Record => + ProviderShared.isRecord(value) && Object.values(value).every(isRequestValue) + +const decodeJson = Schema.decodeUnknownEffect(Schema.UnknownFromJsonString) + +export const decodeRequestBody = (text: string) => + decodeJson(text).pipe( + Effect.mapError(() => ProviderShared.invalidRequest("Request hooks require a JSON object body")), + Effect.flatMap((body) => + isRequestBody(body) + ? Effect.succeed(body) + : Effect.fail(ProviderShared.invalidRequest("Request hooks require a JSON object body")), + ), + ) + +export const jsonRequestBaseParts = (input: JsonRequestInput) => Effect.gen(function* () { const url = applyQuery( renderEndpoint(input.endpoint, { request: input.request, body: input.body }).toString(), input.request.http?.query, ) const body = yield* bodyWithOverlay(input.body, input.request, input.encodeBody) + return { + url, + jsonBody: body.jsonBody, + bodyText: body.bodyText, + headers: { + ...input.headers?.({ request: input.request }), + ...input.request.http?.headers, + }, + } + }) + +export const jsonRequestParts = (input: JsonRequestInput) => + Effect.gen(function* () { + const base = yield* jsonRequestBaseParts(input) + const transformRequest = input.transformRequest + const transformed = transformRequest + ? yield* transformRequest({ headers: base.headers, body: yield* decodeRequestBody(base.bodyText) }) + : { headers: base.headers, body: base.jsonBody } + const bodyText = transformRequest ? ProviderShared.encodeJson(transformed.body) : base.bodyText const headers = yield* Auth.toEffect(input.auth)({ request: input.request, method: "POST", - url, - body: body.bodyText, - headers: Headers.fromInput({ - ...input.headers?.({ request: input.request }), - ...input.request.http?.headers, - }), + url: base.url, + body: bodyText, + headers: Headers.fromInput(transformed.headers), }) - return { url, jsonBody: body.jsonBody, bodyText: body.bodyText, headers } + return { url: base.url, jsonBody: transformed.body, bodyText, headers } }) export interface HttpJsonInput<_Body, Frame> { diff --git a/packages/ai/src/route/transport/index.ts b/packages/ai/src/route/transport/index.ts index cf8fef1d08..33ab5dd595 100644 --- a/packages/ai/src/route/transport/index.ts +++ b/packages/ai/src/route/transport/index.ts @@ -4,6 +4,7 @@ import { Auth } from "../auth" import type { Interface as RequestExecutorInterface } from "../executor" import type { Interface as WebSocketExecutorInterface } from "./websocket" import type { LLMError, LLMRequest } from "../../schema" +import type { RequestTransform } from "../request-transform" export interface TransportRuntime { readonly http: RequestExecutorInterface @@ -27,6 +28,7 @@ export interface TransportPrepareInput { readonly auth: Auth.Definition readonly encodeBody: (body: Body) => string readonly headers?: (input: { readonly request: LLMRequest }) => Record + readonly transformRequest?: RequestTransform } export * as HttpTransport from "./http" diff --git a/packages/ai/src/route/transport/websocket.ts b/packages/ai/src/route/transport/websocket.ts index 310121420c..4973224c9d 100644 --- a/packages/ai/src/route/transport/websocket.ts +++ b/packages/ai/src/route/transport/websocket.ts @@ -1,6 +1,8 @@ import { Cause, Context, Effect, Layer, Queue, Stream } from "effect" import { Headers } from "effect/unstable/http" +import { encodeJson } from "../../protocols/shared" import { LLMError, TransportReason } from "../../schema" +import { Auth } from "../auth" import * as HttpTransport from "./http" import type { Transport } from "./index" @@ -228,13 +230,24 @@ export const json = (input: JsonInput): JsonTransp with: (patch) => json({ ...input, ...patch }), prepare: (prepareInput) => Effect.gen(function* () { - const parts = yield* HttpTransport.jsonRequestParts({ - ...prepareInput, + const parts = yield* HttpTransport.jsonRequestBaseParts(prepareInput) + const message = input.encodeMessage(yield* input.toMessage(parts.jsonBody)) + const transformRequest = prepareInput.transformRequest + const transformed = transformRequest + ? yield* transformRequest({ headers: parts.headers, body: yield* HttpTransport.decodeRequestBody(message) }) + : { headers: parts.headers, body: yield* HttpTransport.decodeRequestBody(message) } + const bodyText = transformRequest ? encodeJson(transformed.body) : message + const headers = yield* Auth.toEffect(prepareInput.auth)({ + request: prepareInput.request, + method: "POST", + url: parts.url, + body: bodyText, + headers: Headers.fromInput(transformed.headers), }) return { url: yield* webSocketUrl(parts.url), - headers: parts.headers, - message: input.encodeMessage(yield* input.toMessage(parts.jsonBody)), + headers, + message: bodyText, } }), frames: (prepared, _request, runtime) => { diff --git a/packages/ai/test/prepare.test.ts b/packages/ai/test/prepare.test.ts index 6923c5a678..81de261742 100644 --- a/packages/ai/test/prepare.test.ts +++ b/packages/ai/test/prepare.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test" -import { Effect, Schema } from "effect" +import { Effect, Schema, Stream } from "effect" import { HttpClientRequest } from "effect/unstable/http" import { LLM, mergeProviderOptions } from "../src" import { AnthropicMessages, OpenAIChat } from "../src/protocols" @@ -136,6 +136,41 @@ describe("request option precedence", () => { ), ) + it.effect("transforms provider-native headers and JSON", () => + Effect.gen(function* () { + const model = OpenAIChat.route + .with({ endpoint: { baseURL: "https://api.openai.test/v1/" }, auth: Auth.bearer("test") }) + .model({ id: "gpt-4o-mini" }) + yield* LLMClient.stream(LLM.request({ model, prompt: "Say hello." }), { + transformRequest: (request) => + Effect.sync(() => { + const body = { ...request.body, store: true } + delete body.stream_options + return { + headers: { ...request.headers, "x-plugin": "enabled" }, + body, + } + }), + }).pipe( + Stream.runDrain, + Effect.provide( + dynamicResponse((input) => + Effect.gen(function* () { + const web = yield* HttpClientRequest.toWeb(input.request).pipe(Effect.orDie) + expect(web.headers.get("authorization")).toBe("Bearer test") + expect(web.headers.get("x-plugin")).toBe("enabled") + expect(decodeJson(input.text)).toMatchObject({ store: true }) + expect(decodeJson(input.text)).not.toHaveProperty("stream_options") + return input.respond(sseEvents(deltaChunk({}, "stop")), { + headers: { "content-type": "text/event-stream" }, + }) + }), + ), + ), + ) + }), + ) + it.effect("rejects raw body overlays for protocol-owned roots", () => Effect.gen(function* () { const model = OpenAIChat.route diff --git a/packages/ai/test/provider/openai-responses.test.ts b/packages/ai/test/provider/openai-responses.test.ts index e345cd59b1..0278f216cd 100644 --- a/packages/ai/test/provider/openai-responses.test.ts +++ b/packages/ai/test/provider/openai-responses.test.ts @@ -251,16 +251,34 @@ describe("OpenAI Responses route", () => { }), ), ) - const response = yield* LLMClient.generate( + const text: string[] = [] + yield* LLMClient.stream( LLM.request({ model: OpenAI.configure({ baseURL: "https://api.openai.test/v1/", apiKey: "test" }).responsesWebSocket( "gpt-4.1-mini", ), prompt: "Say hello.", }), - ).pipe(Effect.provide(LLMClient.layer.pipe(Layer.provide(deps)))) + { + transformRequest: (request) => + Effect.sync(() => { + expect(request.body.type).toBe("response.create") + expect(request.body.stream).toBeUndefined() + const body = { ...request.body, plugin: true } + delete body.store + return { headers: { ...request.headers, "x-plugin": "enabled" }, body } + }), + }, + ).pipe( + Stream.runForEach((event) => + Effect.sync(() => { + if (LLMEvent.is.textDelta(event)) text.push(event.text) + }), + ), + Effect.provide(LLMClient.layer.pipe(Layer.provide(deps))), + ) - expect(response.text).toBe("Hi") + expect(text.join("")).toBe("Hi") expect(opened).toEqual([{ url: "wss://api.openai.test/v1/responses", authorization: "Bearer test" }]) expect(closed).toBe(true) expect(sent).toHaveLength(1) @@ -268,7 +286,7 @@ describe("OpenAI Responses route", () => { type: "response.create", model: "gpt-4.1-mini", input: [{ role: "user", content: [{ type: "input_text", text: "Say hello." }] }], - store: false, + plugin: true, }) }), ) diff --git a/packages/core/src/aisdk.ts b/packages/core/src/aisdk.ts index d81d625057..7434c469c5 100644 --- a/packages/core/src/aisdk.ts +++ b/packages/core/src/aisdk.ts @@ -1,5 +1,6 @@ export * as AISDK from "./aisdk" +import { AsyncLocalStorage } from "node:async_hooks" import { makeLocationNode } from "@opencode-ai/util/effect/app-node" import type { JSONSchema7, @@ -29,7 +30,7 @@ import { type ToolDefinition, type UsageInput, } from "@opencode-ai/ai" -import { Auth, Endpoint, type AnyRoute } from "@opencode-ai/ai/route" +import { Auth, Endpoint, type AnyRoute, type CallOptions } from "@opencode-ai/ai/route" import { ProviderShared } from "@opencode-ai/ai/protocols/shared" import { Cause, Context, Effect, Layer, Option, Schema, Scope, Stream } from "effect" import { ModelV2 } from "./model" @@ -40,6 +41,12 @@ type SDK = any type UserContent = Extract["content"] type AssistantContent = Extract["content"] type ToolResultContent = Extract +type TransformRequest = NonNullable + +interface AISDKPrepared { + readonly call: LanguageModelV3CallOptions + readonly transformRequest?: TransformRequest +} export interface SDKEvent { readonly model: ModelV2.Info @@ -103,7 +110,7 @@ function wrapSSE(res: Response, ms: number, ctl: AbortController) { }) } -function prepareOptions(model: ModelV2.Info, pkg: string) { +function prepareOptions(model: ModelV2.Info, pkg: string, requests: AsyncLocalStorage) { const projected = mapBodyToProviderOptions(model, pkg) const options: Record = { name: model.providerID, @@ -150,6 +157,21 @@ function prepareOptions(model: ModelV2.Info, pkg: string) { } } + const headers = new Headers(opts.headers) + const transformRequest = requests.getStore() + if (transformRequest) { + if (typeof opts.body !== "string") throw new Error("Session request hooks require a JSON request body") + const prepared = await Effect.runPromise( + transformRequest({ + headers: Object.fromEntries(headers.entries()), + body: JSON.parse(opts.body), + }), + ) + opts.headers = prepared.headers + opts.body = JSON.stringify(prepared.body) + } + if (!transformRequest) opts.headers = headers + const res = await (typeof customFetch === "function" ? customFetch : fetch)(input, { ...opts, timeout: false, @@ -194,6 +216,7 @@ export const locationLayer = Layer.effect( let languageHooks: ((event: LanguageEvent) => Effect.Effect | void)[] = [] const languages = new Map() const sdks = new Map() + const requests = new AsyncLocalStorage() const functionIDs = new WeakMap() let nextFunctionID = 0 const cacheKey = (input: unknown) => @@ -267,7 +290,7 @@ export const locationLayer = Layer.effect( }) const packageName = ProviderV2.packageName(model.package) - const options = prepareOptions(model, packageName) + const options = prepareOptions(model, packageName, requests) const sdkKey = cacheKey({ providerID: model.providerID, package: packageName, @@ -292,7 +315,7 @@ export const locationLayer = Layer.effect( return language }), model: Effect.fn("AISDK.model")(function* (model) { - return modelFromLanguage(model, yield* service.language(model)) + return modelFromLanguage(model, yield* service.language(model), requests) }), }) return service @@ -301,7 +324,11 @@ export const locationLayer = Layer.effect( export const defaultLayer = locationLayer -function modelFromLanguage(info: ModelV2.Info, language: LanguageModelV3) { +function modelFromLanguage( + info: ModelV2.Info, + language: LanguageModelV3, + requests: AsyncLocalStorage, +) { const packageName = ProviderV2.packageName(info.package!) const projected = mapBodyToProviderOptions(info, packageName) const optionKey = providerOptionKey(packageName, info.providerID) @@ -341,8 +368,13 @@ function modelFromLanguage(info: ModelV2.Info, language: LanguageModelV3) { }, with: () => route, model: (input) => Model.make({ ...input, provider: "provider" in input ? input.provider : info.providerID, route }), - prepareTransport: (body) => Effect.succeed(body), - streamPrepared: (prepared) => streamLanguage(language, prepared as LanguageModelV3CallOptions), + prepareTransport: (body, _request, options) => + Effect.succeed({ call: body as LanguageModelV3CallOptions, transformRequest: options?.transformRequest }), + streamPrepared: (prepared) => { + const request = prepared as AISDKPrepared + if (!request.transformRequest) return streamLanguage(language, request.call) + return streamLanguage(language, request.call, requests, request.transformRequest) + }, } return Model.make({ id: info.modelID ?? info.id, provider: info.providerID, route }) } @@ -533,13 +565,21 @@ function providerOptions(input: LLMRequest["providerOptions"]): SharedV3Provider return Object.fromEntries(Object.entries(input).map(([key, value]) => [key, jsonObject(value)])) } -function streamLanguage(language: LanguageModelV3, options: LanguageModelV3CallOptions) { +function streamLanguage( + language: LanguageModelV3, + options: LanguageModelV3CallOptions, + requests?: AsyncLocalStorage, + transformRequest?: TransformRequest, +) { const state = { step: 0, toolNames: {} as Record } return Stream.concat( Stream.make(LLMEvent.stepStart({ index: state.step })), Stream.unwrap( Effect.tryPromise({ - try: () => language.doStream(options), + try: () => + requests && transformRequest + ? requests.run(transformRequest, () => language.doStream(options)) + : language.doStream(options), catch: (error) => llmError("doStream", error), }).pipe( Effect.map((result) => diff --git a/packages/core/src/plugin/hooks.ts b/packages/core/src/plugin/hooks.ts index 5538601c91..78f11fe923 100644 --- a/packages/core/src/plugin/hooks.ts +++ b/packages/core/src/plugin/hooks.ts @@ -16,6 +16,10 @@ export interface Domains { type Callback = (event: Event) => Effect.Effect export interface Interface { + readonly has: ( + domain: Domain, + name: Name, + ) => boolean readonly register: ( domain: Domain, name: Name, @@ -60,7 +64,7 @@ const layer = Layer.effect( return event }) - return Service.of({ register, trigger }) + return Service.of({ has: (domain, name) => (callbacks.get(key(domain, name))?.length ?? 0) > 0, register, trigger }) }), ) diff --git a/packages/core/src/session/compaction.ts b/packages/core/src/session/compaction.ts index 7a18f55f38..fe4cd1279d 100644 --- a/packages/core/src/session/compaction.ts +++ b/packages/core/src/session/compaction.ts @@ -1,16 +1,18 @@ export * as SessionCompaction from "./compaction" -import { LLM, LLMClient, LLMError, LLMEvent, Message, type LLMRequest, type Model } from "@opencode-ai/ai" +import { LLM, LLMEvent, Message, type LLMRequest } from "@opencode-ai/ai" +import type { Agent } from "@opencode-ai/schema/agent" +import type { Model } from "@opencode-ai/schema/model" import { SessionError } from "@opencode-ai/schema/session-error" import { Context, Effect, Layer, Stream } from "effect" import { Config } from "../config" import { EventV2 } from "../event" import { makeLocationNode } from "@opencode-ai/util/effect/app-node" -import { llmClient } from "../effect/app-node-platform" +import { Client } from "@opencode-ai/util/client" import { SessionEvent } from "./event" import type { SessionMessage } from "./message" import { SessionModelHeaders } from "./model-headers" -import { Client } from "@opencode-ai/util/client" +import { SessionModelStream } from "./model-stream" import { SessionRunnerModel } from "./runner/model" import { SessionSchema } from "./schema" import { toSessionError } from "./to-session-error" @@ -63,9 +65,7 @@ type Settings = { type Dependencies = { readonly client: string readonly events: EventV2.Interface - readonly llm: { - readonly stream: (request: LLMRequest) => Stream.Stream - } + readonly llm: SessionModelStream.Interface readonly models: SessionRunnerModel.Interface readonly config: Settings } @@ -73,7 +73,9 @@ type Dependencies = { export type AutoInput = { readonly session: SessionSchema.Info readonly messages: readonly SessionMessage.Info[] - readonly model: Model + readonly model: LLMRequest["model"] + readonly agent: Agent.ID + readonly modelRef: Model.Ref readonly cost: ModelV2.Info["cost"] } @@ -81,11 +83,14 @@ export type ManualInput = { readonly session: SessionSchema.Info readonly messages: readonly SessionMessage.Info[] readonly inputID: SessionMessage.ID + readonly agent: Agent.ID } type Plan = { readonly session: SessionSchema.Info - readonly model: Model + readonly model: LLMRequest["model"] + readonly agent: Agent.ID + readonly modelRef: Model.Ref readonly cost: ModelV2.Info["cost"] readonly reason: SessionMessage.Compaction["reason"] readonly prompt: string @@ -257,14 +262,17 @@ const make = (dependencies: Dependencies) => { : Effect.void, ) yield* dependencies.llm - .stream( - LLM.request({ + .stream({ + sessionID: plan.session.id, + agent: plan.agent, + model: plan.modelRef, + request: LLM.request({ model: plan.model, http: { headers: SessionModelHeaders.make(plan.session, dependencies.client) }, messages: [Message.user(plan.prompt)], tools: [], }), - ) + }) .pipe( Stream.runForEach((event) => { if (LLMEvent.is.providerError(event)) @@ -331,6 +339,8 @@ const make = (dependencies: Dependencies) => { session: input.session, model: input.model, cost: input.cost, + agent: input.agent, + modelRef: input.modelRef, reason: "auto", ...content, }) @@ -380,6 +390,8 @@ const make = (dependencies: Dependencies) => { session: input.session, model: resolved.model, cost: resolved.cost, + agent: input.agent, + modelRef: resolved.ref, reason: "manual", inputID: input.inputID, ...content, @@ -396,7 +408,7 @@ export const layer = Layer.effect( Service, Effect.gen(function* () { const events = yield* EventV2.Service - const llm = yield* LLMClient.Service + const llm = yield* SessionModelStream.Service const config = yield* Config.Service const models = yield* SessionRunnerModel.Service const client = yield* Client.Name @@ -407,5 +419,5 @@ export const layer = Layer.effect( export const node = makeLocationNode({ service: Service, layer, - deps: [EventV2.node, llmClient, Config.node, SessionRunnerModel.node, Client.node], + deps: [EventV2.node, SessionModelStream.node, Config.node, SessionRunnerModel.node, Client.node], }) diff --git a/packages/core/src/session/model-stream.ts b/packages/core/src/session/model-stream.ts new file mode 100644 index 0000000000..eb5833e002 --- /dev/null +++ b/packages/core/src/session/model-stream.ts @@ -0,0 +1,50 @@ +export * as SessionModelStream from "./model-stream" + +import { LLMClient, type LLMError, type LLMEvent, type LLMRequest } from "@opencode-ai/ai" +import type { Agent } from "@opencode-ai/schema/agent" +import type { Model } from "@opencode-ai/schema/model" +import type { Session } from "@opencode-ai/schema/session" +import { Context, Effect, Layer, Stream } from "effect" +import { makeLocationNode } from "../effect/app-node" +import { llmClient } from "../effect/app-node-platform" +import { PluginHooks } from "../plugin/hooks" + +export interface Input { + readonly sessionID: Session.ID + readonly agent: Agent.ID + readonly model: Model.Ref + readonly request: LLMRequest +} + +export interface Interface { + readonly stream: (input: Input) => Stream.Stream +} + +export class Service extends Context.Service()("@opencode/v2/SessionModelStream") {} + +const layer = Layer.effect( + Service, + Effect.gen(function* () { + const hooks = yield* PluginHooks.Service + const llm = yield* LLMClient.Service + return Service.of({ + stream: (input) => { + if (!hooks.has("session", "request")) return llm.stream(input.request) + return llm.stream(input.request, { + transformRequest: (request) => + hooks + .trigger("session", "request", { + sessionID: input.sessionID, + agent: input.agent, + model: input.model, + headers: request.headers, + body: request.body, + }) + .pipe(Effect.map((event) => ({ headers: event.headers, body: event.body }))), + }) + }, + }) + }), +) + +export const node = makeLocationNode({ service: Service, layer, deps: [PluginHooks.node, llmClient] }) diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts index 2a8aee08cc..f82a29a2ee 100644 --- a/packages/core/src/session/runner/llm.ts +++ b/packages/core/src/session/runner/llm.ts @@ -1,6 +1,6 @@ export * as SessionRunnerLLM from "./llm" -import { LLMClient, LLMError, LLMEvent, isContextOverflowFailure, type ProviderErrorEvent } from "@opencode-ai/ai" +import { LLMError, LLMEvent, isContextOverflowFailure, type ProviderErrorEvent } from "@opencode-ai/ai" import { SessionError } from "@opencode-ai/schema/session-error" import { Cause, Effect, Exit, Fiber, FiberSet, Layer, Option, Semaphore, Stream } from "effect" import { Database } from "../../database/database" @@ -14,6 +14,7 @@ import { SessionContext } from "../context" import { SessionEvent } from "../event" import { SessionPending } from "../pending" import { SessionModelRequest } from "../model-request" +import { SessionModelStream } from "../model-stream" import { SessionMessage } from "../message" import { SessionSchema } from "../schema" import { SessionStore } from "../store" @@ -22,7 +23,6 @@ import { Service } from "./index" import { createLLMEventPublisher } from "./publish-llm-event" import { Snapshot } from "../../snapshot" import { makeLocationNode } from "@opencode-ai/util/effect/app-node" -import { llmClient } from "../../effect/app-node-platform" import { StepFailedError } from "../error" import { toSessionError } from "../to-session-error" import { SessionRunnerRetry } from "./retry" @@ -32,7 +32,7 @@ const layer = Layer.effect( Service, Effect.gen(function* () { const events = yield* EventV2.Service - const llm = yield* LLMClient.Service + const llm = yield* SessionModelStream.Service const store = yield* SessionStore.Service const context = yield* SessionContext.Service const modelRequests = yield* SessionModelRequest.Service @@ -132,7 +132,9 @@ const layer = Layer.effect( const serialized = (effect: Effect.Effect) => publication.withPermit(effect) const publish = (event: LLMEvent, error?: SessionError.Error) => serialized(publisher.publish(event, error)) let overflowFailure: ProviderErrorEvent | undefined - const providerStream = llm.stream(prepared.request).pipe( + const providerStream = llm + .stream({ sessionID: session.id, agent: agent.id, model: resolved.ref, request: prepared.request }) + .pipe( Stream.runForEach((event) => Effect.gen(function* () { if (overflowFailure || publisher.hasProviderError()) return @@ -493,7 +495,7 @@ export const node = makeLocationNode({ layer, deps: [ EventV2.node, - llmClient, + SessionModelStream.node, SessionContext.node, SessionModelRequest.node, SessionStore.node, diff --git a/packages/core/src/session/title.ts b/packages/core/src/session/title.ts index a506cb1cc1..f5504c2fbc 100644 --- a/packages/core/src/session/title.ts +++ b/packages/core/src/session/title.ts @@ -1,16 +1,16 @@ export * as SessionTitle from "./title" -import { LLM, LLMClient, LLMError, LLMEvent, Message, type LLMRequest } from "@opencode-ai/ai" +import { LLM, LLMEvent, Message } from "@opencode-ai/ai" import { Context, Effect, Layer, Stream } from "effect" import { AgentV2 } from "../agent" import { Database } from "../database/database" import { EventV2 } from "../event" import { makeLocationNode } from "@opencode-ai/util/effect/app-node" import { Client } from "@opencode-ai/util/client" -import { llmClient } from "../effect/app-node-platform" import { SessionEvent } from "./event" import { SessionHistory } from "./history" import { SessionModelHeaders } from "./model-headers" +import { SessionModelStream } from "./model-stream" import { SessionRunnerModel } from "./runner/model" import { SessionSchema } from "./schema" import { SessionUsage } from "./usage" @@ -20,9 +20,7 @@ const MAX_LENGTH = 100 type Dependencies = { readonly client: string readonly events: EventV2.Interface - readonly llm: { - readonly stream: (request: LLMRequest) => Stream.Stream - } + readonly llm: SessionModelStream.Interface readonly agents: AgentV2.Interface readonly models: SessionRunnerModel.Interface } @@ -65,15 +63,18 @@ const make = (dependencies: Dependencies) => { : Effect.void, ) const streamed = yield* dependencies.llm - .stream( - LLM.request({ + .stream({ + sessionID: session.id, + agent: agent.id, + model: resolved.ref, + request: LLM.request({ model: resolved.model, http: { headers: SessionModelHeaders.make(session, dependencies.client) }, system: agent.system, messages: [Message.user(firstUser.text)], tools: [], }), - ) + }) .pipe( Stream.runForEach((event) => { if (LLMEvent.is.providerError(event)) failed = true @@ -108,7 +109,7 @@ export const layer = Layer.effect( Service, Effect.gen(function* () { const events = yield* EventV2.Service - const llm = yield* LLMClient.Service + const llm = yield* SessionModelStream.Service const agents = yield* AgentV2.Service const models = yield* SessionRunnerModel.Service const database = yield* Database.Service @@ -123,5 +124,5 @@ export const layer = Layer.effect( export const node = makeLocationNode({ service: Service, layer, - deps: [EventV2.node, llmClient, AgentV2.node, SessionRunnerModel.node, Database.node, Client.node], + deps: [EventV2.node, SessionModelStream.node, AgentV2.node, SessionRunnerModel.node, Database.node, Client.node], }) diff --git a/packages/core/test/aisdk.test.ts b/packages/core/test/aisdk.test.ts index 0bc16e8f59..7dd668b61d 100644 --- a/packages/core/test/aisdk.test.ts +++ b/packages/core/test/aisdk.test.ts @@ -3,12 +3,17 @@ import { AISDK } from "@opencode-ai/core/aisdk" import { ModelV2 } from "@opencode-ai/core/model" import { ProviderV2 } from "@opencode-ai/core/provider" import { LLM, LLMError, LLMEvent, Message } from "@opencode-ai/ai" -import { LLMClient, RequestExecutor } from "@opencode-ai/ai/route" +import { LLMClient, RequestExecutor, type RequestData } from "@opencode-ai/ai/route" import { expect } from "bun:test" -import { Effect, Layer } from "effect" +import { Effect, Layer, Stream } from "effect" import { testEffect } from "./lib/effect" -const it = testEffect(AISDK.locationLayer) +const it = testEffect( + Layer.merge( + AISDK.locationLayer, + LLMClient.layer.pipe(Layer.provide(Layer.mock(RequestExecutor.Service)({ execute: () => Effect.die("unused") }))), + ), +) const model = (packageName: string, settings: Record = {}) => ModelV2.Info.make({ @@ -99,6 +104,64 @@ it.effect("projects request settings, headers, and body overlays", () => }), ) +it.effect("isolates request transforms across cached AI SDK streams", () => + Effect.gen(function* () { + const aisdk = yield* AISDK.Service + const upstream: Array<{ headers: Headers; body: Record }> = [] + let sdkFetch: typeof fetch | undefined + const input = model("test-sdk") + if (!input.settings) return yield* Effect.die("AI SDK test model settings are missing") + Object.assign(input.settings, { + fetch: async (_input: Parameters[0], init?: RequestInit) => { + upstream.push({ + headers: new Headers(init?.headers), + body: JSON.parse(String(init?.body)), + }) + return new Response() + }, + }) + yield* aisdk.hook.sdk((event) => { + sdkFetch = event.options.fetch + event.sdk = { + languageModel: () => ({ + doStream: async (options: LanguageModelV3CallOptions) => { + if (!sdkFetch) throw new Error("AI SDK fetch was not installed") + await sdkFetch("https://provider.test/model", { + method: "POST", + headers: Object.fromEntries( + Object.entries(options.headers ?? {}).filter( + (entry): entry is [string, string] => entry[1] !== undefined, + ), + ), + body: JSON.stringify({ model: "api-model", remove: true }), + }) + return { stream: new ReadableStream({ start: (controller) => controller.close() }) } + }, + }), + } + }) + const resolved = yield* aisdk.model(input) + const llm = yield* LLMClient.Service + const run = (session: string) => + llm + .stream(LLM.request({ model: resolved, prompt: session }), { + transformRequest: (request) => + Effect.sync(() => { + const body: RequestData["body"] = { ...request.body, session } + delete body.remove + return { headers: { ...request.headers, "x-session": session }, body } + }), + }) + .pipe(Stream.runDrain) + + yield* Effect.all([run("one"), run("two")], { concurrency: "unbounded" }) + + expect(upstream.map((request) => request.headers.get("x-session")).toSorted()).toEqual(["one", "two"]) + expect(upstream.map((request) => request.body.session).toSorted()).toEqual(["one", "two"]) + expect(upstream.every((request) => !("remove" in request.body))).toBe(true) + }), +) + it.effect("maps pro reasoning bodies to AI SDK provider options", () => Effect.gen(function* () { const aisdk = yield* AISDK.Service diff --git a/packages/core/test/session-compaction.test.ts b/packages/core/test/session-compaction.test.ts index c58433b9cd..55e7542f99 100644 --- a/packages/core/test/session-compaction.test.ts +++ b/packages/core/test/session-compaction.test.ts @@ -1,6 +1,7 @@ import { expect, test } from "bun:test" import { LLMClient, LLMEvent, Model, type LLMRequest } from "@opencode-ai/ai" import { OpenAIChat } from "@opencode-ai/ai/protocols" +import { AgentV2 } from "@opencode-ai/core/agent" import { Config } from "@opencode-ai/core/config" import { Database } from "@opencode-ai/core/database/database" import { AppNodeBuilder } from "@opencode-ai/core/effect/app-node-builder" @@ -181,6 +182,7 @@ it.effect("manual compaction summarizes short context instead of no-op", () => expect( yield* compaction.compactManual({ session, + agent: AgentV2.ID.make("build"), messages: [userMessage], inputID: SessionMessage.ID.make("msg_manual_compaction"), }), diff --git a/packages/core/test/session-model-stream.test.ts b/packages/core/test/session-model-stream.test.ts new file mode 100644 index 0000000000..7ae22e7400 --- /dev/null +++ b/packages/core/test/session-model-stream.test.ts @@ -0,0 +1,73 @@ +import { expect } from "bun:test" +import { LLM, LLMClient } from "@opencode-ai/ai" +import { OpenAIChat } from "@opencode-ai/ai/protocols" +import type { RequestData } from "@opencode-ai/ai/route" +import { Agent } from "@opencode-ai/schema/agent" +import { Model } from "@opencode-ai/schema/model" +import { Provider } from "@opencode-ai/schema/provider" +import { Session } from "@opencode-ai/schema/session" +import { Effect, Layer, Stream } from "effect" +import { AppNodeBuilder } from "@opencode-ai/core/effect/app-node-builder" +import { llmClient } from "@opencode-ai/core/effect/app-node-platform" +import { LayerNode } from "@opencode-ai/core/effect/layer-node" +import { PluginHooks } from "@opencode-ai/core/plugin/hooks" +import { SessionModelStream } from "@opencode-ai/core/session/model-stream" +import { testEffect } from "./lib/effect" + +const prepared: RequestData[] = [] +const client = Layer.mock(LLMClient.Service)({ + prepare: () => Effect.die("unused"), + generate: () => Effect.die("unused"), + stream: (_request, options) => + Stream.unwrap( + options?.transformRequest + ? options.transformRequest({ headers: { original: "true" }, body: { remove: true } }).pipe( + Effect.map((request) => { + prepared.push(request) + return Stream.empty + }), + ) + : Effect.die("request transform was not provided"), + ), +}) +const it = testEffect( + AppNodeBuilder.build(LayerNode.group([PluginHooks.node, SessionModelStream.node]), [[llmClient, client]]), +) + +it.effect("forwards session identity and applies request hook mutations", () => + Effect.gen(function* () { + prepared.length = 0 + const hooks = yield* PluginHooks.Service + const stream = yield* SessionModelStream.Service + const sessionID = Session.ID.make("ses_model_stream") + const agent = Agent.ID.make("build") + const model = Model.Ref.make({ + providerID: Provider.ID.make("test"), + id: Model.ID.make("catalog-model"), + }) + yield* hooks.register("session", "request", (event) => + Effect.sync(() => { + expect(event.sessionID).toBe(sessionID) + expect(event.agent).toBe(agent) + expect(event.model).toEqual(model) + event.headers["x-plugin"] = "enabled" + delete event.body.remove + }), + ) + yield* stream + .stream({ + sessionID, + agent, + model, + request: LLM.request({ + model: OpenAIChat.route + .with({ endpoint: { baseURL: "https://api.openai.test/v1/" } }) + .model({ id: "api-model" }), + prompt: "Hello", + }), + }) + .pipe(Stream.runDrain) + + expect(prepared).toEqual([{ headers: { original: "true", "x-plugin": "enabled" }, body: {} }]) + }), +)