diff --git a/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachHandler.kt b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachHandler.kt index c797dad..179884b 100644 --- a/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachHandler.kt +++ b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachHandler.kt @@ -3,19 +3,11 @@ package com.jorisjonkers.personalstack.agentgateway.ws import com.jorisjonkers.personalstack.agentgateway.config.GatewayProperties import com.jorisjonkers.personalstack.agentgateway.observability.AgentGatewayTelemetry import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAgentKindLabel -import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAttachTelemetry import com.jorisjonkers.personalstack.agentgateway.observability.GatewayFailureReasonLabel import com.jorisjonkers.personalstack.agentgateway.observability.GatewayModeLabel -import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOperationLabel -import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOperationTelemetry import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOutcomeLabel -import com.jorisjonkers.personalstack.agentgateway.observability.GatewayReplayTelemetry import com.jorisjonkers.personalstack.agentgateway.tmux.AgentSessionManager -import com.jorisjonkers.personalstack.agentgateway.tmux.LogTailer import com.jorisjonkers.personalstack.agentgateway.tmux.TranscriptStore -import com.jorisjonkers.personalstack.agentgateway.tmux.TranscriptTailer -import com.jorisjonkers.personalstack.agentgateway.tmux.TranscriptTailerOptions -import io.micrometer.observation.Observation import io.micrometer.observation.ObservationRegistry import jakarta.annotation.PreDestroy import org.slf4j.LoggerFactory @@ -25,10 +17,6 @@ import org.springframework.web.socket.TextMessage import org.springframework.web.socket.WebSocketSession import org.springframework.web.socket.handler.TextWebSocketHandler import tools.jackson.databind.ObjectMapper -import java.io.IOException -import java.net.URLDecoder -import java.nio.charset.StandardCharsets -import java.time.Duration import java.time.Instant import java.util.concurrent.ConcurrentHashMap @@ -43,17 +31,22 @@ import java.util.concurrent.ConcurrentHashMap * Code doesn't ripple into the runner image. */ @Component -@Suppress("TooManyFunctions") class AgentAttachHandler( private val sessions: AgentSessionManager, - private val mapper: ObjectMapper, - private val props: GatewayProperties, - private val transcriptStore: TranscriptStore, - private val telemetry: AgentGatewayTelemetry = AgentGatewayTelemetry.NOOP, - private val observationRegistry: ObservationRegistry = ObservationRegistry.NOOP, + mapper: ObjectMapper, + props: GatewayProperties, + transcriptStore: TranscriptStore, + telemetry: AgentGatewayTelemetry = AgentGatewayTelemetry.NOOP, + observationRegistry: ObservationRegistry = ObservationRegistry.NOOP, ) : TextWebSocketHandler() { private val log = LoggerFactory.getLogger(AgentAttachHandler::class.java) private val tailers = ConcurrentHashMap() + private val sender = AgentWebSocketSender(mapper) + private val telemetryRecorder = AgentAttachTelemetryRecorder(telemetry, observationRegistry) + private val durableAttach = + AgentDurableAttachHandler(props, transcriptStore, sender, telemetryRecorder, tailers, observationRegistry) + private val liveAttach = AgentLiveAttachHandler(sessions, props, sender, telemetryRecorder, tailers) + private val inputHandler = AgentInputHandler(sessions, mapper) @PreDestroy fun shutdown() { @@ -65,517 +58,75 @@ class AgentAttachHandler( } } - @Suppress("LongMethod", "ReturnCount") override fun afterConnectionEstablished(session: WebSocketSession) { val startedAt = Instant.now() - val query = queryOf(session) - val requestedMode = requestedModeOf(query) - val agentId = - agentIdOf(session) ?: run { - recordAttachTerminal( - kind = GatewayAgentKindLabel.OTHER, - mode = requestedMode, - outcome = GatewayOutcomeLabel.FAILURE, - reason = GatewayFailureReasonLabel.INVALID_REQUEST, - startedAt = startedAt, + val query = AgentAttachQuery.parse(session) + val requestedMode = AgentAttachQuery.requestedModeOf(query) + val agentId = agentIdOf(session) + val agent = agentId?.let(sessions::get) + + when { + agentId == null -> + closeRejectedAttach( + session, + requestedMode, + GatewayFailureReasonLabel.INVALID_REQUEST, + "missing agentId", + startedAt, ) - session.close(CloseStatus.BAD_DATA.withReason("missing agentId")) - return - } - if (query.malformed) { - recordAttachTerminal( - kind = GatewayAgentKindLabel.OTHER, - mode = requestedMode, - outcome = GatewayOutcomeLabel.FAILURE, - reason = GatewayFailureReasonLabel.INVALID_REQUEST, - startedAt = startedAt, - ) - session.close(CloseStatus.BAD_DATA.withReason("malformed query")) - return - } - val agent = - sessions.get(agentId) ?: run { - recordAttachTerminal( - kind = GatewayAgentKindLabel.OTHER, - mode = requestedMode, - outcome = GatewayOutcomeLabel.FAILURE, - reason = GatewayFailureReasonLabel.NOT_FOUND, - startedAt = startedAt, - ) - session.close(CloseStatus.BAD_DATA.withReason("unknown agent")) - return - } - val kind = GatewayAgentKindLabel.fromRaw(agent.kind.name) - log.info("ws attach to agent {} (tmux={})", agentId, agent.tmuxSession) - - val stableSessionId = agent.stableSessionId - if (stableSessionId != null) { - attachDurable( - DurableAttachContext( - session = session, - stableSessionId = stableSessionId, - epoch = agent.epoch, - kind = kind, - requestedMode = requestedMode, - query = query.values, - startedAt = startedAt, - ), - ) - return - } - - runCatching { - // One current-screen snapshot (with ANSI) so the TUI renders - // immediately; the EOF-based tailer then streams only new bytes, - // so nothing in the snapshot gets replayed from the log. - // - // capture-pane joins pane rows with bare "\n"; the xterm runs - // with convertEol off (correct for the live PTY stream, which - // already emits CR), so each "\n" in the snapshot would drop a - // row without returning to column 0 — the banner staircases to - // the right. Normalise the snapshot's line ends to CRLF so it - // repaints flush-left. Existing CRLF is left intact. - val snapshot = - runCatching { sessions.captureWithEscapes(agentId) } - .getOrDefault("") - .replace("\r\n", "\n") - .replace("\n", "\r\n") - LogTailer.chunked(snapshot, LogTailer.MAX_CHUNK_CHARS) { sendOutput(session, it) } - - val tailer = - LogTailer(agent.logFile, intervalMs = props.tmux.tailIntervalMs) { text -> - sendOutput(session, text) - } - tailers[session.id] = tailer - tailer.start() - recordTailerStartup(kind, requestedMode, GatewayOutcomeLabel.SUCCESS, GatewayFailureReasonLabel.NONE) - recordAttachTerminal( - kind = kind, - mode = requestedMode, - outcome = GatewayOutcomeLabel.SUCCESS, - reason = GatewayFailureReasonLabel.NONE, - startedAt = startedAt, - ) - }.onFailure { - recordAttachTerminal( - kind = kind, - mode = requestedMode, - outcome = GatewayOutcomeLabel.FAILURE, - reason = failureReasonLabel(it), - startedAt = startedAt, - ) - runCatching { session.close(CloseStatus.SERVER_ERROR.withReason("attach failed")) } - .onFailure { closeError -> log.warn("closing failed websocket attach failed: {}", closeError.message) } - } - } - - // Output is relayed as small `{"output": "..."}` frames; LogTailer - // already bounds each chunk so a JSON frame stays under the default - // WebSocket buffer and nothing is buffered whole on the heap. - private fun sendOutput( - session: WebSocketSession, - text: String, - ) { - if (text.isEmpty() || !session.isOpen) return - sendJson(session, mapOf("output" to text)) - } - // Durable attach sends control frames before starting the transcript tailer. - @Suppress("CyclomaticComplexMethod", "LongMethod", "ReturnCount") - private fun attachDurable(context: DurableAttachContext) { - val session = context.session - val stableSessionId = context.stableSessionId - val epoch = context.epoch - val kind = context.kind - val requestedMode = context.requestedMode - val query = context.query - val startedAt = context.startedAt - val metadata = - runCatching { transcriptStore.recoverMetadata(stableSessionId) } - .getOrElse { - val reason = failureReasonLabel(it) - log.warn("metadata recovery for transcript {} failed: {}", stableSessionId, it.message) - recordAttachTerminal(kind, requestedMode, GatewayOutcomeLabel.FAILURE, reason, startedAt) - runCatching { session.close(CloseStatus.SERVER_ERROR.withReason("metadata recovery failed")) } - .onFailure { closeError -> - log.warn("closing websocket after metadata recovery failure failed: {}", closeError.message) - } - return - } - val requestedOffset = parseOffset(query) - val requestedEpoch = parseEpoch(query) - if (requestedOffset.malformed || requestedEpoch.malformed) { - telemetry.recordAttachFailure( - GatewayAttachTelemetry( - kind = kind, - mode = requestedMode, - outcome = GatewayOutcomeLabel.FAILURE, - reason = GatewayFailureReasonLabel.INVALID_REQUEST, - ), - ) - } - val mode = query["mode"]?.uppercase() - val canResume = - requestedEpoch.value == epoch && - requestedOffset.value != null && - requestedOffset.value in metadata.logicalStart..metadata.logicalEnd && - metadata.logicalEnd - requestedOffset.value <= MAX_RESUME_REPLAY_BYTES - val resume = (mode == "RESUME" || mode == null) && canResume - // Cold/snapshot attach used to replay from logicalStart — the entire - // retained transcript (up to the ~50MB cap). For a TUI that repaints - // the whole pane continuously that is megabytes of escape sequences - // the browser must parse on the main thread before it can even echo a - // keystroke, so perceived lag scaled with session age. Bound the cold - // replay to a recent tail window instead: a screenful of repaints - // lives well inside it, and attach cost becomes constant regardless of - // how long the session has run. (Resume from a live client's own - // offset is unchanged, but is itself capped above so a long-absent - // client falls back to this bounded snapshot rather than a huge gap.) - val coldStart = maxOf(metadata.logicalStart, metadata.logicalEnd - MAX_COLD_REPLAY_BYTES) - val replayStart = if (resume) requestedOffset.value else coldStart - val control = if (resume) "RESUME" else "SNAPSHOT" - - runCatching { - sendJson( - session, - mapOf( - "control" to control, - "stableSessionId" to stableSessionId, - "epoch" to epoch, - "start" to metadata.logicalStart, - "end" to metadata.logicalEnd, - ), - requireOpen = true, - ) - if (!resume) sendJson(session, mapOf("reset" to true), requireOpen = true) - if (metadata.logicalStart > 0) sendJson(session, mapOf("trim" to metadata.logicalStart), requireOpen = true) - sendJson(session, mapOf("cursor" to replayStart), requireOpen = true) - }.onFailure { - recordAttachTerminal(kind, requestedMode, GatewayOutcomeLabel.FAILURE, failureReasonLabel(it), startedAt) - runCatching { session.close(CloseStatus.SERVER_ERROR.withReason("websocket send failed")) } - .onFailure { closeError -> log.warn("closing failed websocket attach failed: {}", closeError.message) } - return - } + query.malformed -> + closeRejectedAttach( + session, + requestedMode, + GatewayFailureReasonLabel.INVALID_REQUEST, + "malformed query", + startedAt, + ) - var replaySendFailure = false + agent == null -> + closeRejectedAttach( + session, + requestedMode, + GatewayFailureReasonLabel.NOT_FOUND, + "unknown agent", + startedAt, + ) - val tailer = - TranscriptTailer( - store = transcriptStore, - stableSessionId = stableSessionId, - startOffset = replayStart, - onText = { frame -> - runCatching { - sendJson(session, mapOf("output" to frame.output, "off" to frame.off), requireOpen = true) - }.onFailure { replaySendFailure = true } - .getOrThrow() - }, - options = - TranscriptTailerOptions( - intervalMs = props.tmux.tailIntervalMs, - onTrim = { sendJson(session, mapOf("trim" to it, "cursor" to it), requireOpen = true) }, - observationRegistry = observationRegistry, + agent.stableSessionId != null -> + durableAttach.attach( + DurableAttachContext( + session = session, + stableSessionId = agent.stableSessionId, + epoch = agent.epoch, + kind = GatewayAgentKindLabel.fromRaw(agent.kind.name), + requestedMode = requestedMode, + query = query.values, + startedAt = startedAt, ), - ) - tailers[session.id] = tailer - val replay = tailer.replayAvailable() - val replayReason = - if (replaySendFailure) GatewayFailureReasonLabel.IO_ERROR else failureReasonLabel(replay.failureReason) - recordReplay(replay.bytes, replay.success, replayReason) - if (!replay.success) { - log.warn( - "replay of transcript {} failed before live tailing after {} bytes and {} frames: {}", - stableSessionId, - replay.bytes, - replay.frames, - replay.failureReason, - ) - telemetry.recordReplayFailure( - GatewayReplayTelemetry( - bytes = replay.bytes, - outcome = GatewayOutcomeLabel.FAILURE, - reason = replayReason, - ), - ) - recordReplayOperation( - kind = kind, - mode = requestedMode, - outcome = GatewayOutcomeLabel.FAILURE, - reason = replayReason, - ) - } else { - recordReplayOperation(kind, requestedMode, GatewayOutcomeLabel.SUCCESS, GatewayFailureReasonLabel.NONE) - } - if (replaySendFailure) { - recordAttachTerminal( - kind = kind, - mode = requestedMode, - outcome = GatewayOutcomeLabel.FAILURE, - reason = GatewayFailureReasonLabel.IO_ERROR, - startedAt = startedAt, - ) - runCatching { session.close(CloseStatus.SERVER_ERROR.withReason("websocket send failed")) } - .onFailure { closeError -> log.warn("closing failed websocket attach failed: {}", closeError.message) } - return - } - val replayEnd = - runCatching { transcriptStore.recoverMetadata(stableSessionId).logicalEnd } - .getOrElse { - val reason = failureReasonLabel(it) - log.warn( - "metadata recovery for transcript {} failed before replay completion: {}", - stableSessionId, - it.message, - ) - recordAttachTerminal(kind, requestedMode, GatewayOutcomeLabel.FAILURE, reason, startedAt) - runCatching { session.close(CloseStatus.SERVER_ERROR.withReason("metadata recovery failed")) } - .onFailure { closeError -> - log.warn("closing websocket after metadata recovery failure failed: {}", closeError.message) - } - return - } - runCatching { - sendJson( - session, - mapOf( - "control" to "REPLAY_COMPLETE", - "cursor" to replayEnd, - ), - requireOpen = true, - ) - }.onFailure { - recordAttachTerminal(kind, requestedMode, GatewayOutcomeLabel.FAILURE, failureReasonLabel(it), startedAt) - runCatching { session.close(CloseStatus.SERVER_ERROR.withReason("websocket send failed")) } - .onFailure { closeError -> log.warn("closing failed websocket attach failed: {}", closeError.message) } - return - } - val start = tailer.start() - val startReason = - if (start.success) GatewayFailureReasonLabel.NONE else failureReasonLabel(start.failureReason) - recordTailerStartup( - kind, - requestedMode, - if (start.success) GatewayOutcomeLabel.SUCCESS else GatewayOutcomeLabel.FAILURE, - startReason, - ) - recordAttachTerminal( - kind = kind, - mode = requestedMode, - outcome = if (start.success) GatewayOutcomeLabel.SUCCESS else GatewayOutcomeLabel.FAILURE, - reason = startReason, - startedAt = startedAt, - ) - if (!start.success) { - runCatching { session.close(CloseStatus.SERVER_ERROR.withReason("tailer startup failed")) } - .onFailure { closeError -> log.warn("closing failed websocket attach failed: {}", closeError.message) } - } - } - - private fun sendJson( - session: WebSocketSession, - payload: Map, - requireOpen: Boolean = false, - ) { - if (!session.isOpen) { - if (requireOpen) throw IOException("websocket session is closed") - return - } - val msg = mapper.writeValueAsString(payload) - synchronized(session) { session.sendMessage(TextMessage(msg)) } - } + ) - private fun queryOf(session: WebSocketSession): QueryParseResult { - val raw = session.uri?.rawQuery ?: return QueryParseResult() - val values = mutableMapOf() - var malformed = false - raw - .split('&') - .filter { it.isNotBlank() } - .forEach { part -> - val pieces = part.split('=', limit = 2) - val key = - runCatching { decodeQuery(pieces[0]) } - .getOrElse { - malformed = true - return@forEach - }.takeIf { it.isNotBlank() } - ?: return@forEach - val value = - if (pieces.size == 2) { - runCatching { decodeQuery(pieces[1]) } - .getOrElse { - malformed = true - return@forEach - } - } else { - "" - } - values[key] = value + else -> { + log.info("ws attach to agent {} (tmux={})", agentId, agent.tmuxSession) + liveAttach.attach( + LiveAttachContext( + session = session, + agentId = agentId, + agent = agent, + kind = GatewayAgentKindLabel.fromRaw(agent.kind.name), + requestedMode = requestedMode, + startedAt = startedAt, + ), + ) } - return QueryParseResult(values, malformed) - } - - private fun parseOffset(query: Map): ParsedLong { - val raw = query["offset"] ?: query["cursor"] ?: query["off"] ?: return ParsedLong() - return ParsedLong(value = raw.toLongOrNull(), malformed = raw.toLongOrNull() == null) - } - - private fun parseEpoch(query: Map): ParsedLong { - val raw = query["epoch"] ?: return ParsedLong() - return ParsedLong(value = raw.toLongOrNull(), malformed = raw.toLongOrNull() == null) - } - - private fun requestedModeOf(query: QueryParseResult): GatewayModeLabel { - val mode = query.values["mode"]?.uppercase() - return when { - mode == "SNAPSHOT" -> GatewayModeLabel.SNAPSHOT - mode == "RESUME" -> GatewayModeLabel.RESUME - query.values.keys.any { it == "offset" || it == "cursor" || it == "off" } -> GatewayModeLabel.RESUME - else -> GatewayModeLabel.SNAPSHOT } } - private fun decodeQuery(value: String): String = URLDecoder.decode(value, StandardCharsets.UTF_8) - - private fun recordAttachTerminal( - kind: GatewayAgentKindLabel, - mode: GatewayModeLabel, - outcome: GatewayOutcomeLabel, - reason: GatewayFailureReasonLabel, - startedAt: Instant, - ) { - val observation = - Observation - .start("agent.gateway.session.attach", observationRegistry) - .lowCardinalityKeyValue("kind", kind.label) - .lowCardinalityKeyValue("mode", mode.label) - .lowCardinalityKeyValue("outcome", outcome.label) - .lowCardinalityKeyValue("reason", reason.label) - try { - val event = GatewayAttachTelemetry(kind = kind, mode = mode, outcome = outcome, reason = reason) - telemetry.recordAttachAttempt(event) - if (outcome == GatewayOutcomeLabel.FAILURE) telemetry.recordAttachFailure(event) - telemetry.recordOperation( - GatewayOperationTelemetry( - operation = GatewayOperationLabel.ATTACH, - kind = kind, - mode = mode, - outcome = outcome, - reason = reason, - duration = Duration.between(startedAt, Instant.now()), - ), - ) - } finally { - observation.stop() - } - } - - private fun recordReplay( - bytes: Long, - success: Boolean, - failureReason: GatewayFailureReasonLabel, - ) { - val outcome = if (success) GatewayOutcomeLabel.SUCCESS else GatewayOutcomeLabel.FAILURE - val reason = if (success) GatewayFailureReasonLabel.NONE else failureReason - val observation = - Observation - .start("agent.gateway.replay", observationRegistry) - .lowCardinalityKeyValue("outcome", outcome.label) - .lowCardinalityKeyValue("reason", reason.label) - try { - telemetry.recordReplay( - GatewayReplayTelemetry( - bytes = bytes, - outcome = outcome, - reason = reason, - ), - ) - } finally { - observation.stop() - } - } - - private fun recordReplayOperation( - kind: GatewayAgentKindLabel, - mode: GatewayModeLabel, - outcome: GatewayOutcomeLabel, - reason: GatewayFailureReasonLabel, - ) { - telemetry.recordOperation( - GatewayOperationTelemetry( - operation = GatewayOperationLabel.REPLAY, - kind = kind, - mode = mode, - outcome = outcome, - reason = reason, - duration = Duration.ZERO, - ), - ) - } - - private fun recordTailerStartup( - kind: GatewayAgentKindLabel, - mode: GatewayModeLabel, - outcome: GatewayOutcomeLabel, - reason: GatewayFailureReasonLabel, - ) { - telemetry.recordOperation( - GatewayOperationTelemetry( - operation = GatewayOperationLabel.REPLAY, - kind = kind, - mode = mode, - outcome = outcome, - reason = reason, - duration = Duration.ZERO, - ), - ) - } - - private fun failureReasonLabel(error: Throwable): GatewayFailureReasonLabel = - when (error) { - is IOException -> GatewayFailureReasonLabel.IO_ERROR - is IllegalArgumentException -> GatewayFailureReasonLabel.INVALID_REQUEST - is SecurityException -> GatewayFailureReasonLabel.PERMISSION_DENIED - else -> GatewayFailureReasonLabel.UNKNOWN - } - - private fun failureReasonLabel(reason: String?): GatewayFailureReasonLabel { - val rawReason = reason - return GatewayFailureReasonLabel.fromRaw(rawReason) - } - override fun handleTextMessage( session: WebSocketSession, message: TextMessage, ) { - val agentId = agentIdOf(session) ?: return - val payload = payloadOf(message) ?: return - val resize = payload["resize"] as? Map<*, *> - if (resize != null) { - handleResize(agentId, resize) - } else { - (payload["input"] as? String)?.let { input -> - val enter = payload["enter"] as? Boolean ?: true - sessions.send(agentId, input, enter) - } - } - } - - private fun payloadOf(message: TextMessage): Map<*, *>? = - runCatching { mapper.readValue(message.payload, Map::class.java) } - .getOrElse { - log.warn("bad ws payload: {}", message.payload.take(PAYLOAD_PREVIEW_CHARS)) - null - } - - private fun handleResize( - agentId: String, - resize: Map<*, *>, - ) { - val cols = (resize["cols"] as? Number)?.toInt() ?: return - val rows = (resize["rows"] as? Number)?.toInt() ?: return - runCatching { sessions.resize(agentId, cols, rows) } - .onFailure { log.warn("resize of agent {} failed: {}", agentId, it.message) } + inputHandler.handleTextMessage(session, message) } override fun afterConnectionClosed( @@ -588,46 +139,25 @@ class AgentAttachHandler( } } - private fun agentIdOf(session: WebSocketSession): String? { - val path = session.uri?.path - return path - ?.let { Regex("/ws/agents/([^/]+)/attach").find(it) } - ?.groupValues - ?.get(1) + private fun closeRejectedAttach( + session: WebSocketSession, + requestedMode: GatewayModeLabel, + reason: GatewayFailureReasonLabel, + closeReason: String, + startedAt: Instant, + ) { + telemetryRecorder.recordAttachTerminal( + kind = GatewayAgentKindLabel.OTHER, + mode = requestedMode, + outcome = GatewayOutcomeLabel.FAILURE, + reason = reason, + startedAt = startedAt, + ) + session.close(CloseStatus.BAD_DATA.withReason(closeReason)) } - private data class QueryParseResult( - val values: Map = emptyMap(), - val malformed: Boolean = false, - ) - - private data class ParsedLong( - val value: Long? = null, - val malformed: Boolean = false, - ) - - private data class DurableAttachContext( - val session: WebSocketSession, - val stableSessionId: String, - val epoch: Long, - val kind: GatewayAgentKindLabel, - val requestedMode: GatewayModeLabel, - val query: Map, - val startedAt: Instant, - ) - companion object { - // Upper bound on bytes replayed into the browser on a cold/snapshot - // attach. A full-screen TUI repaint is tens of KiB, so 512 KiB always - // contains at least one complete repaint while keeping main-thread - // parse cost constant instead of O(session history). - internal const val MAX_COLD_REPLAY_BYTES = 512L * 1024L - - // A resuming client whose offset is further than this behind the live - // end has effectively been away too long to "catch up" cheaply; treat - // it as cold and send the bounded snapshot instead of replaying the - // whole gap. - internal const val MAX_RESUME_REPLAY_BYTES = 512L * 1024L - private const val PAYLOAD_PREVIEW_CHARS = 120 + internal const val MAX_COLD_REPLAY_BYTES = AgentAttachLimits.MAX_COLD_REPLAY_BYTES + internal const val MAX_RESUME_REPLAY_BYTES = AgentAttachLimits.MAX_RESUME_REPLAY_BYTES } } diff --git a/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachTelemetryRecorder.kt b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachTelemetryRecorder.kt new file mode 100644 index 0000000..944c3c6 --- /dev/null +++ b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachTelemetryRecorder.kt @@ -0,0 +1,148 @@ +package com.jorisjonkers.personalstack.agentgateway.ws + +import com.jorisjonkers.personalstack.agentgateway.observability.AgentGatewayTelemetry +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAgentKindLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAttachTelemetry +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayFailureReasonLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayModeLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOperationLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOperationTelemetry +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOutcomeLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayReplayTelemetry +import io.micrometer.observation.Observation +import io.micrometer.observation.ObservationRegistry +import java.io.IOException +import java.time.Duration +import java.time.Instant + +internal class AgentAttachTelemetryRecorder( + private val telemetry: AgentGatewayTelemetry, + private val observationRegistry: ObservationRegistry, +) { + fun recordAttachTerminal( + kind: GatewayAgentKindLabel, + mode: GatewayModeLabel, + outcome: GatewayOutcomeLabel, + reason: GatewayFailureReasonLabel, + startedAt: Instant, + ) { + val observation = + Observation + .start("agent.gateway.session.attach", observationRegistry) + .lowCardinalityKeyValue("kind", kind.label) + .lowCardinalityKeyValue("mode", mode.label) + .lowCardinalityKeyValue("outcome", outcome.label) + .lowCardinalityKeyValue("reason", reason.label) + try { + val event = GatewayAttachTelemetry(kind = kind, mode = mode, outcome = outcome, reason = reason) + telemetry.recordAttachAttempt(event) + if (outcome == GatewayOutcomeLabel.FAILURE) telemetry.recordAttachFailure(event) + telemetry.recordOperation( + GatewayOperationTelemetry( + operation = GatewayOperationLabel.ATTACH, + kind = kind, + mode = mode, + outcome = outcome, + reason = reason, + duration = Duration.between(startedAt, Instant.now()), + ), + ) + } finally { + observation.stop() + } + } + + fun recordReplay( + bytes: Long, + success: Boolean, + failureReason: GatewayFailureReasonLabel, + ) { + val outcome = if (success) GatewayOutcomeLabel.SUCCESS else GatewayOutcomeLabel.FAILURE + val reason = if (success) GatewayFailureReasonLabel.NONE else failureReason + val observation = + Observation + .start("agent.gateway.replay", observationRegistry) + .lowCardinalityKeyValue("outcome", outcome.label) + .lowCardinalityKeyValue("reason", reason.label) + try { + telemetry.recordReplay(GatewayReplayTelemetry(bytes = bytes, outcome = outcome, reason = reason)) + } finally { + observation.stop() + } + } + + fun recordReplayFailure( + bytes: Long, + reason: GatewayFailureReasonLabel, + ) { + telemetry.recordReplayFailure( + GatewayReplayTelemetry( + bytes = bytes, + outcome = GatewayOutcomeLabel.FAILURE, + reason = reason, + ), + ) + } + + fun recordAttachFailure( + kind: GatewayAgentKindLabel, + mode: GatewayModeLabel, + reason: GatewayFailureReasonLabel, + ) { + telemetry.recordAttachFailure( + GatewayAttachTelemetry( + kind = kind, + mode = mode, + outcome = GatewayOutcomeLabel.FAILURE, + reason = reason, + ), + ) + } + + fun recordReplayOperation( + kind: GatewayAgentKindLabel, + mode: GatewayModeLabel, + outcome: GatewayOutcomeLabel, + reason: GatewayFailureReasonLabel, + ) { + recordOperation(GatewayOperationLabel.REPLAY, kind, mode, outcome, reason) + } + + fun recordTailerStartup( + kind: GatewayAgentKindLabel, + mode: GatewayModeLabel, + outcome: GatewayOutcomeLabel, + reason: GatewayFailureReasonLabel, + ) { + recordOperation(GatewayOperationLabel.REPLAY, kind, mode, outcome, reason) + } + + private fun recordOperation( + operation: GatewayOperationLabel, + kind: GatewayAgentKindLabel, + mode: GatewayModeLabel, + outcome: GatewayOutcomeLabel, + reason: GatewayFailureReasonLabel, + ) { + telemetry.recordOperation( + GatewayOperationTelemetry( + operation = operation, + kind = kind, + mode = mode, + outcome = outcome, + reason = reason, + duration = Duration.ZERO, + ), + ) + } +} + +internal fun failureReasonLabel(error: Throwable): GatewayFailureReasonLabel = + when (error) { + is IOException -> GatewayFailureReasonLabel.IO_ERROR + is IllegalArgumentException -> GatewayFailureReasonLabel.INVALID_REQUEST + is SecurityException -> GatewayFailureReasonLabel.PERMISSION_DENIED + else -> GatewayFailureReasonLabel.UNKNOWN + } + +internal fun failureReasonLabel(reason: String?): GatewayFailureReasonLabel = GatewayFailureReasonLabel.fromRaw(reason) diff --git a/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachWebSocketSupport.kt b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachWebSocketSupport.kt new file mode 100644 index 0000000..3ba34e0 --- /dev/null +++ b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachWebSocketSupport.kt @@ -0,0 +1,132 @@ +package com.jorisjonkers.personalstack.agentgateway.ws + +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayModeLabel +import org.slf4j.LoggerFactory +import org.springframework.web.socket.CloseStatus +import org.springframework.web.socket.TextMessage +import org.springframework.web.socket.WebSocketSession +import tools.jackson.databind.ObjectMapper +import java.io.IOException +import java.net.URLDecoder +import java.nio.charset.StandardCharsets + +internal object AgentAttachLimits { + // Upper bound on bytes replayed into the browser on a cold/snapshot + // attach. A full-screen TUI repaint is tens of KiB, so 512 KiB always + // contains at least one complete repaint while keeping main-thread + // parse cost constant instead of O(session history). + const val MAX_COLD_REPLAY_BYTES = 512L * 1024L + + // A resuming client whose offset is further than this behind the live + // end has effectively been away too long to "catch up" cheaply; treat + // it as cold and send the bounded snapshot instead of replaying the + // whole gap. + const val MAX_RESUME_REPLAY_BYTES = 512L * 1024L +} + +internal class AgentWebSocketSender( + private val mapper: ObjectMapper, +) { + private val log = LoggerFactory.getLogger(AgentWebSocketSender::class.java) + + fun sendOutput( + session: WebSocketSession, + text: String, + ) { + if (text.isEmpty() || !session.isOpen) return + sendJson(session, mapOf("output" to text)) + } + + fun sendJson( + session: WebSocketSession, + payload: Map, + requireOpen: Boolean = false, + ) { + if (!session.isOpen) { + if (requireOpen) throw IOException("websocket session is closed") + return + } + val msg = mapper.writeValueAsString(payload) + synchronized(session) { session.sendMessage(TextMessage(msg)) } + } + + fun closeServerError( + session: WebSocketSession, + reason: String, + ) { + runCatching { session.close(CloseStatus.SERVER_ERROR.withReason(reason)) } + .onFailure { log.warn("closing failed websocket attach failed: {}", it.message) } + } +} + +internal object AgentAttachQuery { + fun parse(session: WebSocketSession): QueryParseResult { + val raw = session.uri?.rawQuery ?: return QueryParseResult() + val values = mutableMapOf() + var malformed = false + raw + .split('&') + .filter { it.isNotBlank() } + .forEach { part -> + val pieces = part.split('=', limit = 2) + val key = decodePart(pieces[0], onMalformed = { malformed = true }) ?: return@forEach + val value = + if (pieces.size == 2) { + decodePart(pieces[1], onMalformed = { malformed = true }) ?: return@forEach + } else { + "" + } + values[key] = value + } + return QueryParseResult(values, malformed) + } + + fun requestedModeOf(query: QueryParseResult): GatewayModeLabel { + val mode = query.values["mode"]?.uppercase() + return when { + mode == "SNAPSHOT" -> GatewayModeLabel.SNAPSHOT + mode == "RESUME" -> GatewayModeLabel.RESUME + query.values.keys.any { it == "offset" || it == "cursor" || it == "off" } -> GatewayModeLabel.RESUME + else -> GatewayModeLabel.SNAPSHOT + } + } + + fun parseOffset(query: Map): ParsedLong { + val raw = query["offset"] ?: query["cursor"] ?: query["off"] ?: return ParsedLong() + val value = raw.toLongOrNull() + return ParsedLong(value = value, malformed = value == null) + } + + fun parseEpoch(query: Map): ParsedLong { + val raw = query["epoch"] ?: return ParsedLong() + val value = raw.toLongOrNull() + return ParsedLong(value = value, malformed = value == null) + } + + private fun decodePart( + value: String, + onMalformed: () -> Unit, + ): String? = + runCatching { URLDecoder.decode(value, StandardCharsets.UTF_8) } + .getOrElse { + onMalformed() + null + }?.takeIf { it.isNotBlank() } +} + +internal fun agentIdOf(session: WebSocketSession): String? = + session.uri + ?.path + ?.let { Regex("/ws/agents/([^/]+)/attach").find(it) } + ?.groupValues + ?.get(1) + +internal data class QueryParseResult( + val values: Map = emptyMap(), + val malformed: Boolean = false, +) + +internal data class ParsedLong( + val value: Long? = null, + val malformed: Boolean = false, +) diff --git a/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentDurableAttachHandler.kt b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentDurableAttachHandler.kt new file mode 100644 index 0000000..5b87abe --- /dev/null +++ b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentDurableAttachHandler.kt @@ -0,0 +1,301 @@ +package com.jorisjonkers.personalstack.agentgateway.ws + +import com.jorisjonkers.personalstack.agentgateway.config.GatewayProperties +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAgentKindLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayFailureReasonLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayModeLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOutcomeLabel +import com.jorisjonkers.personalstack.agentgateway.tmux.TranscriptStore +import com.jorisjonkers.personalstack.agentgateway.tmux.TranscriptTailer +import com.jorisjonkers.personalstack.agentgateway.tmux.TranscriptTailerOptions +import io.micrometer.observation.ObservationRegistry +import org.slf4j.LoggerFactory +import org.springframework.web.socket.WebSocketSession +import java.time.Instant +import java.util.concurrent.ConcurrentMap + +internal class AgentDurableAttachHandler( + private val props: GatewayProperties, + private val transcriptStore: TranscriptStore, + private val sender: AgentWebSocketSender, + private val telemetry: AgentAttachTelemetryRecorder, + private val tailers: ConcurrentMap, + private val observationRegistry: ObservationRegistry, +) { + private val log = LoggerFactory.getLogger(AgentDurableAttachHandler::class.java) + + fun attach(context: DurableAttachContext) { + val metadataResult = runCatching { transcriptStore.recoverMetadata(context.stableSessionId) } + metadataResult.onFailure { failAttach(context, failureReasonLabel(it), "metadata recovery failed") } + val metadata = metadataResult.getOrNull() + if (metadata != null) { + val request = DurableReplayRequest.from(context, metadata.logicalStart, metadata.logicalEnd) + if (request.malformed) recordInvalidReplayRequest(context) + if (sendReplayControl(context, request)) { + val tailer = createTailer(context, request.replayStart) + tailers[context.session.id] = tailer.delegate + val replaySendFailure = replayAvailable(context, tailer) + if (!replaySendFailure && sendReplayComplete(context)) startLiveTailing(context, tailer.delegate) + } + } + } + + private fun sendReplayControl( + context: DurableAttachContext, + request: DurableReplayRequest, + ): Boolean = + runCatching { + sender.sendJson( + context.session, + mapOf( + "control" to request.control, + "stableSessionId" to context.stableSessionId, + "epoch" to context.epoch, + "start" to request.logicalStart, + "end" to request.logicalEnd, + ), + requireOpen = true, + ) + if (!request.resume) sender.sendJson(context.session, mapOf("reset" to true), requireOpen = true) + if (request.logicalStart > 0) { + sender.sendJson(context.session, mapOf("trim" to request.logicalStart), requireOpen = true) + } + sender.sendJson(context.session, mapOf("cursor" to request.replayStart), requireOpen = true) + }.fold( + onSuccess = { true }, + onFailure = { + failAttach(context, failureReasonLabel(it), "websocket send failed") + false + }, + ) + + private fun createTailer( + context: DurableAttachContext, + replayStart: Long, + ): DurableTranscriptTailer { + var replaySendFailure = false + val tailer = + TranscriptTailer( + store = transcriptStore, + stableSessionId = context.stableSessionId, + startOffset = replayStart, + onText = { frame -> + runCatching { + sender.sendJson( + context.session, + mapOf("output" to frame.output, "off" to frame.off), + requireOpen = true, + ) + }.onFailure { replaySendFailure = true } + .getOrThrow() + }, + options = + TranscriptTailerOptions( + intervalMs = props.tmux.tailIntervalMs, + onTrim = { + sender.sendJson( + context.session, + mapOf("trim" to it, "cursor" to it), + requireOpen = true, + ) + }, + observationRegistry = observationRegistry, + ), + ) + return DurableTranscriptTailer(tailer) { replaySendFailure } + } + + private fun replayAvailable( + context: DurableAttachContext, + tailer: DurableTranscriptTailer, + ): Boolean { + val replay = tailer.delegate.replayAvailable() + val sendFailed = tailer.replaySendFailed() + val replayReason = + if (sendFailed) { + GatewayFailureReasonLabel.IO_ERROR + } else { + failureReasonLabel(replay.failureReason) + } + telemetry.recordReplay(replay.bytes, replay.success, replayReason) + val sendFailure = + if (replay.success) { + recordReplaySuccess(context) + false + } else { + recordReplayFailure(context, replay.bytes, replay.frames, replay.failureReason, replayReason) + if (sendFailed) failAttach(context, replayReason, "websocket send failed") + sendFailed + } + return sendFailure + } + + private fun recordReplaySuccess(context: DurableAttachContext) { + telemetry.recordReplayOperation( + context.kind, + context.requestedMode, + GatewayOutcomeLabel.SUCCESS, + GatewayFailureReasonLabel.NONE, + ) + } + + private fun recordReplayFailure( + context: DurableAttachContext, + bytes: Long, + frames: Long, + failureReason: String?, + replayReason: GatewayFailureReasonLabel, + ) { + log.warn( + "replay of transcript {} failed before live tailing after {} bytes and {} frames: {}", + context.stableSessionId, + bytes, + frames, + failureReason, + ) + telemetry.recordReplayFailure(bytes, replayReason) + telemetry.recordReplayOperation( + context.kind, + context.requestedMode, + GatewayOutcomeLabel.FAILURE, + replayReason, + ) + } + + private fun sendReplayComplete(context: DurableAttachContext): Boolean { + val metadataResult = runCatching { transcriptStore.recoverMetadata(context.stableSessionId) } + metadataResult.onFailure { failAttach(context, failureReasonLabel(it), "metadata recovery failed") } + val replayEnd = metadataResult.getOrNull()?.logicalEnd + return replayEnd != null && sendReplayComplete(context, replayEnd) + } + + private fun sendReplayComplete( + context: DurableAttachContext, + replayEnd: Long, + ): Boolean = + runCatching { + sender.sendJson( + context.session, + mapOf("control" to "REPLAY_COMPLETE", "cursor" to replayEnd), + requireOpen = true, + ) + }.fold( + onSuccess = { true }, + onFailure = { + failAttach(context, failureReasonLabel(it), "websocket send failed") + false + }, + ) + + private fun startLiveTailing( + context: DurableAttachContext, + tailer: TranscriptTailer, + ) { + val start = tailer.start() + val startReason = + if (start.success) { + GatewayFailureReasonLabel.NONE + } else { + failureReasonLabel(start.failureReason) + } + val outcome = if (start.success) GatewayOutcomeLabel.SUCCESS else GatewayOutcomeLabel.FAILURE + telemetry.recordTailerStartup(context.kind, context.requestedMode, outcome, startReason) + telemetry.recordAttachTerminal(context.kind, context.requestedMode, outcome, startReason, context.startedAt) + if (!start.success) sender.closeServerError(context.session, "tailer startup failed") + } + + private fun recordInvalidReplayRequest(context: DurableAttachContext) { + telemetry.recordAttachFailure( + kind = context.kind, + mode = context.requestedMode, + reason = GatewayFailureReasonLabel.INVALID_REQUEST, + ) + } + + private fun failAttach( + context: DurableAttachContext, + reason: GatewayFailureReasonLabel, + closeReason: String, + ) { + telemetry.recordAttachTerminal( + context.kind, + context.requestedMode, + GatewayOutcomeLabel.FAILURE, + reason, + context.startedAt, + ) + sender.closeServerError(context.session, closeReason) + } +} + +private data class DurableTranscriptTailer( + val delegate: TranscriptTailer, + val replaySendFailed: () -> Boolean, +) + +internal data class DurableAttachContext( + val session: WebSocketSession, + val stableSessionId: String, + val epoch: Long, + val kind: GatewayAgentKindLabel, + val requestedMode: GatewayModeLabel, + val query: Map, + val startedAt: Instant, +) + +private data class DurableReplayRequest( + val logicalStart: Long, + val logicalEnd: Long, + val replayStart: Long, + val resume: Boolean, + val malformed: Boolean, +) { + val control: String = if (resume) "RESUME" else "SNAPSHOT" + + companion object { + fun from( + context: DurableAttachContext, + logicalStart: Long, + logicalEnd: Long, + ): DurableReplayRequest { + val requestedOffset = AgentAttachQuery.parseOffset(context.query) + val requestedEpoch = AgentAttachQuery.parseEpoch(context.query) + val mode = context.query["mode"]?.uppercase() + val resumeRequest = DurableReplayResumeRequest(requestedEpoch, requestedOffset, mode) + val resumeWindow = DurableReplayResumeWindow(context.epoch, logicalStart, logicalEnd) + val resume = DurableReplayResumeDecision(resumeRequest, resumeWindow).canResume + val coldStart = maxOf(logicalStart, logicalEnd - AgentAttachLimits.MAX_COLD_REPLAY_BYTES) + return DurableReplayRequest( + logicalStart = logicalStart, + logicalEnd = logicalEnd, + replayStart = if (resume) requestedOffset.value ?: coldStart else coldStart, + resume = resume, + malformed = requestedOffset.malformed || requestedEpoch.malformed, + ) + } + } +} + +private data class DurableReplayResumeRequest( + val requestedEpoch: ParsedLong, + val requestedOffset: ParsedLong, + val mode: String?, +) + +private data class DurableReplayResumeWindow( + val epoch: Long, + val logicalStart: Long, + val logicalEnd: Long, +) + +private data class DurableReplayResumeDecision( + val request: DurableReplayResumeRequest, + val window: DurableReplayResumeWindow, +) { + val canResume: Boolean = + (request.mode == "RESUME" || request.mode == null) && + request.requestedEpoch.value == window.epoch && + request.requestedOffset.value != null && + request.requestedOffset.value in window.logicalStart..window.logicalEnd && + window.logicalEnd - request.requestedOffset.value <= AgentAttachLimits.MAX_RESUME_REPLAY_BYTES +} diff --git a/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentInputHandler.kt b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentInputHandler.kt new file mode 100644 index 0000000..e65da20 --- /dev/null +++ b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentInputHandler.kt @@ -0,0 +1,59 @@ +package com.jorisjonkers.personalstack.agentgateway.ws + +import com.jorisjonkers.personalstack.agentgateway.tmux.AgentSessionManager +import org.slf4j.LoggerFactory +import org.springframework.web.socket.TextMessage +import org.springframework.web.socket.WebSocketSession +import tools.jackson.databind.ObjectMapper + +internal class AgentInputHandler( + private val sessions: AgentSessionManager, + private val mapper: ObjectMapper, +) { + private val log = LoggerFactory.getLogger(AgentInputHandler::class.java) + + fun handleTextMessage( + session: WebSocketSession, + message: TextMessage, + ) { + val agentId = agentIdOf(session) ?: return + val payload = payloadOf(message) ?: return + val resize = payload["resize"] as? Map<*, *> + if (resize != null) { + handleResize(agentId, resize) + } else { + handleInput(agentId, payload) + } + } + + private fun payloadOf(message: TextMessage): Map<*, *>? = + runCatching { mapper.readValue(message.payload, Map::class.java) } + .getOrElse { + log.warn("bad ws payload: {}", message.payload.take(PAYLOAD_PREVIEW_CHARS)) + null + } + + private fun handleResize( + agentId: String, + resize: Map<*, *>, + ) { + val cols = (resize["cols"] as? Number)?.toInt() ?: return + val rows = (resize["rows"] as? Number)?.toInt() ?: return + runCatching { sessions.resize(agentId, cols, rows) } + .onFailure { log.warn("resize of agent {} failed: {}", agentId, it.message) } + } + + private fun handleInput( + agentId: String, + payload: Map<*, *>, + ) { + (payload["input"] as? String)?.let { input -> + val enter = payload["enter"] as? Boolean ?: true + sessions.send(agentId, input, enter) + } + } + + private companion object { + const val PAYLOAD_PREVIEW_CHARS = 120 + } +} diff --git a/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentLiveAttachHandler.kt b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentLiveAttachHandler.kt new file mode 100644 index 0000000..e8475c3 --- /dev/null +++ b/services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentLiveAttachHandler.kt @@ -0,0 +1,91 @@ +package com.jorisjonkers.personalstack.agentgateway.ws + +import com.jorisjonkers.personalstack.agentgateway.config.GatewayProperties +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAgentKindLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayFailureReasonLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayModeLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOutcomeLabel +import com.jorisjonkers.personalstack.agentgateway.tmux.AgentSession +import com.jorisjonkers.personalstack.agentgateway.tmux.AgentSessionManager +import com.jorisjonkers.personalstack.agentgateway.tmux.LogTailer +import org.slf4j.LoggerFactory +import org.springframework.web.socket.CloseStatus +import org.springframework.web.socket.WebSocketSession +import java.time.Instant +import java.util.concurrent.ConcurrentMap + +internal class AgentLiveAttachHandler( + private val sessions: AgentSessionManager, + private val props: GatewayProperties, + private val sender: AgentWebSocketSender, + private val telemetry: AgentAttachTelemetryRecorder, + private val tailers: ConcurrentMap, +) { + private val log = LoggerFactory.getLogger(AgentLiveAttachHandler::class.java) + + fun attach(context: LiveAttachContext) { + runCatching { + sendSnapshot(context.session, context.agentId) + startTailer(context) + telemetry.recordTailerStartup( + context.kind, + context.requestedMode, + GatewayOutcomeLabel.SUCCESS, + GatewayFailureReasonLabel.NONE, + ) + telemetry.recordAttachTerminal( + kind = context.kind, + mode = context.requestedMode, + outcome = GatewayOutcomeLabel.SUCCESS, + reason = GatewayFailureReasonLabel.NONE, + startedAt = context.startedAt, + ) + }.onFailure { error -> + telemetry.recordAttachTerminal( + kind = context.kind, + mode = context.requestedMode, + outcome = GatewayOutcomeLabel.FAILURE, + reason = failureReasonLabel(error), + startedAt = context.startedAt, + ) + closeFailedAttach(context.session) + } + } + + private fun sendSnapshot( + session: WebSocketSession, + agentId: String, + ) { + // capture-pane joins pane rows with bare "\n"; normalize to CRLF so + // xterm repaints the snapshot flush-left before live PTY bytes arrive. + val snapshot = + runCatching { sessions.captureWithEscapes(agentId) } + .getOrDefault("") + .replace("\r\n", "\n") + .replace("\n", "\r\n") + LogTailer.chunked(snapshot, LogTailer.MAX_CHUNK_CHARS) { sender.sendOutput(session, it) } + } + + private fun startTailer(context: LiveAttachContext) { + val tailer = + LogTailer(context.agent.logFile, intervalMs = props.tmux.tailIntervalMs) { text -> + sender.sendOutput(context.session, text) + } + tailers[context.session.id] = tailer + tailer.start() + } + + private fun closeFailedAttach(session: WebSocketSession) { + runCatching { session.close(CloseStatus.SERVER_ERROR.withReason("attach failed")) } + .onFailure { log.warn("closing failed websocket attach failed: {}", it.message) } + } +} + +internal data class LiveAttachContext( + val session: WebSocketSession, + val agentId: String, + val agent: AgentSession, + val kind: GatewayAgentKindLabel, + val requestedMode: GatewayModeLabel, + val startedAt: Instant, +) diff --git a/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachHandlerTest.kt b/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachHandlerTest.kt index fc82c82..b819d6a 100644 --- a/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachHandlerTest.kt +++ b/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachHandlerTest.kt @@ -2,14 +2,9 @@ package com.jorisjonkers.personalstack.agentgateway.ws import com.jorisjonkers.personalstack.agentgateway.config.GatewayProperties import com.jorisjonkers.personalstack.agentgateway.observability.AgentGatewayTelemetry -import com.jorisjonkers.personalstack.agentgateway.observability.GatewayActiveSessionsSample -import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAttachTelemetry import com.jorisjonkers.personalstack.agentgateway.observability.GatewayFailureReasonLabel import com.jorisjonkers.personalstack.agentgateway.observability.GatewayModeLabel -import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOperationTelemetry import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOutcomeLabel -import com.jorisjonkers.personalstack.agentgateway.observability.GatewayReplayTelemetry -import com.jorisjonkers.personalstack.agentgateway.observability.GatewayStorageTelemetrySample import com.jorisjonkers.personalstack.agentgateway.tmux.AgentKind import com.jorisjonkers.personalstack.agentgateway.tmux.AgentSession import com.jorisjonkers.personalstack.agentgateway.tmux.AgentSessionManager @@ -578,40 +573,6 @@ class AgentAttachHandlerTest { verify(exactly = 0) { sessions.send(any(), any(), any()) } } - private class RecordingTelemetry : AgentGatewayTelemetry { - val attachAttempts = CopyOnWriteArrayList() - val attachFailures = CopyOnWriteArrayList() - val replayEvents = CopyOnWriteArrayList() - val replayFailures = CopyOnWriteArrayList() - val operations = CopyOnWriteArrayList() - - override fun recordActiveSessions(sample: GatewayActiveSessionsSample) = Unit - - override fun recordOperation(event: GatewayOperationTelemetry) { - operations += event - } - - override fun recordAttachAttempt(event: GatewayAttachTelemetry) { - attachAttempts += event - } - - override fun recordAttachFailure(event: GatewayAttachTelemetry) { - attachFailures += event - } - - override fun recordReplay(event: GatewayReplayTelemetry) { - replayEvents += event - } - - override fun recordReplayFailure(event: GatewayReplayTelemetry) { - replayFailures += event - } - - override fun recordStorage(sample: GatewayStorageTelemetrySample) = Unit - - override fun recordStorageLimit(sample: GatewayStorageTelemetrySample) = Unit - } - private fun transcriptStore( tmp: Path, segmentBytes: Long = 1024 * 1024, diff --git a/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentInputHandlerTest.kt b/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentInputHandlerTest.kt new file mode 100644 index 0000000..f9b42e1 --- /dev/null +++ b/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentInputHandlerTest.kt @@ -0,0 +1,46 @@ +package com.jorisjonkers.personalstack.agentgateway.ws + +import com.jorisjonkers.personalstack.agentgateway.tmux.AgentSessionManager +import io.mockk.every +import io.mockk.mockk +import io.mockk.verify +import org.junit.jupiter.api.Test +import org.springframework.web.socket.TextMessage +import org.springframework.web.socket.WebSocketSession +import tools.jackson.databind.ObjectMapper +import java.net.URI + +class AgentInputHandlerTest { + private val sessions = mockk(relaxed = true) + private val handler = AgentInputHandler(sessions, ObjectMapper()) + + @Test + fun `resize frame routes to AgentSessionManager resize`() { + handler.handleTextMessage(wsSession("abc"), TextMessage("""{"resize":{"cols":120,"rows":40}}""")) + + verify { sessions.resize("abc", 120, 40) } + verify(exactly = 0) { sessions.send(any(), any(), any()) } + } + + @Test + fun `input frame routes to AgentSessionManager send`() { + handler.handleTextMessage(wsSession("abc"), TextMessage("""{"input":"ls\r","enter":false}""")) + + verify { sessions.send("abc", "ls\r", false) } + verify(exactly = 0) { sessions.resize(any(), any(), any()) } + } + + @Test + fun `malformed json is ignored`() { + handler.handleTextMessage(wsSession("abc"), TextMessage("not json at all")) + + verify(exactly = 0) { sessions.resize(any(), any(), any()) } + verify(exactly = 0) { sessions.send(any(), any(), any()) } + } + + private fun wsSession(id: String): WebSocketSession { + val ws = mockk(relaxed = true) + every { ws.uri } returns URI("ws://host/ws/agents/$id/attach") + return ws + } +} diff --git a/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentLiveAttachHandlerTest.kt b/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentLiveAttachHandlerTest.kt new file mode 100644 index 0000000..af9f2c1 --- /dev/null +++ b/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentLiveAttachHandlerTest.kt @@ -0,0 +1,122 @@ +package com.jorisjonkers.personalstack.agentgateway.ws + +import com.jorisjonkers.personalstack.agentgateway.config.GatewayProperties +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAgentKindLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayFailureReasonLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayModeLabel +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOutcomeLabel +import com.jorisjonkers.personalstack.agentgateway.tmux.AgentKind +import com.jorisjonkers.personalstack.agentgateway.tmux.AgentSession +import com.jorisjonkers.personalstack.agentgateway.tmux.AgentSessionManager +import io.micrometer.observation.ObservationRegistry +import io.mockk.every +import io.mockk.mockk +import io.mockk.slot +import io.mockk.verify +import org.assertj.core.api.Assertions.assertThat +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import org.springframework.web.socket.CloseStatus +import org.springframework.web.socket.TextMessage +import org.springframework.web.socket.WebSocketSession +import tools.jackson.databind.ObjectMapper +import java.io.IOException +import java.net.URI +import java.nio.file.Path +import java.time.Instant +import java.util.concurrent.ConcurrentHashMap + +class AgentLiveAttachHandlerTest { + private val sessions = mockk(relaxed = true) + private val props = + GatewayProperties( + workspaceRoot = "/workspace", + tmux = GatewayProperties.Tmux(socketName = "agent-gw", stateDir = "/tmp"), + cli = GatewayProperties.Cli(claude = "claude", codex = "codex"), + ) + + @Test + fun `live attach sends normalized snapshot and records success`( + @TempDir tmp: Path, + ) { + val telemetry = RecordingTelemetry() + val tailers = ConcurrentHashMap() + val handler = attachHandler(telemetry, tailers) + val ws = wsSession("abc") + val sent = slot() + every { sessions.captureWithEscapes("abc") } returns "row1\nrow2" + every { ws.sendMessage(capture(sent)) } returns Unit + + handler.attach(context(ws, agent(tmp))) + + verify { sessions.captureWithEscapes("abc") } + assertThat(sent.captured.payload).contains("row1\\r\\nrow2") + assertThat(tailers).containsKey("ws-abc") + assertThat(telemetry.attachAttempts.single().outcome).isEqualTo(GatewayOutcomeLabel.SUCCESS) + assertThat(telemetry.operations.map { it.reason }).contains(GatewayFailureReasonLabel.NONE) + + tailers.remove("ws-abc")?.close() + } + + @Test + fun `live attach closes websocket when snapshot send fails`( + @TempDir tmp: Path, + ) { + val telemetry = RecordingTelemetry() + val handler = attachHandler(telemetry, ConcurrentHashMap()) + val ws = wsSession("abc") + every { sessions.captureWithEscapes("abc") } returns "snapshot" + every { ws.sendMessage(any()) } throws IOException("send failed") + + handler.attach(context(ws, agent(tmp))) + + assertThat(telemetry.attachAttempts.single().outcome).isEqualTo(GatewayOutcomeLabel.FAILURE) + verify { ws.close(any()) } + } + + private fun attachHandler( + telemetry: RecordingTelemetry, + tailers: ConcurrentHashMap, + ): AgentLiveAttachHandler = + AgentLiveAttachHandler( + sessions = sessions, + props = props, + sender = AgentWebSocketSender(ObjectMapper()), + telemetry = AgentAttachTelemetryRecorder(telemetry, ObservationRegistry.NOOP), + tailers = tailers, + ) + + private fun context( + ws: WebSocketSession, + agent: AgentSession, + ): LiveAttachContext = + LiveAttachContext( + session = ws, + agentId = agent.id, + agent = agent, + kind = GatewayAgentKindLabel.SHELL, + requestedMode = GatewayModeLabel.SNAPSHOT, + startedAt = Instant.now(), + ) + + private fun agent(tmp: Path): AgentSession { + val logFile = tmp.resolve("agent-abc.log") + logFile.toFile().createNewFile() + return AgentSession( + id = "abc", + kind = AgentKind.SHELL, + tmuxSession = "agent-abc", + logFile = logFile, + cwd = "/workspace", + createdAt = Instant.now(), + ) + } + + private fun wsSession(id: String): WebSocketSession { + val ws = mockk(relaxed = true) + every { ws.uri } returns URI("ws://host/ws/agents/$id/attach") + every { ws.id } returns "ws-$id" + every { ws.isOpen } returns true + return ws + } +} diff --git a/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/RecordingTelemetry.kt b/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/RecordingTelemetry.kt new file mode 100644 index 0000000..e82eca1 --- /dev/null +++ b/services/agent-gateway/src/test/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/RecordingTelemetry.kt @@ -0,0 +1,43 @@ +package com.jorisjonkers.personalstack.agentgateway.ws + +import com.jorisjonkers.personalstack.agentgateway.observability.AgentGatewayTelemetry +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayActiveSessionsSample +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAttachTelemetry +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOperationTelemetry +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayReplayTelemetry +import com.jorisjonkers.personalstack.agentgateway.observability.GatewayStorageTelemetrySample +import java.util.concurrent.CopyOnWriteArrayList + +internal class RecordingTelemetry : AgentGatewayTelemetry { + val attachAttempts = CopyOnWriteArrayList() + val attachFailures = CopyOnWriteArrayList() + val replayEvents = CopyOnWriteArrayList() + val replayFailures = CopyOnWriteArrayList() + val operations = CopyOnWriteArrayList() + + override fun recordActiveSessions(sample: GatewayActiveSessionsSample) = Unit + + override fun recordOperation(event: GatewayOperationTelemetry) { + operations += event + } + + override fun recordAttachAttempt(event: GatewayAttachTelemetry) { + attachAttempts += event + } + + override fun recordAttachFailure(event: GatewayAttachTelemetry) { + attachFailures += event + } + + override fun recordReplay(event: GatewayReplayTelemetry) { + replayEvents += event + } + + override fun recordReplayFailure(event: GatewayReplayTelemetry) { + replayFailures += event + } + + override fun recordStorage(sample: GatewayStorageTelemetrySample) = Unit + + override fun recordStorageLimit(sample: GatewayStorageTelemetrySample) = Unit +}