Skip to content
Open
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
1 change: 1 addition & 0 deletions packages/ai/src/provider-package.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ export interface Settings extends Readonly<Record<string, unknown>> {
readonly baseURL?: string
readonly headers?: Readonly<Record<string, string>>
readonly body?: Readonly<Record<string, unknown>>
readonly chunkTimeout?: number
}

export interface Definition<
Expand Down
10 changes: 8 additions & 2 deletions packages/ai/src/providers/azure.ts
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,6 @@ export type Settings = ProviderPackage.Settings &
readonly useDeploymentBasedUrls?: boolean
readonly providerOptions?: OpenAIProviderOptionsInput
}

const resourceBaseURL = (resourceName: string) => `https://${resourceName.trim()}.openai.azure.com/openai`

const responsesRoute = OpenAIResponses.route.with({
Expand Down Expand Up @@ -152,11 +151,18 @@ export const provider = {
}

const config = (settings: Settings): Config => {
const http =
settings.body === undefined && settings.chunkTimeout === undefined
? undefined
: {
body: settings.body === undefined ? undefined : { ...settings.body },
chunkTimeout: settings.chunkTimeout,
}
const common = {
apiKey: settings.apiKey,
apiVersion: settings.apiVersion,
headers: settings.headers === undefined ? undefined : { ...settings.headers },
http: settings.body === undefined ? undefined : { body: { ...settings.body } },
http,
providerOptions: settings.providerOptions,
queryParams: settings.queryParams === undefined ? undefined : { ...settings.queryParams },
useDeploymentBasedUrls: settings.useDeploymentBasedUrls,
Expand Down
33 changes: 29 additions & 4 deletions packages/ai/src/route/transport/http.ts
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
import { Effect } from "effect"
import { Duration, Effect, Stream } from "effect"
import { Headers, HttpClientRequest } from "effect/unstable/http"
import { Auth } from "../auth.js"
import { render as renderEndpoint } from "../endpoint.js"
import { Framing } from "../framing.js"
import type { HttpMiddleware, Transport, TransportPrepareInput } from "./index.js"
import * as ProviderShared from "../../protocols/shared.js"
import { mergeJsonRecords, type LLMRequest } from "../../schema/index.js"
import { AIError, mergeJsonRecords, TransportError, type LLMRequest } from "../../schema/index.js"
import { RequestExecutor } from "../executor.js"

export type JsonRequestInput<Body> = TransportPrepareInput<Body>
Expand Down Expand Up @@ -87,11 +87,36 @@ export const httpJson = <Body, Frame>(input: HttpJsonInput<Body, Frame>): HttpJs
middleware: prepareInput.middleware,
}
}),
execute: (prepared, _request, runtime) =>
execute: (prepared, request, runtime) =>
Effect.gen(function* () {
const response = yield* runtime.http.execute(prepared.request, prepared.middleware)
const chunkTimeout = request.http?.chunkTimeout
const bytes = RequestExecutor.responseStream(response)
const guarded =
chunkTimeout === undefined || chunkTimeout <= 0
? bytes
: bytes.pipe(
// A stalled provider stream that sends nothing is otherwise bounded only by the
// OS socket; abort once no chunk arrives within the configured window.
Stream.timeoutOrElse({
duration: Duration.millis(chunkTimeout),
orElse: () =>
Stream.fail(
new AIError({
reason: new TransportError({
message: `No data received from ${response.request.url} within the ${chunkTimeout}ms chunkTimeout; the stream may be stalled`,
transport: "http",
operation: "read",
code: "chunk-timeout",
url: response.request.url,
phase: "receive",
}),
}),
),
}),
)
return {
frames: prepared.framing.frame(RequestExecutor.responseStream(response)),
frames: prepared.framing.frame(guarded),
http: RequestExecutor.responseHttp(response),
body: prepared.framing.body,
}
Expand Down
8 changes: 6 additions & 2 deletions packages/ai/src/schema/options.ts
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,9 @@ export class HttpOptions extends Schema.Class<HttpOptions>("AI.HttpOptions")({
body: Schema.optional(JsonSchema),
headers: Schema.optional(Schema.Record(Schema.String, Schema.String)),
query: Schema.optional(Schema.Record(Schema.String, Schema.String)),
chunkTimeout: Schema.optional(Schema.Number).annotate({
description: "Abort the stream when no data chunk arrives within this many milliseconds.",
}),
}) {}

export namespace HttpOptions {
Expand All @@ -60,8 +63,9 @@ export const mergeHttpOptions = (...items: ReadonlyArray<HttpOptions | undefined
const body = mergeJsonRecords(...items.map((item) => item?.body))
const headers = mergeStringRecords(...items.map((item) => item?.headers))
const query = mergeStringRecords(...items.map((item) => item?.query))
if (!body && !headers && !query) return undefined
return new HttpOptions({ body, headers, query })
const chunkTimeout = items.reduceRight((found, item) => found ?? item?.chunkTimeout, undefined as number | undefined)
if (!body && !headers && !query && chunkTimeout === undefined) return undefined
return new HttpOptions({ body, headers, query, chunkTimeout })
}

export class GenerationOptions extends Schema.Class<GenerationOptions>("LLM.GenerationOptions")({
Expand Down
55 changes: 55 additions & 0 deletions packages/ai/test/chunk-timeout.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
import { describe, expect, test } from "bun:test"
import { Effect } from "effect"
import { AIError, LLM } from "../src/index.js"
import { HttpOptions, mergeHttpOptions } from "../src/schema/options.js"
import * as OpenAIChat from "../src/protocols/openai-chat.js"
import { LLMClient } from "../src/route.js"
import { dynamicResponse } from "./lib/http.js"
import { it } from "./lib/effect.js"

const headers = { "content-type": "text/event-stream" }

// Sends one SSE frame, then holds the connection open without further data.
// Only a configured chunkTimeout can terminate the read.
const stalledResponse = dynamicResponse((input) =>
Effect.sync(() => {
const encoder = new TextEncoder()
let sent = false
const stream = new ReadableStream({
pull(controller) {
if (!sent) {
sent = true
controller.enqueue(encoder.encode('data: {"choices":[]}\n\n'))
}
},
})
return input.respond(stream, { headers })
}),
)

describe("chunkTimeout", () => {
it.live("aborts a stalled HTTP SSE stream with a typed transport error", () =>
Effect.gen(function* () {
const route = OpenAIChat.route.with({ http: { chunkTimeout: 100 } })
const request = LLM.request({ model: route.model({ id: "test" }), prompt: "Hello" })
const error = yield* LLMClient.generate(request).pipe(Effect.provide(stalledResponse), Effect.flip)

expect(error).toBeInstanceOf(AIError)
expect(error.reason).toMatchObject({
_tag: "Transport",
operation: "read",
phase: "receive",
code: "chunk-timeout",
})
expect(error.message).toContain("chunkTimeout")
}),
)

test("mergeHttpOptions keeps the rightmost configured chunkTimeout", () => {
const merged = mergeHttpOptions(new HttpOptions({ chunkTimeout: 1000 }), new HttpOptions({ chunkTimeout: 250 }))
expect(merged?.chunkTimeout).toBe(250)
expect(mergeHttpOptions(new HttpOptions({ chunkTimeout: 250 }))?.chunkTimeout).toBe(250)
expect(mergeHttpOptions(new HttpOptions({ body: { a: 1 } }))?.chunkTimeout).toBeUndefined()
expect(mergeHttpOptions(new HttpOptions({ chunkTimeout: 250 }), new HttpOptions())?.chunkTimeout).toBe(250)
})
})
6 changes: 4 additions & 2 deletions services/www/src/docs/content/providers.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ The `providers` object is keyed by the provider ID used in model references, suc
| `env` | Ordered environment variable names that can provide the credential. |
| `package` | Runtime provider package. |
| `canonical` | Built-in provider ID whose catalog defaults this provider inherits. |
| `settings` | JSON options passed to the runtime package. |
| `settings` | JSON options passed to the runtime package, such as `baseURL` and `chunkTimeout`. |
| `headers` | String-valued HTTP headers added to requests. |
| `body` | JSON fields merged into request bodies. |
| `models` | Models to add or override, keyed by the OpenCode model ID. |
Expand All @@ -86,7 +86,9 @@ models, and connection still apply.
}
```

`settings` is package-specific. A field only has an effect when the selected package supports it.
`settings` is package-specific. A field only has an effect when the selected package supports it. Packages that
support `chunkTimeout` abort a stream when no data chunk arrives within the given number of milliseconds, guarding
against a stalled provider.

## Requests

Expand Down
Loading