mirror of
https://github.com/decolua/9router.git
synced 2026-05-08 12:01:28 +00:00
132 lines
3.6 KiB
JavaScript
132 lines
3.6 KiB
JavaScript
// Stream handler with disconnect detection - shared for all providers
|
|
|
|
// Get HH:MM timestamp
|
|
function getTimeString() {
|
|
return new Date().toLocaleTimeString("en-US", { hour12: false, hour: "2-digit", minute: "2-digit" });
|
|
}
|
|
|
|
/**
|
|
* Create stream controller with abort and disconnect detection
|
|
* @param {object} options
|
|
* @param {function} options.onDisconnect - Callback when client disconnects
|
|
* @param {object} options.log - Logger instance
|
|
* @param {string} options.provider - Provider name
|
|
* @param {string} options.model - Model name
|
|
*/
|
|
export function createStreamController({ onDisconnect, log, provider, model } = {}) {
|
|
const abortController = new AbortController();
|
|
const startTime = Date.now();
|
|
let disconnected = false;
|
|
let abortTimeout = null;
|
|
|
|
const logStream = (status) => {
|
|
const duration = Date.now() - startTime;
|
|
const p = provider?.toUpperCase() || "UNKNOWN";
|
|
console.log(`[${getTimeString()}] 🌊 [STREAM] ${p} | ${model || "unknown"} | ${duration}ms | ${status}`);
|
|
};
|
|
|
|
return {
|
|
signal: abortController.signal,
|
|
startTime,
|
|
|
|
isConnected: () => !disconnected,
|
|
|
|
// Call when client disconnects
|
|
handleDisconnect: (reason = "client_closed") => {
|
|
if (disconnected) return;
|
|
disconnected = true;
|
|
|
|
logStream(`disconnect: ${reason}`);
|
|
|
|
// Delay abort to allow cleanup
|
|
abortTimeout = setTimeout(() => {
|
|
abortController.abort();
|
|
}, 500);
|
|
|
|
onDisconnect?.({ reason, duration: Date.now() - startTime });
|
|
},
|
|
|
|
// Call when stream completes normally
|
|
handleComplete: () => {
|
|
if (disconnected) return;
|
|
disconnected = true;
|
|
|
|
logStream("complete");
|
|
|
|
if (abortTimeout) {
|
|
clearTimeout(abortTimeout);
|
|
abortTimeout = null;
|
|
}
|
|
},
|
|
|
|
// Call on error
|
|
handleError: (error) => {
|
|
if (abortTimeout) {
|
|
clearTimeout(abortTimeout);
|
|
abortTimeout = null;
|
|
}
|
|
|
|
if (error.name === "AbortError") {
|
|
logStream("aborted");
|
|
return;
|
|
}
|
|
|
|
logStream(`error: ${error.message}`);
|
|
},
|
|
|
|
abort: () => abortController.abort()
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Create transform stream with disconnect detection
|
|
* Wraps existing transform stream and adds abort capability
|
|
*/
|
|
export function createDisconnectAwareStream(transformStream, streamController) {
|
|
const reader = transformStream.readable.getReader();
|
|
const writer = transformStream.writable.getWriter();
|
|
|
|
return new ReadableStream({
|
|
async pull(controller) {
|
|
if (!streamController.isConnected()) {
|
|
controller.close();
|
|
return;
|
|
}
|
|
|
|
try {
|
|
const { done, value } = await reader.read();
|
|
if (done) {
|
|
streamController.handleComplete();
|
|
controller.close();
|
|
return;
|
|
}
|
|
controller.enqueue(value);
|
|
} catch (error) {
|
|
streamController.handleError(error);
|
|
controller.error(error);
|
|
}
|
|
},
|
|
|
|
cancel(reason) {
|
|
streamController.handleDisconnect(reason || "cancelled");
|
|
reader.cancel();
|
|
writer.abort();
|
|
}
|
|
});
|
|
}
|
|
|
|
/**
|
|
* Pipe provider response through transform with disconnect detection
|
|
* @param {Response} providerResponse - Response from provider
|
|
* @param {TransformStream} transformStream - Transform stream for SSE
|
|
* @param {object} streamController - Stream controller from createStreamController
|
|
*/
|
|
export function pipeWithDisconnect(providerResponse, transformStream, streamController) {
|
|
const transformedBody = providerResponse.body.pipeThrough(transformStream);
|
|
return createDisconnectAwareStream(
|
|
{ readable: transformedBody, writable: { getWriter: () => ({ abort: () => {} }) } },
|
|
streamController
|
|
);
|
|
}
|
|
|