Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion packages/core/src/aisdk.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ import {
import { Auth, Endpoint, RequestExecutor, type AnyRoute } 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 { makeParser } from "effect/unstable/encoding/Sse"
import type { ID, Info } from "./model.js"
import { Provider } from "./provider.js"
import { State } from "./state.js"
Expand Down Expand Up @@ -65,15 +66,23 @@ function wrapSSE(res: Response, ms: number, ctl: AbortController) {
if (!res.headers.get("content-type")?.includes("text/event-stream")) return res

const reader = res.body.getReader()
const decoder = new TextDecoder()
let deadline: number | undefined
const parser = makeParser((event) => {
if (event._tag === "Event") deadline = Date.now() + ms
})
const body = new ReadableStream<Uint8Array>({
async pull(ctrl) {
const expires = deadline ?? Date.now() + ms
deadline = expires
const part = await new Promise<Awaited<ReturnType<typeof reader.read>>>((resolve, reject) => {
const remaining = Math.max(0, expires - Date.now())
const id = setTimeout(() => {
const err = new Error("SSE read timed out")
ctl.abort(err)
void reader.cancel(err)
reject(err)
}, ms)
}, remaining)

reader.read().then(
(part) => {
Expand All @@ -92,6 +101,7 @@ function wrapSSE(res: Response, ms: number, ctl: AbortController) {
return
}

parser.feed(decoder.decode(part.value, { stream: true }))
ctrl.enqueue(part.value)
},
async cancel(reason) {
Expand Down
58 changes: 58 additions & 0 deletions packages/core/test/aisdk.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { APICallError } from "@ai-sdk/provider"
import type { LanguageModelV3, LanguageModelV3StreamPart } from "@ai-sdk/provider"
import { createMistral } from "@ai-sdk/mistral"
import { createOpenAICompatible } from "@ai-sdk/openai-compatible"
import { AISDK } from "@opencode-ai/core/aisdk"
import { SessionRunnerRetry } from "@opencode-ai/core/session/runner/retry"
import { toSessionError } from "@opencode-ai/core/session/to-session-error"
Expand Down Expand Up @@ -412,6 +413,63 @@ it.effect("moves a tool image through the real Mistral provider as a user messag
}),
)

it.effect("does not treat SSE comment heartbeats as model progress", () =>
Effect.gen(function* () {
const aisdk = yield* AISDK.Service
const encoder = new TextEncoder()
let heartbeat: ReturnType<typeof setInterval> | undefined
const customFetch = Object.assign(
async () =>
new Response(
new ReadableStream({
start(controller) {
controller.enqueue(
encoder.encode(
'data: {"id":"response-1","object":"chat.completion.chunk","created":0,"model":"api-model","choices":[{"index":0,"delta":{"content":"partial"},"finish_reason":null}]}\n\n',
),
)
heartbeat = setInterval(() => controller.enqueue(encoder.encode(": keepalive\n\n")), 5)
},
cancel() {
if (heartbeat) clearInterval(heartbeat)
},
}),
{ headers: { "content-type": "text/event-stream" } },
),
{ preconnect: fetch.preconnect },
)
yield* aisdk.hook.sdk((event) => {
event.sdk = createOpenAICompatible({
...event.options,
name: String(event.options.name),
baseURL: String(event.options.baseURL),
})
})
const resolved = yield* aisdk.model(
model("@ai-sdk/openai-compatible", {
apiKey: "test",
baseURL: "https://example.test/v1",
chunkTimeout: 25,
fetch: customFetch,
}),
)
const result = yield* LLMClient.generate(LLM.request({ model: resolved, prompt: "Hello" })).pipe(
Effect.provide(client),
Effect.result,
Effect.ensuring(
Effect.sync(() => {
if (heartbeat) clearInterval(heartbeat)
}),
),
)

expect(result).toMatchObject({
_tag: "Failure",
failure: { reason: { message: expect.stringContaining("SSE read timed out") } },
})
}),
)

it.effect("emits malformed AI SDK tool input without executing it", () =>
Effect.gen(function* () {
const aisdk = yield* AISDK.Service
Expand Down
Loading