2023-10-05 07:33:34 +00:00
|
|
|
import type { Context } from '../../context.ts'
|
|
|
|
import { StreamingApi } from '../../utils/stream.ts'
|
|
|
|
|
|
|
|
interface SSEMessage {
|
|
|
|
data: string
|
|
|
|
event?: string
|
|
|
|
id?: string
|
|
|
|
}
|
|
|
|
|
|
|
|
class SSEStreamingApi extends StreamingApi {
|
|
|
|
constructor(writable: WritableStream) {
|
|
|
|
super(writable)
|
|
|
|
}
|
|
|
|
|
|
|
|
async writeSSE(message: SSEMessage) {
|
2023-10-16 09:54:28 +00:00
|
|
|
const data = message.data
|
|
|
|
.split('\n')
|
|
|
|
.map((line) => {
|
|
|
|
return `data: ${line}`
|
|
|
|
})
|
|
|
|
.join('\n')
|
|
|
|
|
2023-10-05 07:33:34 +00:00
|
|
|
const sseData =
|
2023-10-16 09:54:28 +00:00
|
|
|
[message.event && `event: ${message.event}`, data, message.id && `id: ${message.id}`]
|
2023-10-05 07:33:34 +00:00
|
|
|
.filter(Boolean)
|
2023-10-16 09:54:28 +00:00
|
|
|
.join('\n') + '\n\n'
|
2023-10-05 07:33:34 +00:00
|
|
|
|
|
|
|
await this.write(sseData)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
const setSSEHeaders = (context: Context) => {
|
|
|
|
context.header('Transfer-Encoding', 'chunked')
|
|
|
|
context.header('Content-Type', 'text/event-stream')
|
|
|
|
context.header('Cache-Control', 'no-cache')
|
|
|
|
context.header('Connection', 'keep-alive')
|
|
|
|
}
|
|
|
|
|
|
|
|
export const streamSSE = (c: Context, cb: (stream: SSEStreamingApi) => Promise<void>) => {
|
|
|
|
return c.stream(async (originalStream: StreamingApi) => {
|
|
|
|
const { readable, writable } = new TransformStream()
|
|
|
|
const stream = new SSEStreamingApi(writable)
|
|
|
|
originalStream.pipe(readable)
|
|
|
|
setSSEHeaders(c)
|
|
|
|
|
|
|
|
try {
|
|
|
|
await cb(stream)
|
|
|
|
} catch (err) {
|
|
|
|
console.error('Error during streaming: ', err)
|
|
|
|
stream.close()
|
|
|
|
}
|
|
|
|
})
|
|
|
|
}
|