Lifei Zhou
86 lines
1.9 KiB
TypeScript
86 lines
1.9 KiB
TypeScript
import type { Stream } from '@aaif/goose-sdk';
|
|
|
|
export type ClosableAcpStream = Stream & {
|
|
close: () => void;
|
|
};
|
|
|
|
export function createWebSocketStream(wsUrl: string): ClosableAcpStream {
|
|
const ws = new window.WebSocket(wsUrl);
|
|
|
|
const incoming: unknown[] = [];
|
|
const waiters: Array<() => void> = [];
|
|
let closed = false;
|
|
|
|
function pushMessage(message: unknown): void {
|
|
incoming.push(message);
|
|
waiters.shift()?.();
|
|
}
|
|
|
|
function waitForMessage(): Promise<void> {
|
|
if (incoming.length > 0 || closed) {
|
|
return Promise.resolve();
|
|
}
|
|
return new Promise<void>((resolve) => waiters.push(resolve));
|
|
}
|
|
|
|
const openPromise = new Promise<void>((resolve, reject) => {
|
|
ws.addEventListener('open', () => resolve(), { once: true });
|
|
ws.addEventListener('error', () => reject(new Error('ACP WebSocket connection failed')), {
|
|
once: true,
|
|
});
|
|
});
|
|
|
|
ws.addEventListener('message', (event) => {
|
|
if (typeof event.data !== 'string') {
|
|
return;
|
|
}
|
|
try {
|
|
pushMessage(JSON.parse(event.data));
|
|
} catch {
|
|
// Ignore malformed messages from the transport.
|
|
}
|
|
});
|
|
|
|
const closeWaiters = () => {
|
|
closed = true;
|
|
for (const waiter of waiters) {
|
|
waiter();
|
|
}
|
|
waiters.length = 0;
|
|
};
|
|
|
|
ws.addEventListener('close', closeWaiters);
|
|
ws.addEventListener('error', closeWaiters);
|
|
|
|
const readable = new window.ReadableStream({
|
|
async pull(controller) {
|
|
await waitForMessage();
|
|
while (incoming.length > 0) {
|
|
controller.enqueue(incoming.shift());
|
|
}
|
|
if (closed && incoming.length === 0) {
|
|
controller.close();
|
|
}
|
|
},
|
|
});
|
|
|
|
const writable = new window.WritableStream({
|
|
async write(message) {
|
|
await openPromise;
|
|
ws.send(JSON.stringify(message));
|
|
},
|
|
close() {
|
|
ws.close();
|
|
},
|
|
abort() {
|
|
ws.close();
|
|
},
|
|
});
|
|
|
|
return {
|
|
readable,
|
|
writable,
|
|
close: () => ws.close(),
|
|
} as ClosableAcpStream;
|
|
}
|