Skip to content

Instantly share code, notes, and snippets.

@JurajBegovac
Last active December 26, 2023 05:47
Show Gist options
  • Star 2 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save JurajBegovac/0007ae0a9631fe8606a48791c94ab6c6 to your computer and use it in GitHub Desktop.
Save JurajBegovac/0007ae0a9631fe8606a48791c94ab6c6 to your computer and use it in GitHub Desktop.
Implementation for SSE with ktor for KMM (kotlin mobile multiplatform). ktorVersion = "2.2.1" coroutinesVersion = "1.6.4"
MIT License
Copyright (c) 2022-2023 Juraj Begovac
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
// package
import com.wesports.data.common.ktor.addOrReplaceParameter
import io.ktor.client.HttpClient
import io.ktor.client.request.headers
import io.ktor.client.request.prepareGet
import io.ktor.client.statement.HttpResponse
import io.ktor.client.statement.HttpStatement
import io.ktor.client.statement.bodyAsChannel
import io.ktor.http.HttpHeaders
import io.ktor.http.contentType
import io.ktor.http.isSuccess
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.isActive
const val HEADER_LAST_EVENT_ID = "Last-Event-Id"
typealias Milliseconds = Long
typealias HeadersProvider = suspend (EventId?) -> Map<String, String>
typealias QueryParamsProvider = suspend (EventId?) -> Map<String, String>
class UnauthorizedError : Throwable()
class NotEventStreamError : Throwable()
fun HttpClient.readSse(
url: String,
headersProvider: HeadersProvider = { it?.let { mapOf(HEADER_LAST_EVENT_ID to it) } ?: emptyMap() },
queryParamsProvider: QueryParamsProvider = { emptyMap() },
defaultReconnectDelayMillis: Milliseconds = 3000L
): Flow<SseEvent> {
var reconnectDelay: Milliseconds = defaultReconnectDelayMillis
var lastEventId: String? = null
return flow {
coroutineScope {
while (isActive) {
val customHeaders = headersProvider(lastEventId)
val queryParams = queryParamsProvider(lastEventId)
prepareRequest(
url = url,
headers = customHeaders,
queryParams = queryParams
).execute { response ->
if (!response.status.isSuccess()) {
throw UnauthorizedError()
}
if (!response.isEventStream()) {
throw NotEventStreamError()
}
response.bodyAsChannel()
.readSse(
onSseEvent = { sseEvent ->
lastEventId = sseEvent.id
emit(sseEvent)
},
onRetryChanged = {
reconnectDelay = it
}
)
}
delay(reconnectDelay)
}
}
}
}
private suspend fun HttpClient.prepareRequest(
url: String,
headers: Map<String, String> = emptyMap(),
queryParams: Map<String, String> = emptyMap()
): HttpStatement =
prepareGet(url) {
headers {
append(HttpHeaders.Accept, "text/event-stream")
append(HttpHeaders.CacheControl, "no-cache")
append(HttpHeaders.Connection, "keep-alive")
headers.forEach { (key, value) -> append(key, value) }
}
queryParams.forEach { (key, value) -> addOrReplaceParameter(key, value) }
}
private fun HttpResponse.isEventStream(): Boolean {
val contentType = contentType() ?: return false
return contentType.contentType == "text" && contentType.contentSubtype == "event-stream"
}
private fun HttpRequestBuilder.addOrReplaceParameter(key: String, value: String?): Unit =
value?.let {
url.parameters.remove(key)
url.parameters.append(key, it)
} ?: Unit
// package
typealias EventType = String
typealias EventData = String
typealias EventId = String
data class SseEvent(val id: EventId? = null, val event: EventType? = null, val data: EventData = "")
// package
inline fun parseSseLine(
line: String?,
onSseRawEvent: (SseRawEvent) -> (Unit)
) {
val parts = line.takeIf { !it.isNullOrBlank() }?.split(":", limit = 2)
val field = parts?.getOrNull(0)?.trim()
val value = parts?.getOrNull(1)?.trim().orEmpty()
onSseRawEvent(
when (field) {
null -> SseRawEvent.End
"" -> SseRawEvent.Comment(value)
"id" -> SseRawEvent.Id(value)
"data" -> SseRawEvent.Data(value)
"event" -> SseRawEvent.Event(value)
"retry" -> value.toLongOrNull()?.takeIf { it > 0 }?.let { SseRawEvent.Retry(it) }
?: SseRawEvent.Error(SseRawError.InvalidReconnectionTime(value))
else -> SseRawEvent.Error(SseRawError.InvalidField(field, value))
}
)
}
sealed interface SseRawEvent {
data class Id(val value: EventId) : SseRawEvent
data class Event(val value: EventType) : SseRawEvent
data class Data(val value: EventData) : SseRawEvent
data class Comment(val value: String) : SseRawEvent
data class Retry(val value: Milliseconds) : SseRawEvent
object End : SseRawEvent
data class Error(val error: SseRawError) : SseRawEvent
}
sealed class SseRawError : Throwable() {
data class InvalidField(val field: String, val value: String?) : SseRawError()
data class InvalidReconnectionTime(val value: String) : SseRawError()
}
// package
import io.ktor.utils.io.ByteReadChannel
import io.ktor.utils.io.readUTF8Line
suspend inline fun ByteReadChannel.readSse(
onSseEvent: (SseEvent) -> (Unit),
onRetryChanged: (Milliseconds) -> (Unit)
) {
var id: EventId? = null
var event: EventType? = null
var data: EventData? = null
while (!isClosedForRead) {
parseSseLine(
line = readUTF8Line(),
onSseRawEvent = { sseRawEvent ->
when (sseRawEvent) {
SseRawEvent.End -> {
if (data != null) {
onSseEvent(SseEvent(id, event, data!!))
id = null
event = null
data = null
} else {
// do nothing - maybe it's end after comment
}
}
is SseRawEvent.Id -> id = sseRawEvent.value
is SseRawEvent.Event -> event = sseRawEvent.value
is SseRawEvent.Data -> data = sseRawEvent.value
is SseRawEvent.Comment -> {
// do nothing
}
is SseRawEvent.Error -> {
// do nothing for now
}
is SseRawEvent.Retry -> onRetryChanged(sseRawEvent.value)
}
}
)
}
}
@pschichtel
Copy link

@JurajBegovac would it be an option to give this code a license?

@JurajBegovac
Copy link
Author

@pschichtel Why do you need a license, and do you intend to use the code? Please note that this code isn't anything special and isn't a library.

@pschichtel
Copy link

@JurajBegovac yes, I have a use-case in the context of the openai streaming API. The code is special in the sense, that is seems to be the one available complete-enough implementation of SSE for ktor. Currently there is no license defined for the code here, which implies that only you as the author may use it. Giving it a license (MIT, Apache, Mozilla, ...) would make reuse easier.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment