chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,554 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import type {
|
||||
SSEStreamPart,
|
||||
StreamSubscription,
|
||||
StreamSubscriptionFactory,
|
||||
} from "../src/v3/apiClient/runStream.js";
|
||||
import { RunSubscription, SSEStreamSubscription } from "../src/v3/apiClient/runStream.js";
|
||||
import type { SubscribeRunRawShape } from "../src/v3/schemas/api.js";
|
||||
|
||||
// Test implementations
|
||||
// Update TestStreamSubscription to return a ReadableStream
|
||||
class TestStreamSubscription implements StreamSubscription {
|
||||
constructor(private chunks: unknown[]) {}
|
||||
|
||||
async subscribe(): Promise<ReadableStream<SSEStreamPart<unknown>>> {
|
||||
return new ReadableStream({
|
||||
start: async (controller) => {
|
||||
for (let i = 0; i < this.chunks.length; i++) {
|
||||
controller.enqueue({
|
||||
id: `msg-${i}`,
|
||||
chunk: this.chunks[i],
|
||||
timestamp: Date.now() + i,
|
||||
});
|
||||
}
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamSubscriptionFactory can remain the same
|
||||
class TestStreamSubscriptionFactory implements StreamSubscriptionFactory {
|
||||
private streams = new Map<string, unknown[]>();
|
||||
|
||||
setStreamChunks(runId: string, streamKey: string, chunks: unknown[]) {
|
||||
this.streams.set(`${runId}:${streamKey}`, chunks);
|
||||
}
|
||||
|
||||
createSubscription(runId: string, streamKey: string): StreamSubscription {
|
||||
const chunks = this.streams.get(`${runId}:${streamKey}`) ?? [];
|
||||
return new TestStreamSubscription(chunks);
|
||||
}
|
||||
}
|
||||
|
||||
// Remove the RunShapeProvider implementations and replace with stream creators
|
||||
function createTestShapeStream(
|
||||
shapes: SubscribeRunRawShape[]
|
||||
): ReadableStream<SubscribeRunRawShape> {
|
||||
return new ReadableStream({
|
||||
start: async (controller) => {
|
||||
// Emit all shapes immediately
|
||||
for (const shape of shapes) {
|
||||
controller.enqueue(shape);
|
||||
}
|
||||
controller.close();
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
function createDelayedTestShapeStream(
|
||||
shapes: SubscribeRunRawShape[]
|
||||
): ReadableStream<SubscribeRunRawShape> {
|
||||
return new ReadableStream({
|
||||
start: async (controller) => {
|
||||
// Emit first shape immediately
|
||||
if (shapes.length > 0) {
|
||||
controller.enqueue(shapes[0]);
|
||||
}
|
||||
|
||||
let currentShapeIndex = 1;
|
||||
|
||||
// Emit remaining shapes with delay
|
||||
const interval = setInterval(() => {
|
||||
if (currentShapeIndex >= shapes.length) {
|
||||
clearInterval(interval);
|
||||
controller.close();
|
||||
return;
|
||||
}
|
||||
controller.enqueue(shapes[currentShapeIndex++]!);
|
||||
}, 100);
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
describe("RunSubscription", () => {
|
||||
it("should handle basic run subscription", async () => {
|
||||
const shapes = [
|
||||
{
|
||||
id: "123",
|
||||
friendlyId: "run_123",
|
||||
taskIdentifier: "test-task",
|
||||
status: "COMPLETED_SUCCESSFULLY",
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
completedAt: new Date(),
|
||||
number: 1,
|
||||
usageDurationMs: 100,
|
||||
costInCents: 0,
|
||||
baseCostInCents: 0,
|
||||
isTest: false,
|
||||
runTags: [],
|
||||
realtimeStreams: [],
|
||||
},
|
||||
];
|
||||
|
||||
const subscription = new RunSubscription({
|
||||
runShapeStream: createTestShapeStream(shapes),
|
||||
stopRunShapeStream: () => {},
|
||||
streamFactory: new TestStreamSubscriptionFactory(),
|
||||
closeOnComplete: true,
|
||||
abortController: new AbortController(),
|
||||
});
|
||||
|
||||
const results = await convertAsyncIterableToArray(subscription);
|
||||
|
||||
expect(results).toHaveLength(1);
|
||||
expect(results[0]).toMatchObject({
|
||||
id: "run_123",
|
||||
taskIdentifier: "test-task",
|
||||
status: "COMPLETED",
|
||||
});
|
||||
});
|
||||
|
||||
it("should handle payload and outputs", async () => {
|
||||
const shapes: SubscribeRunRawShape[] = [
|
||||
{
|
||||
id: "123",
|
||||
friendlyId: "run_123",
|
||||
taskIdentifier: "test-task",
|
||||
status: "COMPLETED_SUCCESSFULLY",
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
completedAt: new Date(),
|
||||
number: 1,
|
||||
usageDurationMs: 100,
|
||||
costInCents: 0,
|
||||
baseCostInCents: 0,
|
||||
isTest: false,
|
||||
runTags: [],
|
||||
payload: JSON.stringify({ test: "payload" }),
|
||||
payloadType: "application/json",
|
||||
output: JSON.stringify({ test: "output" }),
|
||||
outputType: "application/json",
|
||||
realtimeStreams: [],
|
||||
},
|
||||
];
|
||||
|
||||
const subscription = new RunSubscription({
|
||||
runShapeStream: createTestShapeStream(shapes),
|
||||
stopRunShapeStream: () => {},
|
||||
streamFactory: new TestStreamSubscriptionFactory(),
|
||||
closeOnComplete: true,
|
||||
abortController: new AbortController(),
|
||||
});
|
||||
|
||||
const results = await convertAsyncIterableToArray(subscription);
|
||||
|
||||
expect(results).toHaveLength(1);
|
||||
expect(results[0]).toMatchObject({
|
||||
id: "run_123",
|
||||
taskIdentifier: "test-task",
|
||||
status: "COMPLETED",
|
||||
payload: { test: "payload" },
|
||||
output: { test: "output" },
|
||||
});
|
||||
});
|
||||
|
||||
it("should keep stream open when closeOnComplete is false", async () => {
|
||||
const shapes: SubscribeRunRawShape[] = [
|
||||
{
|
||||
id: "123",
|
||||
friendlyId: "run_123",
|
||||
taskIdentifier: "test-task",
|
||||
status: "EXECUTING",
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
completedAt: new Date(),
|
||||
number: 1,
|
||||
usageDurationMs: 100,
|
||||
costInCents: 0,
|
||||
baseCostInCents: 0,
|
||||
isTest: false,
|
||||
runTags: [],
|
||||
realtimeStreams: [],
|
||||
},
|
||||
{
|
||||
id: "123",
|
||||
friendlyId: "run_123",
|
||||
taskIdentifier: "test-task",
|
||||
status: "COMPLETED_SUCCESSFULLY",
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
completedAt: new Date(),
|
||||
number: 1,
|
||||
usageDurationMs: 200,
|
||||
costInCents: 0,
|
||||
baseCostInCents: 0,
|
||||
isTest: false,
|
||||
runTags: [],
|
||||
realtimeStreams: [],
|
||||
},
|
||||
];
|
||||
|
||||
const subscription = new RunSubscription({
|
||||
runShapeStream: createDelayedTestShapeStream(shapes),
|
||||
stopRunShapeStream: () => {},
|
||||
streamFactory: new TestStreamSubscriptionFactory(),
|
||||
closeOnComplete: false,
|
||||
abortController: new AbortController(),
|
||||
});
|
||||
|
||||
// Collect 2 results
|
||||
const results = await collectNResults(subscription, 2);
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results[0]).toMatchObject({
|
||||
id: "run_123",
|
||||
taskIdentifier: "test-task",
|
||||
status: "EXECUTING",
|
||||
});
|
||||
expect(results[1]).toMatchObject({
|
||||
id: "run_123",
|
||||
taskIdentifier: "test-task",
|
||||
status: "COMPLETED",
|
||||
});
|
||||
});
|
||||
|
||||
it("should handle stream data", async () => {
|
||||
const streamFactory = new TestStreamSubscriptionFactory();
|
||||
|
||||
// Set up test chunks
|
||||
streamFactory.setStreamChunks("run_123", "openai", [
|
||||
{ id: "chunk1", content: "Hello" },
|
||||
{ id: "chunk2", content: "World" },
|
||||
]);
|
||||
|
||||
const shapes = [
|
||||
{
|
||||
id: "123",
|
||||
friendlyId: "run_123",
|
||||
taskIdentifier: "openai-streaming",
|
||||
status: "EXECUTING",
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
number: 1,
|
||||
usageDurationMs: 100,
|
||||
costInCents: 0,
|
||||
baseCostInCents: 0,
|
||||
isTest: false,
|
||||
runTags: [],
|
||||
metadata: JSON.stringify({}),
|
||||
metadataType: "application/json",
|
||||
realtimeStreams: ["openai"],
|
||||
},
|
||||
];
|
||||
|
||||
const subscription = new RunSubscription({
|
||||
runShapeStream: createTestShapeStream(shapes),
|
||||
stopRunShapeStream: () => {},
|
||||
streamFactory,
|
||||
abortController: new AbortController(),
|
||||
});
|
||||
|
||||
const results = await collectNResults(
|
||||
subscription.withStreams<{ openai: { id: string; content: string } }>(),
|
||||
3 // 1 run + 2 stream chunks
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(3);
|
||||
expect(results[0]).toMatchObject({
|
||||
type: "run",
|
||||
run: { id: "run_123", taskIdentifier: "openai-streaming", status: "EXECUTING" },
|
||||
});
|
||||
expect(results[1]).toMatchObject({
|
||||
type: "openai",
|
||||
chunk: { id: "chunk1", content: "Hello" },
|
||||
run: { id: "run_123", taskIdentifier: "openai-streaming", status: "EXECUTING" },
|
||||
});
|
||||
expect(results[2]).toMatchObject({
|
||||
type: "openai",
|
||||
chunk: { id: "chunk2", content: "World" },
|
||||
run: { id: "run_123", taskIdentifier: "openai-streaming", status: "EXECUTING" },
|
||||
});
|
||||
});
|
||||
|
||||
it("should only create one stream for multiple runs of the same id", async () => {
|
||||
const streamFactory = new TestStreamSubscriptionFactory();
|
||||
let streamCreationCount = 0;
|
||||
|
||||
// Override createSubscription to count calls
|
||||
const originalCreate = streamFactory.createSubscription.bind(streamFactory);
|
||||
streamFactory.createSubscription = (runId: string, streamKey: string) => {
|
||||
streamCreationCount++;
|
||||
return originalCreate(runId, streamKey);
|
||||
};
|
||||
|
||||
// Set up test chunks
|
||||
streamFactory.setStreamChunks("run_123", "openai", [
|
||||
{ id: "chunk1", content: "Hello" },
|
||||
{ id: "chunk2", content: "World" },
|
||||
]);
|
||||
|
||||
const shapes = [
|
||||
// First run update
|
||||
{
|
||||
id: "123",
|
||||
friendlyId: "run_123",
|
||||
taskIdentifier: "openai-streaming",
|
||||
status: "EXECUTING",
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
number: 1,
|
||||
usageDurationMs: 100,
|
||||
costInCents: 0,
|
||||
baseCostInCents: 0,
|
||||
isTest: false,
|
||||
runTags: [],
|
||||
metadata: JSON.stringify({}),
|
||||
metadataType: "application/json",
|
||||
realtimeStreams: ["openai"],
|
||||
},
|
||||
// Second run update with same stream key
|
||||
{
|
||||
id: "123",
|
||||
friendlyId: "run_123",
|
||||
taskIdentifier: "openai-streaming",
|
||||
status: "EXECUTING",
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
number: 1,
|
||||
usageDurationMs: 200, // Different to show it's a new update
|
||||
costInCents: 0,
|
||||
baseCostInCents: 0,
|
||||
isTest: false,
|
||||
runTags: [],
|
||||
metadata: JSON.stringify({}),
|
||||
metadataType: "application/json",
|
||||
realtimeStreams: ["openai"],
|
||||
},
|
||||
];
|
||||
|
||||
const subscription = new RunSubscription({
|
||||
runShapeStream: createTestShapeStream(shapes),
|
||||
stopRunShapeStream: () => {},
|
||||
streamFactory,
|
||||
abortController: new AbortController(),
|
||||
});
|
||||
|
||||
const results = await collectNResults(
|
||||
subscription.withStreams<{ openai: { id: string; content: string } }>(),
|
||||
4 // 2 runs + 2 stream chunks
|
||||
);
|
||||
|
||||
// Verify we only created one stream
|
||||
expect(streamCreationCount).toBe(1);
|
||||
|
||||
// Verify we got all the expected events
|
||||
expect(results).toHaveLength(4);
|
||||
expect(results[0]).toMatchObject({
|
||||
type: "run",
|
||||
run: {
|
||||
id: "run_123",
|
||||
taskIdentifier: "openai-streaming",
|
||||
status: "EXECUTING",
|
||||
durationMs: 100,
|
||||
},
|
||||
});
|
||||
expect(results[1]).toMatchObject({
|
||||
type: "openai",
|
||||
chunk: { id: "chunk1", content: "Hello" },
|
||||
run: { id: "run_123", durationMs: 100 },
|
||||
});
|
||||
expect(results[2]).toMatchObject({
|
||||
type: "openai",
|
||||
chunk: { id: "chunk2", content: "World" },
|
||||
run: { id: "run_123", durationMs: 100 },
|
||||
});
|
||||
expect(results[3]).toMatchObject({
|
||||
type: "run",
|
||||
run: {
|
||||
id: "run_123",
|
||||
taskIdentifier: "openai-streaming",
|
||||
status: "EXECUTING",
|
||||
durationMs: 200,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it("should handle multiple streams simultaneously", async () => {
|
||||
const streamFactory = new TestStreamSubscriptionFactory();
|
||||
|
||||
// Set up test chunks for two different streams
|
||||
streamFactory.setStreamChunks("run_123", "openai", [
|
||||
{ id: "openai1", content: "Hello" },
|
||||
{ id: "openai2", content: "World" },
|
||||
]);
|
||||
streamFactory.setStreamChunks("run_123", "anthropic", [
|
||||
{ id: "claude1", message: "Hi" },
|
||||
{ id: "claude2", message: "There" },
|
||||
]);
|
||||
|
||||
const shapes = [
|
||||
{
|
||||
id: "123",
|
||||
friendlyId: "run_123",
|
||||
taskIdentifier: "multi-streaming",
|
||||
status: "EXECUTING",
|
||||
createdAt: new Date(),
|
||||
updatedAt: new Date(),
|
||||
number: 1,
|
||||
usageDurationMs: 100,
|
||||
costInCents: 0,
|
||||
baseCostInCents: 0,
|
||||
isTest: false,
|
||||
runTags: [],
|
||||
metadata: JSON.stringify({}),
|
||||
metadataType: "application/json",
|
||||
realtimeStreams: ["openai", "anthropic"],
|
||||
},
|
||||
];
|
||||
|
||||
const subscription = new RunSubscription({
|
||||
runShapeStream: createTestShapeStream(shapes),
|
||||
stopRunShapeStream: () => {},
|
||||
streamFactory,
|
||||
abortController: new AbortController(),
|
||||
});
|
||||
|
||||
const results = await collectNResults(
|
||||
subscription.withStreams<{
|
||||
openai: { id: string; content: string };
|
||||
anthropic: { id: string; message: string };
|
||||
}>(),
|
||||
5 // 1 run + 2 openai chunks + 2 anthropic chunks
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(5);
|
||||
expect(results[0]).toMatchObject({
|
||||
type: "run",
|
||||
run: { id: "run_123", taskIdentifier: "multi-streaming", status: "EXECUTING" },
|
||||
});
|
||||
|
||||
// Filter and verify openai chunks
|
||||
const openaiChunks = results.filter((r) => r.type === "openai");
|
||||
expect(openaiChunks).toHaveLength(2);
|
||||
expect(openaiChunks[0]).toMatchObject({
|
||||
type: "openai",
|
||||
chunk: { id: "openai1", content: "Hello" },
|
||||
run: { id: "run_123" },
|
||||
});
|
||||
expect(openaiChunks[1]).toMatchObject({
|
||||
type: "openai",
|
||||
chunk: { id: "openai2", content: "World" },
|
||||
run: { id: "run_123" },
|
||||
});
|
||||
|
||||
// Filter and verify anthropic chunks
|
||||
const anthropicChunks = results.filter((r) => r.type === "anthropic");
|
||||
expect(anthropicChunks).toHaveLength(2);
|
||||
expect(anthropicChunks[0]).toMatchObject({
|
||||
type: "anthropic",
|
||||
chunk: { id: "claude1", message: "Hi" },
|
||||
run: { id: "run_123" },
|
||||
});
|
||||
expect(anthropicChunks[1]).toMatchObject({
|
||||
type: "anthropic",
|
||||
chunk: { id: "claude2", message: "There" },
|
||||
run: { id: "run_123" },
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("SSEStreamSubscription", () => {
|
||||
let originalFetch: typeof global.fetch;
|
||||
|
||||
beforeEach(() => {
|
||||
originalFetch = global.fetch;
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
global.fetch = originalFetch;
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it("does not retry the initial fetch on 401", async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(new Response(null, { status: 401 }));
|
||||
global.fetch = fetchMock;
|
||||
|
||||
const sub = new SSEStreamSubscription("https://api.test/realtime/v1/streams/run_x/chat", {
|
||||
headers: { Authorization: "Bearer expired" },
|
||||
});
|
||||
|
||||
const stream = await sub.subscribe();
|
||||
const reader = stream.getReader();
|
||||
await expect(reader.read()).rejects.toMatchObject({ status: 401 });
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it("does not retry the initial fetch on 403", async () => {
|
||||
const fetchMock = vi.fn().mockResolvedValue(new Response(null, { status: 403 }));
|
||||
global.fetch = fetchMock;
|
||||
|
||||
const sub = new SSEStreamSubscription("https://api.test/realtime/v1/streams/run_x/chat", {
|
||||
headers: { Authorization: "Bearer denied" },
|
||||
});
|
||||
|
||||
const stream = await sub.subscribe();
|
||||
const reader = stream.getReader();
|
||||
await expect(reader.read()).rejects.toMatchObject({ status: 403 });
|
||||
expect(fetchMock).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
});
|
||||
|
||||
export async function convertAsyncIterableToArray<T>(iterable: AsyncIterable<T>): Promise<T[]> {
|
||||
const result: T[] = [];
|
||||
for await (const item of iterable) {
|
||||
result.push(item);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
async function collectNResults<T>(
|
||||
iterable: AsyncIterable<T>,
|
||||
count: number,
|
||||
timeoutMs: number = 1000
|
||||
): Promise<T[]> {
|
||||
const results: T[] = [];
|
||||
const promise = new Promise<T[]>((resolve) => {
|
||||
(async () => {
|
||||
for await (const result of iterable) {
|
||||
results.push(result);
|
||||
if (results.length === count) {
|
||||
resolve(results);
|
||||
break;
|
||||
}
|
||||
}
|
||||
})();
|
||||
});
|
||||
|
||||
return Promise.race([
|
||||
promise,
|
||||
new Promise<T[]>((_, reject) =>
|
||||
setTimeout(
|
||||
() =>
|
||||
reject(
|
||||
new Error(
|
||||
`Timeout waiting for ${count} results after ${timeoutMs}ms, but only had ${results.length}`
|
||||
)
|
||||
),
|
||||
timeoutMs
|
||||
)
|
||||
),
|
||||
]);
|
||||
}
|
||||
Reference in New Issue
Block a user