openAgent

1export * as AISDK from "./aisdk"
2 
3import { makeLocationNode } from "./effect/app-node"
4import type { LanguageModelV3 } from "@ai-sdk/provider"
5import { Cause, Context, Effect, Layer, Schema, Scope } from "effect"
6import { ModelV2 } from "./model"
7import { ProviderV2 } from "./provider"
8import { State } from "./state"
9 
10type SDK = any
11 
12export interface SDKEvent {
13  readonly model: ModelV2.Info
14  readonly package: string
15  readonly options: Record<string, any>
16  sdk?: SDK
17}
18 
19export interface LanguageEvent {
20  readonly model: ModelV2.Info
21  readonly sdk: SDK
22  readonly options: Record<string, any>
23  language?: LanguageModelV3
24}
25 
26function wrapSSE(res: Response, ms: number, ctl: AbortController) {
27  if (typeof ms !== "number" || ms <= 0) return res
28  if (!res.body) return res
29  if (!res.headers.get("content-type")?.includes("text/event-stream")) return res
30 
31  const reader = res.body.getReader()
32  const body = new ReadableStream<Uint8Array>({
33    async pull(ctrl) {
34      const part = await new Promise<Awaited<ReturnType<typeof reader.read>>>((resolve, reject) => {
35        const id = setTimeout(() => {
36          const err = new Error("SSE read timed out")
37          ctl.abort(err)
38          void reader.cancel(err)
39          reject(err)
40        }, ms)
41 
42        reader.read().then(
43          (part) => {
44            clearTimeout(id)
45            resolve(part)
46          },
47          (err) => {
48            clearTimeout(id)
49            reject(err)
50          },
51        )
52      })
53 
54      if (part.done) {
55        ctrl.close()
56        return
57      }
58 
59      ctrl.enqueue(part.value)
60    },
61    async cancel(reason) {
62      ctl.abort(reason)
63      await reader.cancel(reason)
64    },
65  })
66 
67  return new Response(body, {
68    headers: new Headers(res.headers),
69    status: res.status,
70    statusText: res.statusText,
71  })
72}
73 
74function prepareOptions(model: ModelV2.Info, pkg: string) {
75  const options: Record<string, any> = {
76    name: model.providerID,
77    ...(model.api.type === "aisdk" ? (model.api.settings ?? {}) : {}),
78    ...model.request.body,
79  }
80  if (model.api.type === "aisdk" && model.api.url) options.baseURL = model.api.url
81 
82  const customFetch = options.fetch
83  const chunkTimeout = options.chunkTimeout
84  delete options.chunkTimeout
85  options.fetch = async (input: Parameters<typeof fetch>[0], init?: RequestInit) => {
86    const opts = { ...(init ?? {}) }
87    const signals = [
88      opts.signal,
89      typeof chunkTimeout === "number" && chunkTimeout > 0 ? new AbortController() : undefined,
90      options.timeout !== undefined && options.timeout !== null && options.timeout !== false
91        ? AbortSignal.timeout(options.timeout)
92        : undefined,
93    ].filter((item): item is AbortSignal | AbortController => Boolean(item))
94    const chunkAbortCtl = signals.find((item): item is AbortController => item instanceof AbortController)
95    const abortSignals = signals.map((item) => (item instanceof AbortController ? item.signal : item))
96    if (abortSignals.length === 1) opts.signal = abortSignals[0]
97    if (abortSignals.length > 1) opts.signal = AbortSignal.any(abortSignals)
98 
99    if (
100      (pkg === "@ai-sdk/openai" || pkg === "@ai-sdk/azure" || pkg === "@ai-sdk/amazon-bedrock/mantle") &&
101      opts.body &&
102      opts.method === "POST"
103    ) {
104      const body = JSON.parse(opts.body as string)
105      if (body.store !== true && Array.isArray(body.input)) {
106        for (const item of body.input) {
107          if ("id" in item) delete item.id
108        }
109        opts.body = JSON.stringify(body)
110      }
111    }
112 
113    const res = await (typeof customFetch === "function" ? customFetch : fetch)(input, {
114      ...opts,
115      timeout: false,
116    })
117    if (!chunkAbortCtl || typeof chunkTimeout !== "number") return res
118    return wrapSSE(res, chunkTimeout, chunkAbortCtl)
119  }
120 
121  return options
122}
123 
124export class InitError extends Schema.TaggedErrorClass<InitError>()("AISDK.InitError", {
125  providerID: ProviderV2.ID,
126  cause: Schema.Defect(),
127}) {}
128 
129function initError(providerID: ProviderV2.ID) {
130  return Effect.catchCause((cause) => Effect.fail(new InitError({ providerID, cause: Cause.squash(cause) })))
131}
132 
133export interface Interface {
134  readonly hook: {
135    readonly sdk: (
136      callback: (event: SDKEvent) => Effect.Effect<void> | void,
137    ) => Effect.Effect<State.Registration, never, Scope.Scope>
138    readonly language: (
139      callback: (event: LanguageEvent) => Effect.Effect<void> | void,
140    ) => Effect.Effect<State.Registration, never, Scope.Scope>
141  }
142  readonly runSDK: (event: SDKEvent) => Effect.Effect<SDKEvent>
143  readonly runLanguage: (event: LanguageEvent) => Effect.Effect<LanguageEvent>
144  readonly language: (model: ModelV2.Info) => Effect.Effect<LanguageModelV3, InitError>
145}
146 
147export class Service extends Context.Service<Service, Interface>()("@opencode/v2/AISDK") {}
148 
149export const locationLayer = Layer.effect(
150  Service,
151  Effect.gen(function* () {
152    let sdkHooks: ((event: SDKEvent) => Effect.Effect<void> | void)[] = []
153    let languageHooks: ((event: LanguageEvent) => Effect.Effect<void> | void)[] = []
154    const languages = new Map<string, LanguageModelV3>()
155    const sdks = new Map<string, SDK>()
156 
157    const register = <Event>(
158      hooks: () => ((event: Event) => Effect.Effect<void> | void)[],
159      update: (hooks: ((event: Event) => Effect.Effect<void> | void)[]) => void,
160    ) =>
161      Effect.fn("AISDK.hook")(function* (callback: (event: Event) => Effect.Effect<void> | void) {
162        const scope = yield* Scope.Scope
163        let active = true
164        update([...hooks(), callback])
165        const dispose = Effect.sync(() => {
166          if (!active) return
167          active = false
168          update(hooks().filter((item) => item !== callback))
169        })
170        yield* Scope.addFinalizer(scope, dispose)
171        return { dispose }
172      })
173 
174    const run = Effect.fnUntraced(function* <Event>(
175      hooks: readonly ((event: Event) => Effect.Effect<void> | void)[],
176      event: Event,
177    ) {
178      for (const hook of hooks) {
179        const result = hook(event)
180        if (Effect.isEffect(result)) yield* result
181      }
182      return event
183    })
184 
185    const service = Service.of({
186      hook: {
187        sdk: register(
188          () => sdkHooks,
189          (next) => (sdkHooks = next),
190        ),
191        language: register(
192          () => languageHooks,
193          (next) => (languageHooks = next),
194        ),
195      },
196      runSDK: (event) => run(sdkHooks, event),
197      runLanguage: (event) => run(languageHooks, event),
198      language: Effect.fn("AISDK.language")(function* (model) {
199        const key = `${model.providerID}/${model.id}/${model.request.variant ?? "default"}`
200        const existing = languages.get(key)
201        if (existing) return existing
202        if (model.api.type !== "aisdk")
203          return yield* new InitError({
204            providerID: model.providerID,
205            cause: new Error(`Unsupported api ${model.api.type}`),
206          })
207 
208        const options = prepareOptions(model, model.api.package)
209        const sdkKey = JSON.stringify({
210          providerID: model.providerID,
211          api: model.api,
212          options,
213        })
214        const sdk =
215          sdks.get(sdkKey) ??
216          (yield* service.runSDK({ model, package: model.api.package, options }).pipe(initError(model.providerID))).sdk
217        if (!sdk)
218          return yield* new InitError({
219            providerID: model.providerID,
220            cause: new Error("No AISDK provider plugin returned an SDK"),
221          })
222        sdks.set(sdkKey, sdk)
223        const result = yield* service.runLanguage({ model, sdk, options }).pipe(initError(model.providerID))
224        const language = yield* Effect.sync(() => result.language ?? sdk.languageModel(model.api.id)).pipe(
225          initError(model.providerID),
226        )
227        languages.set(key, language)
228        return language
229      }),
230    })
231    return service
232  }),
233)
234 
235export const node = makeLocationNode({ service: Service, layer: locationLayer, deps: [] })
236