feat(plugin): wire session request hook

This commit is contained in:
Aiden Cline 2026-07-17 14:32:08 -05:00 committed by 𝓛𝓲𝓽𝓽𝓵𝓮 𝓕𝓻𝓪𝓷𝓴
commit 7edaa05869
17 changed files with 447 additions and 73 deletions

View file

@ -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<LanguageModelV3Message, { role: "user" }>["content"]
type AssistantContent = Extract<LanguageModelV3Message, { role: "assistant" }>["content"]
type ToolResultContent = Extract<AssistantContent[number], { type: "tool-result" }>
type TransformRequest = NonNullable<CallOptions["transformRequest"]>
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<TransformRequest>) {
const projected = mapBodyToProviderOptions(model, pkg)
const options: Record<string, any> = {
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> | void)[] = []
const languages = new Map<string, LanguageModelV3>()
const sdks = new Map<string, SDK>()
const requests = new AsyncLocalStorage<TransformRequest>()
const functionIDs = new WeakMap<object, number>()
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<TransformRequest>,
) {
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?: TransformRequest,
) {
const state = { step: 0, toolNames: {} as Record<string, string> }
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) =>

View file

@ -16,6 +16,10 @@ export interface Domains {
type Callback<Event> = (event: Event) => Effect.Effect<void>
export interface Interface {
readonly has: <Domain extends keyof Domains, Name extends keyof Domains[Domain]>(
domain: Domain,
name: Name,
) => boolean
readonly register: <Domain extends keyof Domains, Name extends keyof Domains[Domain]>(
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 })
}),
)

View file

@ -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<LLMEvent, LLMError>
}
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],
})

View file

@ -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<LLMEvent, LLMError>
}
export class Service extends Context.Service<Service, Interface>()("@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] })

View file

@ -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 = <A, E, R>(effect: Effect.Effect<A, E, R>) => 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,

View file

@ -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<LLMEvent, LLMError>
}
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],
})

View file

@ -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<string, unknown> = {}) =>
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<string, unknown> }> = []
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<typeof fetch>[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

View file

@ -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"),
}),

View file

@ -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: {} }])
}),
)