From 61e9e2045d4c8c737064a4134312c5b554d4d5cf Mon Sep 17 00:00:00 2001 From: JorisJonkers Agent Date: Sun, 12 Jul 2026 11:42:14 +0000 Subject: [PATCH 1/4] fix: route RunnerPodChatGenerator through AgentGatewayClient with typed agent kind RunnerPodChatGenerator previously bypassed AgentGatewayClient.startHeadlessJob() and called the gateway directly via raw HTTP with an untyped String kind. This meant the two-phase durability commit, idle-endpoint split, and KB hook controls added in agents-api #22 did not apply to chat-generation headless jobs. Changes: - AgentGatewayClient.HeadlessJobRequest: add enableKbHooks (default false) so callers can opt headless jobs into KB auto-recall/capture hooks; the gateway injects KB_AUTO_MCP_DISABLED=1 when the flag is absent (agent-runtime v0.18.1). - HttpAgentGatewayClient: forward enableKbHooks in HeadlessRequestBody. - ChatGenerationProperties: change runnerPodAgentKind from String to WorkspaceAgentKind (typed); add runnerPodEnableKbHooks (default false). - RunnerPodChatGenerator: inject AgentGatewayClient for job launch instead of a raw RestClient POST; resolve the full Workspace object (required by the client contract); keep RestClient only for the SSE stream endpoint that has no equivalent on the port. The class doc is updated to reflect the routing change. - Tests: new assertions in HttpAgentGatewayClientTest for enableKbHooks forwarding; RunnerPodChatGeneratorTest updated for the new constructor and extended with three new tests covering typed-kind routing, enableKbHooks propagation, and the missing-endpoint guard. Closes #14 --- .../agents/config/ChatGenerationProperties.kt | 8 +- .../agents/domain/port/AgentGatewayClient.kt | 6 + .../integration/HttpAgentGatewayClient.kt | 2 + .../integration/RunnerPodChatGenerator.kt | 72 +++++------ .../integration/HttpAgentGatewayClientTest.kt | 33 +++++ .../integration/RunnerPodChatGeneratorTest.kt | 113 +++++++++++++++++- 6 files changed, 192 insertions(+), 42 deletions(-) diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/config/ChatGenerationProperties.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/config/ChatGenerationProperties.kt index 2957825..65dfddc 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/config/ChatGenerationProperties.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/config/ChatGenerationProperties.kt @@ -1,5 +1,6 @@ package com.jorisjonkers.personalstack.agents.config +import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceAgentKind import org.springframework.boot.context.properties.ConfigurationProperties /** @@ -13,10 +14,15 @@ import org.springframework.boot.context.properties.ConfigurationProperties * * `runnerPodWorkspaceId` and `runnerPodAgentKind` are only consulted when * `backend=runner-pod`. + * + * `runnerPodEnableKbHooks` opts the chat-generation headless job into KB + * auto-recall/capture hooks. Defaults to false: the gateway sets + * KB_AUTO_MCP_DISABLED=1 so chat workers do not fire accidental KB writes. */ @ConfigurationProperties(prefix = "chat.generation") data class ChatGenerationProperties( val backend: String = "lightrag", val runnerPodWorkspaceId: String? = null, - val runnerPodAgentKind: String = "CLAUDE", + val runnerPodAgentKind: WorkspaceAgentKind = WorkspaceAgentKind.CLAUDE, + val runnerPodEnableKbHooks: Boolean = false, ) diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/AgentGatewayClient.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/AgentGatewayClient.kt index b690466..17a1626 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/AgentGatewayClient.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/AgentGatewayClient.kt @@ -124,6 +124,12 @@ interface AgentGatewayClient { val stableSessionId: WorkspaceAgentSessionId? = null, val epoch: Long? = null, val continuation: ContinuationMetadata? = null, + /** + * When false (the default), the gateway injects KB_AUTO_MCP_DISABLED=1 + * so the headless worker does not fire auto-KB recall/capture hooks. + * Set to true only for runs that explicitly need KB hook access. + */ + val enableKbHooks: Boolean = false, ) fun startHeadlessJob(request: HeadlessJobRequest): HeadlessJob diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClient.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClient.kt index 927b989..8b51e8d 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClient.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClient.kt @@ -268,6 +268,7 @@ class HttpAgentGatewayClient( val epoch: Long? = null, val continuation: ContinuationBody? = null, val timeoutSeconds: Long? = null, + val enableKbHooks: Boolean = false, ) private data class HeadlessJobDto( @@ -291,6 +292,7 @@ class HttpAgentGatewayClient( epoch = request.epoch, continuation = request.continuation?.toBody(), timeoutSeconds = request.timeoutSeconds, + enableKbHooks = request.enableKbHooks, ), ).retrieve() .body(HeadlessJobDto::class.java) diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt index 26cc496..455fe81 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt @@ -5,6 +5,7 @@ import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper import com.jorisjonkers.personalstack.agents.config.ChatGenerationProperties import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceId +import com.jorisjonkers.personalstack.agents.domain.port.AgentGatewayClient import com.jorisjonkers.personalstack.agents.domain.port.ChatGenerationPort import com.jorisjonkers.personalstack.agents.domain.port.WorkspaceRepository import org.slf4j.LoggerFactory @@ -16,16 +17,21 @@ import java.util.UUID /** * [ChatGenerationPort] adapter that routes a chat prompt through a headless - * agent job on the runner Pod. + * agent job on the runner Pod via the shared [AgentGatewayClient]. * * Flow: * 1. Resolve the target workspace via [ChatGenerationProperties.runnerPodWorkspaceId]. - * 2. POST `{gatewayEndpoint}/agents/headless` with `kind` + `prompt`; read - * back the job `id`. + * 2. Submit the job via [AgentGatewayClient.startHeadlessJob] with the configured + * [ChatGenerationProperties.runnerPodAgentKind]; receive back the job id. * 3. GET `{gatewayEndpoint}/agents/headless/{id}/stream` as text/event-stream; * parse Spring SSE format (event:/data: pairs) and delegate each line to * [parseStreamJsonLine] to extract assistant text chunks. * + * Using the shared gateway client (step 2) ensures the job is submitted through + * the same typed request path as [StartHeadlessJobCommandHandler], so + * [AgentGatewayClient.HeadlessJobRequest.enableKbHooks] and any future + * request-level fields are honoured uniformly. + * * Active only when `chat.generation.backend=runner-pod`. */ @Component @@ -35,6 +41,7 @@ import java.util.UUID havingValue = "runner-pod", ) class RunnerPodChatGenerator( + private val gateway: AgentGatewayClient, private val restClient: RestClient, private val workspaces: WorkspaceRepository, private val props: ChatGenerationProperties, @@ -45,68 +52,55 @@ class RunnerPodChatGenerator( prompt: String, onChunk: (String) -> Unit, ): String { - val gatewayEndpoint = resolveGatewayEndpoint() ?: return "" - val jobId = startHeadlessJob(gatewayEndpoint, prompt) ?: return "" - return streamJobOutput(gatewayEndpoint, jobId, onChunk) + val workspace = resolveWorkspace() ?: return "" + val jobId = startHeadlessJob(workspace, prompt) ?: return "" + return streamJobOutput(workspace.gatewayEndpoint!!, jobId, onChunk) } - private fun resolveGatewayEndpoint(): String? { + private fun resolveWorkspace(): com.jorisjonkers.personalstack.agents.domain.model.Workspace? { val rawId = props.runnerPodWorkspaceId if (rawId.isNullOrBlank()) { log.warn("chat.generation.runner-pod-workspace-id is not configured — no answer produced") return null } - var endpoint: String? = null - runCatching { + return runCatching { val workspaceId = WorkspaceId(UUID.fromString(rawId)) val workspace = workspaces.findById(workspaceId) if (workspace == null) { log.warn("RunnerPodChatGenerator: workspace {} not found — no answer produced", rawId) - return@runCatching + return@runCatching null } - val ep = workspace.gatewayEndpoint - if (ep.isNullOrBlank()) { + if (workspace.gatewayEndpoint.isNullOrBlank()) { log.warn( "RunnerPodChatGenerator: workspace {} has no gateway endpoint — no answer produced", rawId, ) - return@runCatching + return@runCatching null } - endpoint = ep + workspace }.onFailure { ex -> log.warn("RunnerPodChatGenerator: failed to resolve workspace {}: {}", rawId, ex.message) - } - return endpoint + }.getOrNull() } private fun startHeadlessJob( - gatewayEndpoint: String, + workspace: com.jorisjonkers.personalstack.agents.domain.model.Workspace, prompt: String, - ): String? { - var jobId: String? = null + ): String? = runCatching { - val response = - restClient - .post() - .uri("$gatewayEndpoint/agents/headless") - .body( - mapOf( - "kind" to props.runnerPodAgentKind, - "prompt" to prompt, - "partialMessages" to true, - ), - ).retrieve() - .body(Map::class.java) - jobId = response?.get("id")?.toString() - if (jobId.isNullOrBlank()) { - log.warn("RunnerPodChatGenerator: gateway returned no job id") - jobId = null - } + val job = + gateway.startHeadlessJob( + AgentGatewayClient.HeadlessJobRequest( + workspace = workspace, + kind = props.runnerPodAgentKind, + prompt = prompt, + enableKbHooks = props.runnerPodEnableKbHooks, + ), + ) + job.id }.onFailure { ex -> log.warn("RunnerPodChatGenerator: failed to start headless job: {}", ex.message) - } - return jobId - } + }.getOrNull() private fun streamJobOutput( gatewayEndpoint: String, diff --git a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClientTest.kt b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClientTest.kt index c2e5896..d4bb90b 100644 --- a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClientTest.kt +++ b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClientTest.kt @@ -245,6 +245,39 @@ class HttpAgentGatewayClientTest { server.verify() } + @Test + fun `startHeadlessJob forwards enableKbHooks to gateway`() { + val builder = RestClient.builder() + val server = MockRestServiceServer.bindTo(builder).build() + val client = HttpAgentGatewayClient(builder.build()) + val ws = workspace() + val prompt = "run with kb hooks" + + server + .expect(requestTo("http://runner:8090/agents/headless")) + .andExpect(method(HttpMethod.POST)) + .andExpect(jsonPath("$.enableKbHooks").value(true)) + .andRespond( + withSuccess( + """{"id":"job-kbhooks-1","status":"RUNNING","exitCode":null,"output":null}""", + MediaType.APPLICATION_JSON, + ), + ) + + val job = + client.startHeadlessJob( + AgentGatewayClient.HeadlessJobRequest( + workspace = ws, + kind = WorkspaceAgentKind.CLAUDE, + prompt = prompt, + enableKbHooks = true, + ), + ) + + assertThat(job.id).isEqualTo("job-kbhooks-1") + server.verify() + } + private fun workspace() = Workspace( id = WorkspaceId.parse("55555555-5555-4555-8555-555555555555"), diff --git a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt index 9f2457c..fca8b02 100644 --- a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt +++ b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt @@ -1,13 +1,21 @@ package com.jorisjonkers.personalstack.agents.infrastructure.integration import com.jorisjonkers.personalstack.agents.config.ChatGenerationProperties +import com.jorisjonkers.personalstack.agents.domain.model.Workspace +import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceAgentKind import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceId +import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceStatus +import com.jorisjonkers.personalstack.agents.domain.port.AgentGatewayClient import com.jorisjonkers.personalstack.agents.domain.port.WorkspaceRepository 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.springframework.web.client.RestClient +import java.time.Instant +import java.util.UUID class RunnerPodChatGeneratorTest { // region: parseStreamJsonLine unit tests @@ -93,9 +101,10 @@ class RunnerPodChatGeneratorTest { @Test fun `generate returns empty string when runnerPodWorkspaceId is null`() { val props = ChatGenerationProperties(backend = "runner-pod", runnerPodWorkspaceId = null) + val gateway = mockk() val restClient = mockk() val workspaceRepo = mockk() - val generator = RunnerPodChatGenerator(restClient, workspaceRepo, props) + val generator = RunnerPodChatGenerator(gateway, restClient, workspaceRepo, props) val chunks = mutableListOf() val result = generator.generate("some prompt") { chunks.add(it) } @@ -108,11 +117,12 @@ class RunnerPodChatGeneratorTest { fun `generate returns empty string when workspace is not found`() { val wsId = "11111111-1111-4111-8111-111111111111" val props = ChatGenerationProperties(backend = "runner-pod", runnerPodWorkspaceId = wsId) + val gateway = mockk() val restClient = mockk() val workspaceRepo = mockk() every { workspaceRepo.findById(WorkspaceId.parse(wsId)) } returns null - val generator = RunnerPodChatGenerator(restClient, workspaceRepo, props) + val generator = RunnerPodChatGenerator(gateway, restClient, workspaceRepo, props) val chunks = mutableListOf() val result = generator.generate("some prompt") { chunks.add(it) } @@ -121,4 +131,103 @@ class RunnerPodChatGeneratorTest { } // endregion + + // region: gateway routing and typed kind + + @Test + fun `startHeadlessJob routes through AgentGatewayClient with typed kind`() { + val wsId = "22222222-2222-4222-8222-222222222222" + val props = + ChatGenerationProperties( + backend = "runner-pod", + runnerPodWorkspaceId = wsId, + runnerPodAgentKind = WorkspaceAgentKind.CODEX, + runnerPodEnableKbHooks = false, + ) + val workspace = stubWorkspace(wsId, "http://gateway:8080") + val workspaceRepo = mockk() + every { workspaceRepo.findById(WorkspaceId.parse(wsId)) } returns workspace + + val capturedRequest = slot() + val gateway = mockk() + every { gateway.startHeadlessJob(capture(capturedRequest)) } throws RuntimeException("stream not needed") + + val restClient = mockk() + val generator = RunnerPodChatGenerator(gateway, restClient, workspaceRepo, props) + + generator.generate("my prompt") {} + + // gateway was called + verify(exactly = 1) { gateway.startHeadlessJob(any()) } + // typed kind is forwarded + assertThat(capturedRequest.captured.kind).isEqualTo(WorkspaceAgentKind.CODEX) + // prompt is forwarded + assertThat(capturedRequest.captured.prompt).isEqualTo("my prompt") + // kb hooks default off + assertThat(capturedRequest.captured.enableKbHooks).isFalse() + // workspace object is the resolved one + assertThat(capturedRequest.captured.workspace).isSameAs(workspace) + } + + @Test + fun `startHeadlessJob forwards enableKbHooks=true when configured`() { + val wsId = "33333333-3333-4333-8333-333333333333" + val props = + ChatGenerationProperties( + backend = "runner-pod", + runnerPodWorkspaceId = wsId, + runnerPodAgentKind = WorkspaceAgentKind.CLAUDE, + runnerPodEnableKbHooks = true, + ) + val workspace = stubWorkspace(wsId, "http://gateway:8080") + val workspaceRepo = mockk() + every { workspaceRepo.findById(WorkspaceId.parse(wsId)) } returns workspace + + val capturedRequest = slot() + val gateway = mockk() + every { gateway.startHeadlessJob(capture(capturedRequest)) } throws RuntimeException("stream not needed") + + val restClient = mockk() + val generator = RunnerPodChatGenerator(gateway, restClient, workspaceRepo, props) + + generator.generate("kb prompt") {} + + assertThat(capturedRequest.captured.enableKbHooks).isTrue() + } + + @Test + fun `generate returns empty when workspace has no gateway endpoint`() { + val wsId = "44444444-4444-4444-8444-444444444444" + val props = ChatGenerationProperties(backend = "runner-pod", runnerPodWorkspaceId = wsId) + val workspace = stubWorkspace(wsId, gatewayEndpoint = null) + val workspaceRepo = mockk() + every { workspaceRepo.findById(WorkspaceId.parse(wsId)) } returns workspace + + val gateway = mockk() + val restClient = mockk() + val generator = RunnerPodChatGenerator(gateway, restClient, workspaceRepo, props) + + val result = generator.generate("prompt") {} + + assertThat(result).isEmpty() + verify(exactly = 0) { gateway.startHeadlessJob(any()) } + } + + // endregion + + private fun stubWorkspace( + id: String, + gatewayEndpoint: String?, + ) = Workspace( + id = WorkspaceId.parse(id), + name = "test-workspace", + repoUrl = null, + branch = null, + podName = "pod-1", + pvcName = "pvc-1", + gatewayEndpoint = gatewayEndpoint, + status = WorkspaceStatus.READY, + createdAt = Instant.EPOCH, + updatedAt = Instant.EPOCH, + ) } From aec8476abc1d9b1dded6d4ec56f0b752994fccc4 Mon Sep 17 00:00:00 2001 From: JorisJonkers Agent Date: Sun, 12 Jul 2026 11:50:16 +0000 Subject: [PATCH 2/4] fix: remove unused UUID import from RunnerPodChatGeneratorTest --- .../infrastructure/integration/RunnerPodChatGeneratorTest.kt | 1 - 1 file changed, 1 deletion(-) diff --git a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt index fca8b02..ba27b72 100644 --- a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt +++ b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt @@ -15,7 +15,6 @@ import org.assertj.core.api.Assertions.assertThat import org.junit.jupiter.api.Test import org.springframework.web.client.RestClient import java.time.Instant -import java.util.UUID class RunnerPodChatGeneratorTest { // region: parseStreamJsonLine unit tests From 917b20a7288161995fabd14fd6551847fdc8fb05 Mon Sep 17 00:00:00 2001 From: JorisJonkers Agent Date: Sun, 12 Jul 2026 12:01:19 +0000 Subject: [PATCH 3/4] fix: avoid !! on nullable gatewayEndpoint in RunnerPodChatGenerator Return Pair from resolveWorkspace() so the validated non-null endpoint is threaded through without an unsafe !! call. --- .../integration/RunnerPodChatGenerator.kt | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt index 455fe81..2edc633 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt @@ -52,12 +52,17 @@ class RunnerPodChatGenerator( prompt: String, onChunk: (String) -> Unit, ): String { - val workspace = resolveWorkspace() ?: return "" + val (workspace, gatewayEndpoint) = resolveWorkspace() ?: return "" val jobId = startHeadlessJob(workspace, prompt) ?: return "" - return streamJobOutput(workspace.gatewayEndpoint!!, jobId, onChunk) + return streamJobOutput(gatewayEndpoint, jobId, onChunk) } - private fun resolveWorkspace(): com.jorisjonkers.personalstack.agents.domain.model.Workspace? { + /** + * Resolves the target workspace and validates that a gateway endpoint is present. + * Returns a pair of ([Workspace], gatewayEndpoint) or null when the workspace is + * not configured, not found, or has no gateway endpoint yet. + */ + private fun resolveWorkspace(): Pair? { val rawId = props.runnerPodWorkspaceId if (rawId.isNullOrBlank()) { log.warn("chat.generation.runner-pod-workspace-id is not configured — no answer produced") @@ -70,14 +75,15 @@ class RunnerPodChatGenerator( log.warn("RunnerPodChatGenerator: workspace {} not found — no answer produced", rawId) return@runCatching null } - if (workspace.gatewayEndpoint.isNullOrBlank()) { + val endpoint = workspace.gatewayEndpoint + if (endpoint.isNullOrBlank()) { log.warn( "RunnerPodChatGenerator: workspace {} has no gateway endpoint — no answer produced", rawId, ) return@runCatching null } - workspace + workspace to endpoint }.onFailure { ex -> log.warn("RunnerPodChatGenerator: failed to resolve workspace {}: {}", rawId, ex.message) }.getOrNull() From 821e147ca4442f38bc8a29d1ed65d906b9638b2b Mon Sep 17 00:00:00 2001 From: JorisJonkers Agent Date: Sun, 12 Jul 2026 12:10:26 +0000 Subject: [PATCH 4/4] fix: restore partialMessages flag and add happy-path streaming test Codex review flagged two issues: 1. The migration from raw HTTP to AgentGatewayClient dropped partialMessages=true. The SseChatAccumulator relies on token-level deltas; without partial-messages mode the stream would deliver no incremental chunks. Added partialMessages to HeadlessJobRequest/HeadlessRequestBody and set it true in RunnerPodChatGenerator. 2. Routing tests threw on startHeadlessJob, never exercising the SSE streaming path. Added a MockRestServiceServer-based happy-path test that asserts the correct stream URL is called and the result event is parsed and returned. --- .../agents/domain/port/AgentGatewayClient.kt | 6 +++ .../integration/HttpAgentGatewayClient.kt | 2 + .../integration/RunnerPodChatGenerator.kt | 4 ++ .../integration/HttpAgentGatewayClientTest.kt | 4 +- .../integration/RunnerPodChatGeneratorTest.kt | 53 +++++++++++++++++++ 5 files changed, 68 insertions(+), 1 deletion(-) diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/AgentGatewayClient.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/AgentGatewayClient.kt index 17a1626..e42b186 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/AgentGatewayClient.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/AgentGatewayClient.kt @@ -130,6 +130,12 @@ interface AgentGatewayClient { * Set to true only for runs that explicitly need KB hook access. */ val enableKbHooks: Boolean = false, + /** + * Opts into token-level streaming via Claude `--include-partial-messages`. + * Should be set to true when the caller is consuming the SSE stream directly + * (e.g. [RunnerPodChatGenerator]), so token deltas arrive incrementally. + */ + val partialMessages: Boolean = false, ) fun startHeadlessJob(request: HeadlessJobRequest): HeadlessJob diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClient.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClient.kt index 8b51e8d..fb6567a 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClient.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClient.kt @@ -269,6 +269,7 @@ class HttpAgentGatewayClient( val continuation: ContinuationBody? = null, val timeoutSeconds: Long? = null, val enableKbHooks: Boolean = false, + val partialMessages: Boolean = false, ) private data class HeadlessJobDto( @@ -293,6 +294,7 @@ class HttpAgentGatewayClient( continuation = request.continuation?.toBody(), timeoutSeconds = request.timeoutSeconds, enableKbHooks = request.enableKbHooks, + partialMessages = request.partialMessages, ), ).retrieve() .body(HeadlessJobDto::class.java) diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt index 2edc633..61efd32 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGenerator.kt @@ -101,6 +101,10 @@ class RunnerPodChatGenerator( kind = props.runnerPodAgentKind, prompt = prompt, enableKbHooks = props.runnerPodEnableKbHooks, + // Token-level streaming is always enabled for chat-generation runs so + // the SseChatAccumulator receives incremental deltas rather than waiting + // for the full assistant turn before delivering any output. + partialMessages = true, ), ) job.id diff --git a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClientTest.kt b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClientTest.kt index d4bb90b..f26f4d5 100644 --- a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClientTest.kt +++ b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/HttpAgentGatewayClientTest.kt @@ -246,7 +246,7 @@ class HttpAgentGatewayClientTest { } @Test - fun `startHeadlessJob forwards enableKbHooks to gateway`() { + fun `startHeadlessJob forwards enableKbHooks and partialMessages to gateway`() { val builder = RestClient.builder() val server = MockRestServiceServer.bindTo(builder).build() val client = HttpAgentGatewayClient(builder.build()) @@ -257,6 +257,7 @@ class HttpAgentGatewayClientTest { .expect(requestTo("http://runner:8090/agents/headless")) .andExpect(method(HttpMethod.POST)) .andExpect(jsonPath("$.enableKbHooks").value(true)) + .andExpect(jsonPath("$.partialMessages").value(true)) .andRespond( withSuccess( """{"id":"job-kbhooks-1","status":"RUNNING","exitCode":null,"output":null}""", @@ -271,6 +272,7 @@ class HttpAgentGatewayClientTest { kind = WorkspaceAgentKind.CLAUDE, prompt = prompt, enableKbHooks = true, + partialMessages = true, ), ) diff --git a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt index ba27b72..8c34fdc 100644 --- a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt +++ b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/RunnerPodChatGeneratorTest.kt @@ -13,6 +13,12 @@ import io.mockk.slot import io.mockk.verify import org.assertj.core.api.Assertions.assertThat import org.junit.jupiter.api.Test +import org.springframework.http.HttpMethod +import org.springframework.http.MediaType +import org.springframework.test.web.client.MockRestServiceServer +import org.springframework.test.web.client.match.MockRestRequestMatchers.method +import org.springframework.test.web.client.match.MockRestRequestMatchers.requestTo +import org.springframework.test.web.client.response.MockRestResponseCreators.withSuccess import org.springframework.web.client.RestClient import java.time.Instant @@ -164,6 +170,8 @@ class RunnerPodChatGeneratorTest { assertThat(capturedRequest.captured.prompt).isEqualTo("my prompt") // kb hooks default off assertThat(capturedRequest.captured.enableKbHooks).isFalse() + // partial-messages always on for SSE streaming + assertThat(capturedRequest.captured.partialMessages).isTrue() // workspace object is the resolved one assertThat(capturedRequest.captured.workspace).isSameAs(workspace) } @@ -214,6 +222,51 @@ class RunnerPodChatGeneratorTest { // endregion + @Test + fun `generate streams SSE output on successful job launch`() { + val wsId = "55555555-5555-4555-8555-555555555555" + val props = + ChatGenerationProperties( + backend = "runner-pod", + runnerPodWorkspaceId = wsId, + runnerPodAgentKind = WorkspaceAgentKind.CLAUDE, + ) + val workspace = stubWorkspace(wsId, "http://gateway:8080") + val workspaceRepo = mockk() + every { workspaceRepo.findById(WorkspaceId.parse(wsId)) } returns workspace + + val job = + AgentGatewayClient.HeadlessJob( + id = "job-stream-1", + status = AgentGatewayClient.HeadlessStatus.RUNNING, + exitCode = null, + output = null, + ) + val gateway = mockk() + every { gateway.startHeadlessJob(any()) } returns job + + // Use a mock RestClient that returns a single SSE result line + val sseBody = + """ + event:line + data:{"type":"result","subtype":"success","result":"streaming answer"} + + """.trimIndent() + val restClientBuilder = RestClient.builder() + val server = MockRestServiceServer.bindTo(restClientBuilder).build() + server + .expect(requestTo("http://gateway:8080/agents/headless/job-stream-1/stream")) + .andExpect(method(HttpMethod.GET)) + .andRespond(withSuccess(sseBody, MediaType.TEXT_EVENT_STREAM)) + + val generator = RunnerPodChatGenerator(gateway, restClientBuilder.build(), workspaceRepo, props) + val chunks = mutableListOf() + val result = generator.generate("stream prompt") { chunks.add(it) } + + assertThat(result).isEqualTo("streaming answer") + server.verify() + } + private fun stubWorkspace( id: String, gatewayEndpoint: String?,