From 6805d8a624d2d77bc1c69dbe1c568228914447d5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sat, 23 May 2026 20:04:28 -0400 Subject: [PATCH] Add provider streaming and cancel contracts --- .../src/main/services/provider-adapters.ts | 181 ++++++++--------- .../src/main/services/provider-http.ts | 182 ++++++++++++++++++ .../src/main/services/provider-streaming.ts | 169 ++++++++++++++++ .../src/main/services/sdlc-service.ts | 4 +- .../tests/contract/provider-contracts.test.ts | 2 + .../tests/unit/provider-streaming.test.ts | 106 ++++++++++ 6 files changed, 557 insertions(+), 87 deletions(-) create mode 100644 AgentCompany-New/src/main/services/provider-http.ts create mode 100644 AgentCompany-New/src/main/services/provider-streaming.ts create mode 100644 AgentCompany-New/tests/unit/provider-streaming.test.ts diff --git a/AgentCompany-New/src/main/services/provider-adapters.ts b/AgentCompany-New/src/main/services/provider-adapters.ts index 0ce92470..90fb182a 100644 --- a/AgentCompany-New/src/main/services/provider-adapters.ts +++ b/AgentCompany-New/src/main/services/provider-adapters.ts @@ -9,7 +9,18 @@ import { type ProviderRefreshHandler, type ProviderRevokeHandler, } from "./provider-auth-lifecycle"; -import { parseProviderRateLimitHeaders } from "./provider-rate-limit-registry"; +import { getJson, postEventStream, postJson, type JsonRecord } from "./provider-http"; +import { + cancelUnsupported, + mapAnthropicMessageStream, + mapOpenAiChatCompletionStream, + mapOpenAiResponseStream, + streamUnsupported, + type ProviderCancelResult, + type ProviderStreamResult, +} from "./provider-streaming"; + +export { ProviderHttpError, providerRateLimitStateFromError } from "./provider-http"; export interface ProviderExecutionInput { authProfile: ProviderAuthProfile; @@ -24,6 +35,13 @@ export interface ProviderExecutionResult { rateLimitState: ProviderObservedRateLimitState | null; } +export type ProviderStreamInput = ProviderExecutionInput; + +export interface ProviderCancelInput { + authProfile: ProviderAuthProfile; + responseId: string; +} + export interface ProviderHealthInput { authProfile: ProviderAuthProfile; model: string | null; @@ -57,6 +75,8 @@ export interface ProviderAdapter { listModels(ctx: ProviderContext): Promise; getRateLimitState(ctx: ProviderContext): Promise; checkHealth(input: ProviderHealthInput): Promise; + streamResponse(input: ProviderStreamInput): Promise; + cancelResponse(input: ProviderCancelInput): Promise; createResponse?: (input: ProviderExecutionInput) => Promise; } @@ -70,29 +90,6 @@ export type ProviderAuthProfile = | { type: "aws_bedrock"; region: string; accessKeyId: string; secretAccessKey: string; sessionToken: string | null } | { type: "vertex_ai"; projectId: string; region: string; serviceAccountJson: string }; -type JsonRecord = Record; - -const PROVIDER_REQUEST_TIMEOUT_MS = 60_000; - -interface ProviderJsonResponse { - data: JsonRecord; - rateLimitState: ProviderObservedRateLimitState | null; -} - -export class ProviderHttpError extends Error { - constructor( - message: string, - readonly rateLimitState: ProviderObservedRateLimitState | null, - ) { - super(message); - this.name = "ProviderHttpError"; - } -} - -export function providerRateLimitStateFromError(error: unknown): ProviderObservedRateLimitState | null { - return error instanceof ProviderHttpError ? error.rateLimitState : null; -} - export const providerAdapters: ProviderAdapterRegistry = { openai: { providerKey: "openai", @@ -104,6 +101,26 @@ export const providerAdapters: ProviderAdapterRegistry = { listModels: () => listStaticModels("openai"), getRateLimitState: providerRateLimitUnknown, checkHealth: (input) => checkOpenAiHealth(input.authProfile, "https://api.openai.com/v1/models"), + async streamResponse(input) { + const profile = requireApiKeyProfile(input.authProfile, "OpenAI"); + const response = await postEventStream("https://api.openai.com/v1/responses", { + authorization: `Bearer ${profile.apiKey}`, + }, { + model: input.model, + input: input.prompt, + stream: true, + ...(input.instructions ? { instructions: input.instructions } : {}), + }); + return { events: mapOpenAiResponseStream(response.messages), rateLimitState: response.rateLimitState }; + }, + async cancelResponse(input) { + const profile = requireApiKeyProfile(input.authProfile, "OpenAI"); + const responseId = requireResponseId(input.responseId); + const response = await postJson(`https://api.openai.com/v1/responses/${encodeURIComponent(responseId)}/cancel`, { + authorization: `Bearer ${profile.apiKey}`, + }, {}); + return { responseId, rawResponseJson: JSON.stringify(response.data), rateLimitState: response.rateLimitState }; + }, async createResponse(input) { const profile = requireApiKeyProfile(input.authProfile, "OpenAI"); const response = await postJson("https://api.openai.com/v1/responses", { @@ -135,6 +152,21 @@ export const providerAdapters: ProviderAdapterRegistry = { }); return healthy(response.rateLimitState); }, + async streamResponse(input) { + const profile = requireApiKeyProfile(input.authProfile, "Anthropic"); + const response = await postEventStream("https://api.anthropic.com/v1/messages", { + "x-api-key": profile.apiKey, + "anthropic-version": "2023-06-01", + }, { + model: input.model, + max_tokens: 4096, + stream: true, + ...(input.instructions ? { system: input.instructions } : {}), + messages: [{ role: "user", content: input.prompt }], + }); + return { events: mapAnthropicMessageStream(response.messages), rateLimitState: response.rateLimitState }; + }, + cancelResponse: cancelUnsupported("Anthropic"), async createResponse(input) { const profile = requireApiKeyProfile(input.authProfile, "Anthropic"); const response = await postJson("https://api.anthropic.com/v1/messages", { @@ -175,6 +207,8 @@ export const providerAdapters: ProviderAdapterRegistry = { parseAuthProfile: parseAzureProfile, listModels: () => listStaticModels("azure-openai"), getRateLimitState: providerRateLimitUnknown, + streamResponse: streamUnsupported("Azure OpenAI"), + cancelResponse: cancelUnsupported("Azure OpenAI"), async checkHealth(input) { const profile = requireAzureProfile(input.authProfile); const response = await getJson(`${profile.baseUrl}/openai/deployments?api-version=2024-10-21`, { "api-key": profile.apiKey }); @@ -190,6 +224,8 @@ export const providerAdapters: ProviderAdapterRegistry = { parseAuthProfile: parseBedrockProfile, listModels: () => listStaticModels("aws-bedrock"), getRateLimitState: providerRateLimitUnknown, + streamResponse: streamUnsupported("AWS Bedrock"), + cancelResponse: cancelUnsupported("AWS Bedrock"), checkHealth: async (input) => { requireBedrockProfile(input.authProfile); return healthy(null); @@ -204,6 +240,8 @@ export const providerAdapters: ProviderAdapterRegistry = { parseAuthProfile: parseVertexProfile, listModels: () => listStaticModels("vertex-ai"), getRateLimitState: providerRateLimitUnknown, + streamResponse: streamUnsupported("Vertex AI"), + cancelResponse: cancelUnsupported("Vertex AI"), checkHealth: async (input) => { requireVertexProfile(input.authProfile); return healthy(null); @@ -211,64 +249,6 @@ export const providerAdapters: ProviderAdapterRegistry = { }, }; -async function postJson(endpoint: string, headers: Record, body: JsonRecord): Promise { - return requestJson(endpoint, { - method: "POST", - headers: { - "content-type": "application/json", - ...headers, - }, - body: JSON.stringify(body), - }); -} - -async function getJson(endpoint: string, headers: Record): Promise { - return requestJson(endpoint, { method: "GET", headers }); -} - -async function requestJson(endpoint: string, init: RequestInit): Promise { - const controller = new AbortController(); - const timeout = setTimeout(() => controller.abort(), PROVIDER_REQUEST_TIMEOUT_MS); - let response: Response; - try { - response = await fetch(endpoint, { ...init, signal: controller.signal }); - } finally { - clearTimeout(timeout); - } - const rateLimitState = parseProviderRateLimitHeaders(response.headers); - const raw = await response.text(); - const data = parseProviderJson(raw, response.ok); - if (!response.ok) { - const message = readProviderError(data) ?? (raw.trim() || `Provider request failed with HTTP ${response.status}.`); - throw new ProviderHttpError(message, rateLimitState); - } - return { data, rateLimitState }; -} - -function parseProviderJson(raw: string, requireJson: boolean): JsonRecord { - if (!raw.trim()) return {}; - let value: unknown; - try { - value = JSON.parse(raw) as unknown; - } catch (error) { - if (!requireJson) return {}; - throw error; - } - if (!isRecord(value)) throw new Error("Provider returned invalid JSON."); - return value; -} - -function readProviderError(data: JsonRecord): string | null { - const error = data.error; - if (typeof error === "string") return error; - if (isRecord(error) && typeof error.message === "string") return error.message; - return null; -} - -function isRecord(value: unknown): value is JsonRecord { - return typeof value === "object" && value !== null && !Array.isArray(value); -} - function collectText(value: unknown): string[] { if (typeof value === "string") return [value]; if (Array.isArray(value)) return value.flatMap((item) => collectText(item)); @@ -277,6 +257,10 @@ function collectText(value: unknown): string[] { return direct.length > 0 ? direct : Object.values(value).flatMap((item) => collectText(item)); } +function isRecord(value: unknown): value is JsonRecord { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + function requireText(parts: string[], providerKey: string): string { const text = parts.map((part) => part.trim()).filter(Boolean).join("\n"); if (!text) throw new Error(`${providerKey} returned no text content.`); @@ -304,12 +288,22 @@ function openAiCompatibleAdapter({ listModels: () => listStaticModels(providerKey), getRateLimitState: providerRateLimitUnknown, checkHealth: (input) => checkOpenAiHealth(input.authProfile, healthUrl), + async streamResponse(input) { + const profile = requireApiKeyProfile(input.authProfile, label); + const messages = responseMessages(input.instructions, input.prompt); + const response = await postEventStream(endpoint, { + authorization: `Bearer ${profile.apiKey}`, + }, { + model: input.model, + messages, + stream: true, + }); + return { events: mapOpenAiChatCompletionStream(response.messages), rateLimitState: response.rateLimitState }; + }, + cancelResponse: cancelUnsupported(label), async createResponse(input) { const profile = requireApiKeyProfile(input.authProfile, label); - const messages = [ - ...(input.instructions ? [{ role: "system", content: input.instructions }] : []), - { role: "user", content: input.prompt }, - ]; + const messages = responseMessages(input.instructions, input.prompt); const response = await postJson(endpoint, { authorization: `Bearer ${profile.apiKey}`, }, { @@ -323,6 +317,13 @@ function openAiCompatibleAdapter({ }; } +function responseMessages(instructions: string | null, prompt: string): Array<{ role: "system" | "user"; content: string }> { + return [ + ...(instructions ? [{ role: "system" as const, content: instructions }] : []), + { role: "user", content: prompt }, + ]; +} + async function checkOpenAiHealth(authProfile: ProviderAuthProfile, healthUrl: string): Promise { const profile = requireApiKeyProfile(authProfile, "OpenAI-compatible provider"); const response = await getJson(healthUrl, { authorization: `Bearer ${profile.apiKey}` }); @@ -341,6 +342,8 @@ function browserTokenAdapter(providerKey: string, authType: ProviderAuthType): P parseAuthProfile, listModels: () => listStaticModels(providerKey), getRateLimitState: providerRateLimitUnknown, + streamResponse: streamUnsupported(providerKey), + cancelResponse: cancelUnsupported(providerKey), checkHealth: async (input) => { if (authType === PROVIDER_AUTH_TYPE.ACCOUNT_SUBSCRIPTION) requireAccountTokenProfile(input.authProfile); else requireAccessTokenProfile(input.authProfile); @@ -461,6 +464,12 @@ function requireVertexProfile(profile: ProviderAuthProfile): Extract; + +export interface ProviderJsonResponse { + data: JsonRecord; + rateLimitState: ProviderObservedRateLimitState | null; +} + +export interface ProviderEventStreamMessage { + eventName: string | null; + data: JsonRecord; + rawData: string; +} + +export interface ProviderEventStreamResponse { + messages: AsyncIterable; + rateLimitState: ProviderObservedRateLimitState | null; +} + +const PROVIDER_REQUEST_TIMEOUT_MS = 60_000; + +export class ProviderHttpError extends Error { + constructor( + message: string, + readonly rateLimitState: ProviderObservedRateLimitState | null, + ) { + super(message); + this.name = "ProviderHttpError"; + } +} + +export function providerRateLimitStateFromError(error: unknown): ProviderObservedRateLimitState | null { + return error instanceof ProviderHttpError ? error.rateLimitState : null; +} + +export async function postJson(endpoint: string, headers: Record, body: JsonRecord): Promise { + return requestJson(endpoint, { + method: "POST", + headers: { + "content-type": "application/json", + ...headers, + }, + body: JSON.stringify(body), + }); +} + +export async function getJson(endpoint: string, headers: Record): Promise { + return requestJson(endpoint, { method: "GET", headers }); +} + +export async function postEventStream( + endpoint: string, + headers: Record, + body: JsonRecord, +): Promise { + const response = await requestProvider(endpoint, { + method: "POST", + headers: { + accept: "text/event-stream", + "content-type": "application/json", + ...headers, + }, + body: JSON.stringify(body), + }); + const rateLimitState = parseProviderRateLimitHeaders(response.headers); + if (!response.ok) { + const raw = await response.text(); + const data = parseProviderJson(raw, false); + const message = readProviderError(data) ?? (raw.trim() || `Provider request failed with HTTP ${response.status}.`); + throw new ProviderHttpError(message, rateLimitState); + } + if (!response.body) throw new ProviderHttpError("Provider returned an empty event stream.", rateLimitState); + return { messages: readServerSentEvents(response.body), rateLimitState }; +} + +async function requestJson(endpoint: string, init: RequestInit): Promise { + const response = await requestProvider(endpoint, init); + const rateLimitState = parseProviderRateLimitHeaders(response.headers); + const raw = await response.text(); + const data = parseProviderJson(raw, response.ok); + if (!response.ok) { + const message = readProviderError(data) ?? (raw.trim() || `Provider request failed with HTTP ${response.status}.`); + throw new ProviderHttpError(message, rateLimitState); + } + return { data, rateLimitState }; +} + +async function requestProvider(endpoint: string, init: RequestInit): Promise { + const controller = new AbortController(); + const timeout = setTimeout(() => controller.abort(), PROVIDER_REQUEST_TIMEOUT_MS); + try { + return await fetch(endpoint, { ...init, signal: controller.signal }); + } finally { + clearTimeout(timeout); + } +} + +export async function* readServerSentEvents(body: ReadableStream): AsyncIterable { + const decoder = new TextDecoder(); + const state: EventStreamParseState = { eventName: null, dataLines: [] }; + let lineBuffer = ""; + const reader = body.getReader(); + try { + for (;;) { + const { value, done } = await reader.read(); + if (done) break; + lineBuffer += decoder.decode(value, { stream: true }); + const lines = lineBuffer.split("\n"); + lineBuffer = lines.pop() ?? ""; + for (const line of lines) { + const message = parseServerSentEventLine(state, line); + if (message) yield message; + } + } + lineBuffer += decoder.decode(); + if (lineBuffer.length > 0) { + const message = parseServerSentEventLine(state, lineBuffer); + if (message) yield message; + } + const message = flushServerSentEvent(state); + if (message) yield message; + } finally { + reader.releaseLock(); + } +} + +function parseServerSentEventLine(state: EventStreamParseState, rawLine: string): ProviderEventStreamMessage | null { + const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine; + if (line === "") return flushServerSentEvent(state); + if (line.startsWith(":")) return null; + const separatorIndex = line.indexOf(":"); + const field = separatorIndex === -1 ? line : line.slice(0, separatorIndex); + const rawValue = separatorIndex === -1 ? "" : line.slice(separatorIndex + 1); + const value = rawValue.startsWith(" ") ? rawValue.slice(1) : rawValue; + if (field === "event") state.eventName = value; + if (field === "data") state.dataLines.push(value); + return null; +} + +function flushServerSentEvent(state: EventStreamParseState): ProviderEventStreamMessage | null { + if (state.dataLines.length === 0) { + state.eventName = null; + return null; + } + const rawData = state.dataLines.join("\n"); + const eventName = state.eventName; + state.eventName = null; + state.dataLines = []; + if (rawData === "[DONE]") return null; + return { eventName, data: parseProviderJson(rawData, true), rawData }; +} + +function parseProviderJson(raw: string, requireJson: boolean): JsonRecord { + if (!raw.trim()) return {}; + let value: unknown; + try { + value = JSON.parse(raw) as unknown; + } catch (error) { + if (!requireJson) return {}; + throw error; + } + if (!isRecord(value)) throw new Error("Provider returned invalid JSON."); + return value; +} + +function readProviderError(data: JsonRecord): string | null { + const error = data.error; + if (typeof error === "string") return error; + if (isRecord(error) && typeof error.message === "string") return error.message; + return null; +} + +export function isRecord(value: unknown): value is JsonRecord { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +interface EventStreamParseState { + eventName: string | null; + dataLines: string[]; +} diff --git a/AgentCompany-New/src/main/services/provider-streaming.ts b/AgentCompany-New/src/main/services/provider-streaming.ts new file mode 100644 index 00000000..0922429a --- /dev/null +++ b/AgentCompany-New/src/main/services/provider-streaming.ts @@ -0,0 +1,169 @@ +import type { ProviderObservedRateLimitState } from "../../shared/provider-runtime-domain"; +import type { JsonRecord, ProviderEventStreamMessage } from "./provider-http"; +import { isRecord } from "./provider-http"; + +export interface ProviderStreamEvent { + type: "message_start" | "text_delta" | "message_stop" | "error"; + textDelta: string | null; + rawEventJson: string; +} + +export interface ProviderStreamResult { + events: AsyncIterable; + rateLimitState: ProviderObservedRateLimitState | null; +} + +export interface ProviderCancelResult { + responseId: string; + rawResponseJson: string; + rateLimitState: ProviderObservedRateLimitState | null; +} + +export async function* mapOpenAiResponseStream(messages: AsyncIterable): AsyncIterable { + let started = false; + for await (const message of messages) { + const type = eventType(message); + if (type === "response.created") { + if (!started) { + started = true; + yield streamEvent("message_start", null, message); + } + continue; + } + if (type === "response.output_text.delta") { + if (!started) { + started = true; + yield streamEvent("message_start", null, message); + } + const delta = stringField(message.data, "delta"); + if (delta) yield streamEvent("text_delta", delta, message); + continue; + } + if (type === "response.completed" || type === "response.incomplete") { + if (!started) { + started = true; + yield streamEvent("message_start", null, message); + } + yield streamEvent("message_stop", null, message); + continue; + } + if (type === "response.failed" || type === "error") { + yield streamEvent("error", readErrorText(message.data), message); + } + } +} + +export async function* mapOpenAiChatCompletionStream(messages: AsyncIterable): AsyncIterable { + let started = false; + for await (const message of messages) { + const delta = readChatDeltaContent(message.data); + if (delta) { + if (!started) { + started = true; + yield streamEvent("message_start", null, message); + } + yield streamEvent("text_delta", delta, message); + } + if (hasChatFinish(message.data)) { + if (!started) { + started = true; + yield streamEvent("message_start", null, message); + } + yield streamEvent("message_stop", null, message); + } + } +} + +export async function* mapAnthropicMessageStream(messages: AsyncIterable): AsyncIterable { + let started = false; + for await (const message of messages) { + const type = eventType(message); + if (type === "message_start") { + started = true; + yield streamEvent("message_start", null, message); + continue; + } + if (type === "content_block_delta") { + const delta = readAnthropicTextDelta(message.data); + if (delta) { + if (!started) { + started = true; + yield streamEvent("message_start", null, message); + } + yield streamEvent("text_delta", delta, message); + } + continue; + } + if (type === "message_stop") { + if (!started) { + started = true; + yield streamEvent("message_start", null, message); + } + yield streamEvent("message_stop", null, message); + continue; + } + if (type === "error") yield streamEvent("error", readErrorText(message.data), message); + } +} + +export function streamUnsupported(providerLabel: string): () => Promise { + return async () => { + throw new Error(`${providerLabel} streaming is unavailable.`); + }; +} + +export function cancelUnsupported(providerLabel: string): () => Promise { + return async () => { + throw new Error(`${providerLabel} response cancellation is unavailable.`); + }; +} + +function streamEvent(type: ProviderStreamEvent["type"], textDelta: string | null, message: ProviderEventStreamMessage): ProviderStreamEvent { + return { type, textDelta, rawEventJson: message.rawData }; +} + +function eventType(message: ProviderEventStreamMessage): string | null { + const dataType = stringField(message.data, "type"); + return dataType ?? message.eventName; +} + +function readChatDeltaContent(data: JsonRecord): string | null { + for (const choice of recordArray(data.choices)) { + const delta = recordField(choice, "delta"); + const content = delta ? stringField(delta, "content") : null; + if (content) return content; + } + return null; +} + +function hasChatFinish(data: JsonRecord): boolean { + return recordArray(data.choices).some((choice) => typeof choice.finish_reason === "string" && choice.finish_reason.length > 0); +} + +function readAnthropicTextDelta(data: JsonRecord): string | null { + const delta = recordField(data, "delta"); + if (!delta) return null; + return stringField(delta, "text"); +} + +function readErrorText(data: JsonRecord): string | null { + const error = data.error; + if (typeof error === "string") return error; + if (isRecord(error) && typeof error.message === "string") return error.message; + return null; +} + +function stringField(record: JsonRecord, key: string): string | null { + const value = record[key]; + return typeof value === "string" ? value : null; +} + +function recordField(record: JsonRecord, key: string): JsonRecord | null { + const value = record[key]; + return isRecord(value) ? value : null; +} + +function recordArray(value: unknown): JsonRecord[] { + if (!Array.isArray(value)) return []; + return value.filter(isRecord); +} diff --git a/AgentCompany-New/src/main/services/sdlc-service.ts b/AgentCompany-New/src/main/services/sdlc-service.ts index d0c0095f..4203dc9b 100644 --- a/AgentCompany-New/src/main/services/sdlc-service.ts +++ b/AgentCompany-New/src/main/services/sdlc-service.ts @@ -248,7 +248,9 @@ export class SdlcService { && typeof adapter.refresh === "function" && typeof adapter.revoke === "function" && typeof adapter.listModels === "function" - && typeof adapter.getRateLimitState === "function"; + && typeof adapter.getRateLimitState === "function" + && typeof adapter.streamResponse === "function" + && typeof adapter.cancelResponse === "function"; }); const schemaVersion = count(this.db, "select coalesce(max(version), 0) as count from schema_migrations"); const rateLimitRegistry = count(this.db, ` diff --git a/AgentCompany-New/tests/contract/provider-contracts.test.ts b/AgentCompany-New/tests/contract/provider-contracts.test.ts index 12e6f36f..d52a02b6 100644 --- a/AgentCompany-New/tests/contract/provider-contracts.test.ts +++ b/AgentCompany-New/tests/contract/provider-contracts.test.ts @@ -87,6 +87,8 @@ describe("provider contracts", () => { const authProfile = adapter.parseAuthProfile(PROVIDER_SECRET_SAMPLES[provider.key]); await expect(adapter.listModels({ authProfile, model: provider.defaultModel })).resolves.toEqual(provider.models); await expect(adapter.getRateLimitState({ authProfile, model: provider.defaultModel })).resolves.toBeNull(); + expect(typeof adapter.streamResponse).toBe("function"); + expect(typeof adapter.cancelResponse).toBe("function"); } }); diff --git a/AgentCompany-New/tests/unit/provider-streaming.test.ts b/AgentCompany-New/tests/unit/provider-streaming.test.ts new file mode 100644 index 00000000..669437e2 --- /dev/null +++ b/AgentCompany-New/tests/unit/provider-streaming.test.ts @@ -0,0 +1,106 @@ +import { describe, expect, test } from "vitest"; +import { readServerSentEvents, type ProviderEventStreamMessage } from "../../src/main/services/provider-http"; +import { + cancelUnsupported, + mapAnthropicMessageStream, + mapOpenAiChatCompletionStream, + mapOpenAiResponseStream, + streamUnsupported, +} from "../../src/main/services/provider-streaming"; + +describe("provider streaming", () => { + test("parses server sent events and skips done sentinels", async () => { + const source = [ + "event: response.output_text.delta", + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"Hel\"}", + "", + "data: [DONE]", + "", + ].join("\n"); + + const messages = await collect(readServerSentEvents(streamFromText(source, 19))); + + expect(messages).toEqual([{ + eventName: "response.output_text.delta", + data: { type: "response.output_text.delta", delta: "Hel" }, + rawData: "{\"type\":\"response.output_text.delta\",\"delta\":\"Hel\"}", + }]); + }); + + test("maps OpenAI Responses stream deltas", async () => { + const events = await collect(mapOpenAiResponseStream(asyncMessages([ + message("response.created", { type: "response.created" }), + message("response.output_text.delta", { type: "response.output_text.delta", delta: "Hel" }), + message("response.output_text.delta", { type: "response.output_text.delta", delta: "lo" }), + message("response.completed", { type: "response.completed" }), + ]))); + + expect(events.map((event) => [event.type, event.textDelta])).toEqual([ + ["message_start", null], + ["text_delta", "Hel"], + ["text_delta", "lo"], + ["message_stop", null], + ]); + }); + + test("maps OpenAI-compatible chat completion deltas", async () => { + const events = await collect(mapOpenAiChatCompletionStream(asyncMessages([ + message(null, { choices: [{ delta: { content: "Hel" } }] }), + message(null, { choices: [{ delta: { content: "lo" } }] }), + message(null, { choices: [{ delta: {}, finish_reason: "stop" }] }), + ]))); + + expect(events.map((event) => [event.type, event.textDelta])).toEqual([ + ["message_start", null], + ["text_delta", "Hel"], + ["text_delta", "lo"], + ["message_stop", null], + ]); + }); + + test("maps Anthropic message stream deltas", async () => { + const events = await collect(mapAnthropicMessageStream(asyncMessages([ + message("message_start", { type: "message_start" }), + message("content_block_delta", { type: "content_block_delta", delta: { type: "text_delta", text: "Hel" } }), + message("content_block_delta", { type: "content_block_delta", delta: { type: "text_delta", text: "lo" } }), + message("message_stop", { type: "message_stop" }), + ]))); + + expect(events.map((event) => [event.type, event.textDelta])).toEqual([ + ["message_start", null], + ["text_delta", "Hel"], + ["text_delta", "lo"], + ["message_stop", null], + ]); + }); + + test("fails unsupported stream and cancel contracts early", async () => { + await expect(streamUnsupported("Provider")()).rejects.toThrow("Provider streaming is unavailable."); + await expect(cancelUnsupported("Provider")()).rejects.toThrow("Provider response cancellation is unavailable."); + }); +}); + +async function collect(items: AsyncIterable): Promise { + const collected: T[] = []; + for await (const item of items) collected.push(item); + return collected; +} + +async function* asyncMessages(messages: ProviderEventStreamMessage[]): AsyncIterable { + for (const item of messages) yield item; +} + +function message(eventName: string | null, data: Record): ProviderEventStreamMessage { + return { eventName, data, rawData: JSON.stringify(data) }; +} + +function streamFromText(source: string, splitAt: number): ReadableStream { + const encoder = new TextEncoder(); + return new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(source.slice(0, splitAt))); + controller.enqueue(encoder.encode(source.slice(splitAt))); + controller.close(); + }, + }); +}