Skip to content

Commit cc086a7

Browse files
jorisjonkers-dev-agents[bot]ExtraToast
jorisjonkers-dev-agents[bot]
and
ExtraToast
authored
refactor: split AgentAttachHandler into attach collaborators (#18)
* refactor: split AgentAttachHandler into attach collaborators The websocket coordinator held 21 functions and an 89-line afterConnectionEstablished behind three documented suppressions. Durable-attach, live-attach, input-handler, websocket-support, and telemetry-recorder collaborators take over; all three @Suppress annotations are gone and the collaborators are unit-tested. * fix: address detekt and ktlint findings from CI * fix: remaining ktlint style and test reference findings --------- Co-authored-by: JorisJonkers Agent <agents@jorisjonkers.dev>
1 parent 6b21045 commit cc086a7

10 files changed

Lines changed: 1023 additions & 590 deletions

File tree

‎services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachHandler.kt‎

Lines changed: 81 additions & 551 deletions
Large diffs are not rendered by default.

‎services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachTelemetryRecorder.kt‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
package com.jorisjonkers.personalstack.agentgateway.ws
2+
3+
import com.jorisjonkers.personalstack.agentgateway.observability.AgentGatewayTelemetry
4+
import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAgentKindLabel
5+
import com.jorisjonkers.personalstack.agentgateway.observability.GatewayAttachTelemetry
6+
import com.jorisjonkers.personalstack.agentgateway.observability.GatewayFailureReasonLabel
7+
import com.jorisjonkers.personalstack.agentgateway.observability.GatewayModeLabel
8+
import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOperationLabel
9+
import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOperationTelemetry
10+
import com.jorisjonkers.personalstack.agentgateway.observability.GatewayOutcomeLabel
11+
import com.jorisjonkers.personalstack.agentgateway.observability.GatewayReplayTelemetry
12+
import io.micrometer.observation.Observation
13+
import io.micrometer.observation.ObservationRegistry
14+
import java.io.IOException
15+
import java.time.Duration
16+
import java.time.Instant
17+
18+
internal class AgentAttachTelemetryRecorder(
19+
private val telemetry: AgentGatewayTelemetry,
20+
private val observationRegistry: ObservationRegistry,
21+
) {
22+
fun recordAttachTerminal(
23+
kind: GatewayAgentKindLabel,
24+
mode: GatewayModeLabel,
25+
outcome: GatewayOutcomeLabel,
26+
reason: GatewayFailureReasonLabel,
27+
startedAt: Instant,
28+
) {
29+
val observation =
30+
Observation
31+
.start("agent.gateway.session.attach", observationRegistry)
32+
.lowCardinalityKeyValue("kind", kind.label)
33+
.lowCardinalityKeyValue("mode", mode.label)
34+
.lowCardinalityKeyValue("outcome", outcome.label)
35+
.lowCardinalityKeyValue("reason", reason.label)
36+
try {
37+
val event = GatewayAttachTelemetry(kind = kind, mode = mode, outcome = outcome, reason = reason)
38+
telemetry.recordAttachAttempt(event)
39+
if (outcome == GatewayOutcomeLabel.FAILURE) telemetry.recordAttachFailure(event)
40+
telemetry.recordOperation(
41+
GatewayOperationTelemetry(
42+
operation = GatewayOperationLabel.ATTACH,
43+
kind = kind,
44+
mode = mode,
45+
outcome = outcome,
46+
reason = reason,
47+
duration = Duration.between(startedAt, Instant.now()),
48+
),
49+
)
50+
} finally {
51+
observation.stop()
52+
}
53+
}
54+
55+
fun recordReplay(
56+
bytes: Long,
57+
success: Boolean,
58+
failureReason: GatewayFailureReasonLabel,
59+
) {
60+
val outcome = if (success) GatewayOutcomeLabel.SUCCESS else GatewayOutcomeLabel.FAILURE
61+
val reason = if (success) GatewayFailureReasonLabel.NONE else failureReason
62+
val observation =
63+
Observation
64+
.start("agent.gateway.replay", observationRegistry)
65+
.lowCardinalityKeyValue("outcome", outcome.label)
66+
.lowCardinalityKeyValue("reason", reason.label)
67+
try {
68+
telemetry.recordReplay(GatewayReplayTelemetry(bytes = bytes, outcome = outcome, reason = reason))
69+
} finally {
70+
observation.stop()
71+
}
72+
}
73+
74+
fun recordReplayFailure(
75+
bytes: Long,
76+
reason: GatewayFailureReasonLabel,
77+
) {
78+
telemetry.recordReplayFailure(
79+
GatewayReplayTelemetry(
80+
bytes = bytes,
81+
outcome = GatewayOutcomeLabel.FAILURE,
82+
reason = reason,
83+
),
84+
)
85+
}
86+
87+
fun recordAttachFailure(
88+
kind: GatewayAgentKindLabel,
89+
mode: GatewayModeLabel,
90+
reason: GatewayFailureReasonLabel,
91+
) {
92+
telemetry.recordAttachFailure(
93+
GatewayAttachTelemetry(
94+
kind = kind,
95+
mode = mode,
96+
outcome = GatewayOutcomeLabel.FAILURE,
97+
reason = reason,
98+
),
99+
)
100+
}
101+
102+
fun recordReplayOperation(
103+
kind: GatewayAgentKindLabel,
104+
mode: GatewayModeLabel,
105+
outcome: GatewayOutcomeLabel,
106+
reason: GatewayFailureReasonLabel,
107+
) {
108+
recordOperation(GatewayOperationLabel.REPLAY, kind, mode, outcome, reason)
109+
}
110+
111+
fun recordTailerStartup(
112+
kind: GatewayAgentKindLabel,
113+
mode: GatewayModeLabel,
114+
outcome: GatewayOutcomeLabel,
115+
reason: GatewayFailureReasonLabel,
116+
) {
117+
recordOperation(GatewayOperationLabel.REPLAY, kind, mode, outcome, reason)
118+
}
119+
120+
private fun recordOperation(
121+
operation: GatewayOperationLabel,
122+
kind: GatewayAgentKindLabel,
123+
mode: GatewayModeLabel,
124+
outcome: GatewayOutcomeLabel,
125+
reason: GatewayFailureReasonLabel,
126+
) {
127+
telemetry.recordOperation(
128+
GatewayOperationTelemetry(
129+
operation = operation,
130+
kind = kind,
131+
mode = mode,
132+
outcome = outcome,
133+
reason = reason,
134+
duration = Duration.ZERO,
135+
),
136+
)
137+
}
138+
}
139+
140+
internal fun failureReasonLabel(error: Throwable): GatewayFailureReasonLabel =
141+
when (error) {
142+
is IOException -> GatewayFailureReasonLabel.IO_ERROR
143+
is IllegalArgumentException -> GatewayFailureReasonLabel.INVALID_REQUEST
144+
is SecurityException -> GatewayFailureReasonLabel.PERMISSION_DENIED
145+
else -> GatewayFailureReasonLabel.UNKNOWN
146+
}
147+
148+
internal fun failureReasonLabel(reason: String?): GatewayFailureReasonLabel = GatewayFailureReasonLabel.fromRaw(reason)

‎services/agent-gateway/src/main/kotlin/com/jorisjonkers/personalstack/agentgateway/ws/AgentAttachWebSocketSupport.kt‎

Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
package com.jorisjonkers.personalstack.agentgateway.ws
2+
3+
import com.jorisjonkers.personalstack.agentgateway.observability.GatewayModeLabel
4+
import org.slf4j.LoggerFactory
5+
import org.springframework.web.socket.CloseStatus
6+
import org.springframework.web.socket.TextMessage
7+
import org.springframework.web.socket.WebSocketSession
8+
import tools.jackson.databind.ObjectMapper
9+
import java.io.IOException
10+
import java.net.URLDecoder
11+
import java.nio.charset.StandardCharsets
12+
13+
internal object AgentAttachLimits {
14+
// Upper bound on bytes replayed into the browser on a cold/snapshot
15+
// attach. A full-screen TUI repaint is tens of KiB, so 512 KiB always
16+
// contains at least one complete repaint while keeping main-thread
17+
// parse cost constant instead of O(session history).
18+
const val MAX_COLD_REPLAY_BYTES = 512L * 1024L
19+
20+
// A resuming client whose offset is further than this behind the live
21+
// end has effectively been away too long to "catch up" cheaply; treat
22+
// it as cold and send the bounded snapshot instead of replaying the
23+
// whole gap.
24+
const val MAX_RESUME_REPLAY_BYTES = 512L * 1024L
25+
}
26+
27+
internal class AgentWebSocketSender(
28+
private val mapper: ObjectMapper,
29+
) {
30+
private val log = LoggerFactory.getLogger(AgentWebSocketSender::class.java)
31+
32+
fun sendOutput(
33+
session: WebSocketSession,
34+
text: String,
35+
) {
36+
if (text.isEmpty() || !session.isOpen) return
37+
sendJson(session, mapOf("output" to text))
38+
}
39+
40+
fun sendJson(
41+
session: WebSocketSession,
42+
payload: Map<String, Any?>,
43+
requireOpen: Boolean = false,
44+
) {
45+
if (!session.isOpen) {
46+
if (requireOpen) throw IOException("websocket session is closed")
47+
return
48+
}
49+
val msg = mapper.writeValueAsString(payload)
50+
synchronized(session) { session.sendMessage(TextMessage(msg)) }
51+
}
52+
53+
fun closeServerError(
54+
session: WebSocketSession,
55+
reason: String,
56+
) {
57+
runCatching { session.close(CloseStatus.SERVER_ERROR.withReason(reason)) }
58+
.onFailure { log.warn("closing failed websocket attach failed: {}", it.message) }
59+
}
60+
}
61+
62+
internal object AgentAttachQuery {
63+
fun parse(session: WebSocketSession): QueryParseResult {
64+
val raw = session.uri?.rawQuery ?: return QueryParseResult()
65+
val values = mutableMapOf<String, String>()
66+
var malformed = false
67+
raw
68+
.split('&')
69+
.filter { it.isNotBlank() }
70+
.forEach { part ->
71+
val pieces = part.split('=', limit = 2)
72+
val key = decodePart(pieces[0], onMalformed = { malformed = true }) ?: return@forEach
73+
val value =
74+
if (pieces.size == 2) {
75+
decodePart(pieces[1], onMalformed = { malformed = true }) ?: return@forEach
76+
} else {
77+
""
78+
}
79+
values[key] = value
80+
}
81+
return QueryParseResult(values, malformed)
82+
}
83+
84+
fun requestedModeOf(query: QueryParseResult): GatewayModeLabel {
85+
val mode = query.values["mode"]?.uppercase()
86+
return when {
87+
mode == "SNAPSHOT" -> GatewayModeLabel.SNAPSHOT
88+
mode == "RESUME" -> GatewayModeLabel.RESUME
89+
query.values.keys.any { it == "offset" || it == "cursor" || it == "off" } -> GatewayModeLabel.RESUME
90+
else -> GatewayModeLabel.SNAPSHOT
91+
}
92+
}
93+
94+
fun parseOffset(query: Map<String, String>): ParsedLong {
95+
val raw = query["offset"] ?: query["cursor"] ?: query["off"] ?: return ParsedLong()
96+
val value = raw.toLongOrNull()
97+
return ParsedLong(value = value, malformed = value == null)
98+
}
99+
100+
fun parseEpoch(query: Map<String, String>): ParsedLong {
101+
val raw = query["epoch"] ?: return ParsedLong()
102+
val value = raw.toLongOrNull()
103+
return ParsedLong(value = value, malformed = value == null)
104+
}
105+
106+
private fun decodePart(
107+
value: String,
108+
onMalformed: () -> Unit,
109+
): String? =
110+
runCatching { URLDecoder.decode(value, StandardCharsets.UTF_8) }
111+
.getOrElse {
112+
onMalformed()
113+
null
114+
}?.takeIf { it.isNotBlank() }
115+
}
116+
117+
internal fun agentIdOf(session: WebSocketSession): String? =
118+
session.uri
119+
?.path
120+
?.let { Regex("/ws/agents/([^/]+)/attach").find(it) }
121+
?.groupValues
122+
?.get(1)
123+
124+
internal data class QueryParseResult(
125+
val values: Map<String, String> = emptyMap(),
126+
val malformed: Boolean = false,
127+
)
128+
129+
internal data class ParsedLong(
130+
val value: Long? = null,
131+
val malformed: Boolean = false,
132+
)

0 commit comments

Comments
 (0)