diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/HeadlessJobSessionPersistence.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/HeadlessJobSessionPersistence.kt new file mode 100644 index 0000000..fa29581 --- /dev/null +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/HeadlessJobSessionPersistence.kt @@ -0,0 +1,71 @@ +package com.jorisjonkers.personalstack.agents.application.command + +import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceAgentSession +import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceAgentSessionStatus +import com.jorisjonkers.personalstack.agents.domain.port.AgentGatewayClient +import com.jorisjonkers.personalstack.agents.domain.port.WorkspaceAgentSessionRepository +import org.springframework.stereotype.Component +import org.springframework.transaction.annotation.Propagation +import org.springframework.transaction.annotation.Transactional +import java.time.Instant + +/** + * Isolated persistence helper for the two-phase headless job commit. + * + * Extracted into a separate Spring component so that each method is invoked + * through the proxy — Spring AOP intercepts external bean calls only, so + * calling @Transactional methods on `this` inside [StartHeadlessJobCommandHandler] + * would bypass the proxy and lose the REQUIRES_NEW guarantee. + */ +@Component +class HeadlessJobSessionPersistence( + private val sessions: WorkspaceAgentSessionRepository, +) { + /** + * Phase 1: saves the STARTING placeholder before the gateway call. + * Committed in its own transaction so the row is durable on disk even if + * the process crashes during the gateway round-trip. + */ + @Transactional(propagation = Propagation.REQUIRES_NEW) + fun saveStartingSession(session: WorkspaceAgentSession): WorkspaceAgentSession = sessions.save(session) + + /** + * Phase 2: updates the STARTING session with the gateway job id and its + * reported status. Maps the gateway’s HeadlessStatus so that jobs that + * complete or fail synchronously (COMPLETED, FAILED, CANCELLED) are not + * recorded as RUNNING in the DB. + * Committed in its own transaction independent of the outer call. + */ + @Transactional(propagation = Propagation.REQUIRES_NEW) + fun persistSession( + pendingSession: WorkspaceAgentSession, + job: AgentGatewayClient.HeadlessJob, + ) { + sessions.save( + pendingSession.copy( + gatewayAgentId = job.id, + status = gatewayStatusToSessionStatus(job.status), + updatedAt = Instant.now(), + ), + ) + } + + /** + * Marks the STARTING session as FAILED when the gateway call fails. + * Prevents STARTING rows from lingering as irreconcilable stubs. + */ + @Transactional(propagation = Propagation.REQUIRES_NEW) + fun markSessionFailed(session: WorkspaceAgentSession) { + sessions.save(session.markFailed(now = Instant.now())) + } + + private fun gatewayStatusToSessionStatus( + gatewayStatus: AgentGatewayClient.HeadlessStatus, + ): WorkspaceAgentSessionStatus = + when (gatewayStatus) { + AgentGatewayClient.HeadlessStatus.RUNNING -> WorkspaceAgentSessionStatus.RUNNING + AgentGatewayClient.HeadlessStatus.COMPLETED -> WorkspaceAgentSessionStatus.STOPPED + AgentGatewayClient.HeadlessStatus.CANCELLED -> WorkspaceAgentSessionStatus.STOPPED + AgentGatewayClient.HeadlessStatus.FAILED -> WorkspaceAgentSessionStatus.FAILED + } +} diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/SendUserInputCommandHandler.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/SendUserInputCommandHandler.kt index e457ad1..063bda3 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/SendUserInputCommandHandler.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/SendUserInputCommandHandler.kt @@ -2,6 +2,7 @@ package com.jorisjonkers.personalstack.agents.application.command import com.jorisjonkers.personalstack.agents.application.exception.AgentRunnerUnavailableException import com.jorisjonkers.personalstack.agents.application.rag.ContextBuilder +import com.jorisjonkers.personalstack.agents.application.rag.ScopeInference import com.jorisjonkers.personalstack.agents.application.sessionbinding.EnsureRunnerSessionBoundInput import com.jorisjonkers.personalstack.agents.application.sessionbinding.RunnerSessionBindingResult import com.jorisjonkers.personalstack.agents.application.sessionbinding.RunnerSessionBindingService @@ -38,7 +39,8 @@ class SendUserInputCommandHandler( ), ) - val augmented = contextBuilder.augment(command.text) + val scope = ScopeInference.scopeFor(current.workspace) + val augmented = contextBuilder.augment(command.text, scope) gateway.sendInput(current.workspace, current.gatewayAgent.id, augmented, command.enter) } diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/StartHeadlessJobCommandHandler.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/StartHeadlessJobCommandHandler.kt index 8f7fa8d..cde20b9 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/StartHeadlessJobCommandHandler.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/command/StartHeadlessJobCommandHandler.kt @@ -7,15 +7,15 @@ import com.jorisjonkers.personalstack.agents.application.observability.ModeLabel import com.jorisjonkers.personalstack.agents.application.observability.OperationLabel import com.jorisjonkers.personalstack.agents.application.observability.OperationTelemetry import com.jorisjonkers.personalstack.agents.application.observability.OutcomeLabel +import com.jorisjonkers.personalstack.agents.application.rag.ContextBuilder +import com.jorisjonkers.personalstack.agents.application.rag.ScopeInference import com.jorisjonkers.personalstack.agents.application.workspacerunner.WorkspaceRunnerLifecycleService import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceAgentSession import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceAgentSessionStatus import com.jorisjonkers.personalstack.agents.domain.port.AgentGatewayClient -import com.jorisjonkers.personalstack.agents.domain.port.WorkspaceAgentSessionRepository import com.jorisjonkers.personalstack.common.command.CommandHandler import org.slf4j.LoggerFactory import org.springframework.stereotype.Component -import org.springframework.transaction.annotation.Transactional import java.time.Duration import java.time.Instant @@ -24,25 +24,37 @@ import java.time.Instant * one-shot headless job to the gateway, and persists a * [WorkspaceAgentSession] to track it. * - * Idle sweep protection (skipping workspaces with running headless jobs) - * depends on N3's `run_mode` column — activate after N3 merges. + * Two-phase commit: the session is first saved as STARTING with no + * gatewayAgentId (phase 1, committed independently via + * [HeadlessJobSessionPersistence]). After the gateway call succeeds the + * session is updated to RUNNING with the returned job id (phase 2). If + * the process crashes between the two phases the session remains STARTING + * and can be reconciled on the next startup sweep. If the gateway call + * itself fails the session is updated to FAILED so the STARTING row does + * not linger. + * + * The three transactional operations are delegated to + * [HeadlessJobSessionPersistence] to ensure Spring AOP can intercept each + * call through the proxy — self-invocation on `this` bypasses the proxy + * and loses the REQUIRES_NEW commit guarantee. */ @Component class StartHeadlessJobCommandHandler( - private val sessions: WorkspaceAgentSessionRepository, private val gateway: AgentGatewayClient, private val runnerLifecycle: WorkspaceRunnerLifecycleService, + private val persistence: HeadlessJobSessionPersistence, + private val contextBuilder: ContextBuilder, private val telemetry: AgentsApiTelemetry = AgentsApiTelemetry.NOOP, ) : CommandHandler { private val log = LoggerFactory.getLogger(StartHeadlessJobCommandHandler::class.java) - @Transactional override fun handle(command: StartHeadlessJobCommand) { val startedAt = Instant.now() runCatching { val runner = bootRunner(command) - val job = launchJob(runner, command) - persistSession(command, runner, job) + val pendingSession = persistence.saveStartingSession(buildStartingSession(command, runner)) + val job = launchJob(runner, command, pendingSession) + persistence.persistSession(pendingSession, job) log.info("headless job {} launched for workspace {}", job.id, runner.workspace.id.value) }.onSuccess { recordCommandOperation(startedAt, OutcomeLabel.SUCCESS, FailureReasonLabel.NONE) @@ -70,22 +82,49 @@ class StartHeadlessJobCommandHandler( } } + private fun buildStartingSession( + command: StartHeadlessJobCommand, + runner: WorkspaceRunnerLifecycleService.BootOutcome.Ready, + ): WorkspaceAgentSession { + val now = Instant.now() + return WorkspaceAgentSession( + id = command.sessionId, + workspaceId = runner.workspace.id, + kind = command.kind, + gatewayAgentId = null, + status = WorkspaceAgentSessionStatus.STARTING, + createdAt = now, + updatedAt = now, + runMode = HEADLESS_RUN_MODE, + currentSetupId = command.setupId ?: runner.workspace.currentRunnerSetupId, + currentSetupVersion = command.setupVersion ?: runner.workspace.currentRunnerSetupVersion, + epoch = 1, + generation = 1, + ) + } + private fun launchJob( runner: WorkspaceRunnerLifecycleService.BootOutcome.Ready, command: StartHeadlessJobCommand, + pendingSession: WorkspaceAgentSession, ): AgentGatewayClient.HeadlessJob = runCatching { + val scope = ScopeInference.scopeFor(runner.workspace) + val augmentedPrompt = contextBuilder.augment(command.prompt, scope) gateway.startHeadlessJob( AgentGatewayClient.HeadlessJobRequest( workspace = runner.workspace, kind = command.kind, - prompt = command.prompt, + prompt = augmentedPrompt, timeoutSeconds = command.timeoutSeconds, stableSessionId = command.sessionId, epoch = 1, ), ) }.getOrElse { ex -> + // Gateway call failed — mark the already-saved STARTING session as FAILED + // so it does not linger as an irreconcilable STARTING row. + persistence.markSessionFailed(pendingSession) throw AgentRunnerUnavailableException( workspaceId = runner.workspace.id, runnerStatus = "HeadlessLaunchFailed", @@ -93,30 +132,6 @@ class StartHeadlessJobCommandHandler( ) } - private fun persistSession( - command: StartHeadlessJobCommand, - runner: WorkspaceRunnerLifecycleService.BootOutcome.Ready, - job: AgentGatewayClient.HeadlessJob, - ) { - val now = Instant.now() - sessions.save( - WorkspaceAgentSession( - id = command.sessionId, - workspaceId = runner.workspace.id, - kind = command.kind, - gatewayAgentId = job.id, - status = WorkspaceAgentSessionStatus.RUNNING, - createdAt = now, - updatedAt = now, - runMode = HEADLESS_RUN_MODE, - currentSetupId = command.setupId ?: runner.workspace.currentRunnerSetupId, - currentSetupVersion = command.setupVersion ?: runner.workspace.currentRunnerSetupVersion, - epoch = 1, - generation = 1, - ), - ) - } - companion object { const val HEADLESS_RUN_MODE = "HEADLESS" } diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/idle/IdleScaleDownScheduler.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/idle/IdleScaleDownScheduler.kt index 0fde9fe..f1f2a06 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/idle/IdleScaleDownScheduler.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/idle/IdleScaleDownScheduler.kt @@ -141,8 +141,18 @@ class IdleScaleDownScheduler( * A stale runner may only be recycled once every running agent has been * idle (no output) for the grace window. A session with no live agent is * nothing to interrupt. If any agent is still producing output — or its - * idle time can't be confirmed from the gateway — we wait for the next + * idle state can't be confirmed from the gateway — we wait for the next * sweep rather than risk cutting off an in-progress task. + * + * Interactive sessions: use [AgentGatewayClient.agentIdle] which hits the + * interactive /agents/{id} endpoint. + * + * Headless sessions: use [AgentGatewayClient.pollHeadlessJob] and treat + * any terminal status (COMPLETED/FAILED/CANCELLED) as safe-to-recycle. + * The headless job id is stored in [WorkspaceAgentSession.gatewayAgentId]; + * calling the interactive /agents/{id} endpoint with a headless job id + * would return a 404 and force-cancel the idle check, keeping the runner + * alive indefinitely. */ private fun staleRunnerSafeToRecycle( workspace: Workspace, @@ -153,11 +163,29 @@ class IdleScaleDownScheduler( val grace = runtime.staleRecycleQuiet return running.all { session -> val agentId = session.gatewayAgentId ?: return@all false - val idle = gateway.agentIdle(workspace, agentId) - idle != null && idle >= grace + if (session.runMode == HEADLESS_RUN_MODE) { + headlessJobTerminal(workspace, agentId) + } else { + val idle = gateway.agentIdle(workspace, agentId) + idle != null && idle >= grace + } } } + /** + * Returns true when the headless job has reached a terminal state + * (COMPLETED, FAILED, or CANCELLED). Returns false if the job is still + * RUNNING or if the gateway cannot be reached. + */ + private fun headlessJobTerminal( + workspace: Workspace, + headlessJobId: String, + ): Boolean = + runCatching { + val job = gateway.pollHeadlessJob(workspace, headlessJobId) + job.status != AgentGatewayClient.HeadlessStatus.RUNNING + }.getOrElse { false } + private fun effectiveLastSeen(workspace: Workspace): Instant = tracker.lastSeen(workspace.id) ?: workspace.updatedAt private fun scaleDown(workspace: Workspace) { @@ -218,4 +246,8 @@ class IdleScaleDownScheduler( ), ) } + + private companion object { + const val HEADLESS_RUN_MODE = "HEADLESS" + } } diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/rag/ContextBuilder.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/rag/ContextBuilder.kt index 0524af7..deb7d4c 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/rag/ContextBuilder.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/application/rag/ContextBuilder.kt @@ -38,9 +38,12 @@ class ContextBuilder( private val injectedHits = registry.counter("rag.hits.injected") private val injectedChars = registry.counter("rag.chars.injected") - fun augment(userPrompt: String): String { + fun augment( + userPrompt: String, + scope: String? = null, + ): String { if (!props.retrievalEnabled || sources.isEmpty()) return userPrompt - val chunks = buildChunks(dedupedAndFiltered(userPrompt)) + val chunks = buildChunks(dedupedAndFiltered(userPrompt, scope)) // Empty when all snippets failed the score floor, the merged list was // empty, or the character budget was too tight for even the first chunk. if (chunks.isEmpty()) return userPrompt @@ -83,12 +86,15 @@ class ContextBuilder( */ private fun coexistenceRank(snippets: List): List = snippets - private fun dedupedAndFiltered(query: String): List { + private fun dedupedAndFiltered( + query: String, + scope: String? = null, + ): List { val seenIds = mutableSetOf() val seenTexts = mutableSetOf() return coexistenceRank( sources - .flatMap { it.retrieve(query, props.maxSnippets) } + .flatMap { it.retrieve(query, props.maxSnippets, scope) } .filter { it.score >= props.minScore } .sortedByDescending { it.score }, ).filter { s -> diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/RetrievalPort.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/RetrievalPort.kt index 024f24e..baf5d48 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/RetrievalPort.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/domain/port/RetrievalPort.kt @@ -17,6 +17,7 @@ interface RetrievalPort { fun retrieve( query: String, limit: Int, + scope: String? = null, ): List } diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/KnowledgeMcpTransport.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/KnowledgeMcpTransport.kt index 5a087ce..cf23f82 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/KnowledgeMcpTransport.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/KnowledgeMcpTransport.kt @@ -92,12 +92,20 @@ class KnowledgeMcpTransport( fun recall( query: String, limit: Int, + scope: String? = null, ): List = runCatching { + val args = + buildMap { + put("query", query) + put("limit", limit) + put("mode", props.recallMode) + if (!scope.isNullOrBlank()) put("scope", scope) + } val resp = callTool( RECALL_TOOL, - mapOf("query" to query, "limit" to limit, "mode" to props.recallMode), + args, ) val result = resp?.get("result") ?: return@runCatching emptyList() // Prefer structuredContent.hits (MCP 2025-06) for typed parsing; diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/KnowledgeRecallClient.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/KnowledgeRecallClient.kt index 2d46f46..6777301 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/KnowledgeRecallClient.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/KnowledgeRecallClient.kt @@ -29,8 +29,9 @@ class KnowledgeRecallClient( override fun retrieve( query: String, limit: Int, + scope: String?, ): List { if (!props.retrievalEnabled) return emptyList() - return transport.recall(query, limit) + return transport.recall(query, limit, scope) } } diff --git a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/LightRagClient.kt b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/LightRagClient.kt index 667d70d..9673384 100644 --- a/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/LightRagClient.kt +++ b/api/src/main/kotlin/com/jorisjonkers/personalstack/agents/infrastructure/integration/LightRagClient.kt @@ -61,6 +61,7 @@ class LightRagClient( override fun retrieve( query: String, limit: Int, + scope: String?, ): List { if (!props.retrievalEnabled) return emptyList() return runCatching { diff --git a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/command/SendUserInputCommandHandlerTest.kt b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/command/SendUserInputCommandHandlerTest.kt index 5fdbc4a..31dc569 100644 --- a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/command/SendUserInputCommandHandlerTest.kt +++ b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/command/SendUserInputCommandHandlerTest.kt @@ -32,7 +32,7 @@ class SendUserInputCommandHandlerTest { private val gateway = mockk(relaxed = true) private val contextBuilder = mockk { - every { augment(any()) } answers { firstArg() } + every { augment(any(), any()) } answers { firstArg() } } private val binding = mockk() private val handler = SendUserInputCommandHandler(turns, gateway, contextBuilder, binding) @@ -61,7 +61,7 @@ class SendUserInputCommandHandlerTest { bound(ws, session) val savedTurn = slot() every { turns.save(capture(savedTurn)) } answers { savedTurn.captured } - every { contextBuilder.augment("hello") } returns "kb hit\n\nhello" + every { contextBuilder.augment("hello", any()) } returns "kb hit\n\nhello" handler.handle(SendUserInputCommand(sessionId = session.id, text = "hello")) diff --git a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/command/StartHeadlessJobCommandHandlerTest.kt b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/command/StartHeadlessJobCommandHandlerTest.kt index 03204cc..418fcaa 100644 --- a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/command/StartHeadlessJobCommandHandlerTest.kt +++ b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/command/StartHeadlessJobCommandHandlerTest.kt @@ -7,6 +7,7 @@ import com.jorisjonkers.personalstack.agents.application.observability.ModeLabel import com.jorisjonkers.personalstack.agents.application.observability.OperationLabel import com.jorisjonkers.personalstack.agents.application.observability.OperationTelemetry import com.jorisjonkers.personalstack.agents.application.observability.OutcomeLabel +import com.jorisjonkers.personalstack.agents.application.rag.ContextBuilder import com.jorisjonkers.personalstack.agents.application.workspacerunner.RunnerUnavailableReason import com.jorisjonkers.personalstack.agents.application.workspacerunner.WorkspaceRunnerLifecycleService import com.jorisjonkers.personalstack.agents.domain.model.AgentSetupId @@ -19,7 +20,6 @@ import com.jorisjonkers.personalstack.agents.domain.model.WorkspaceAgentSessionS 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.WorkspaceAgentSessionRepository import io.mockk.every import io.mockk.mockk import io.mockk.slot @@ -30,11 +30,19 @@ import org.junit.jupiter.api.assertThrows import java.time.Instant class StartHeadlessJobCommandHandlerTest { - private val sessions = mockk() private val gateway = mockk() private val runnerLifecycle = mockk() private val telemetry = RecordingTelemetry() - private val handler = StartHeadlessJobCommandHandler(sessions, gateway, runnerLifecycle, telemetry) + private val persistence = mockk() + private val contextBuilder = mockk() + private val handler = + StartHeadlessJobCommandHandler( + gateway, + runnerLifecycle, + persistence, + contextBuilder, + telemetry, + ) @Test fun `handle launches headless job and persists session when runner is ready`() { @@ -65,8 +73,14 @@ class StartHeadlessJobCommandHandlerTest { exitCode = null, output = null, ) - val saved = slot() - every { sessions.save(capture(saved)) } answers { firstArg() } + // ContextBuilder passes through the prompt unchanged in this test + every { contextBuilder.augment(prompt, any()) } returns prompt + // Phase 1: saveStartingSession returns a STARTING session + val savedStarting = slot() + every { persistence.saveStartingSession(capture(savedStarting)) } answers { firstArg() } + // Phase 2: persistSession records the RUNNING update + val savedJob = slot() + every { persistence.persistSession(any(), capture(savedJob)) } returns Unit handler.handle( StartHeadlessJobCommand( @@ -79,11 +93,16 @@ class StartHeadlessJobCommandHandlerTest { ), ) - assertThat(saved.captured.gatewayAgentId).isEqualTo("hls-abc123") - assertThat(saved.captured.status).isEqualTo(WorkspaceAgentSessionStatus.RUNNING) - assertThat(saved.captured.runMode).isEqualTo(StartHeadlessJobCommandHandler.HEADLESS_RUN_MODE) - assertThat(saved.captured.currentSetupId).isEqualTo(AgentSetupId("gpu")) - assertThat(saved.captured.currentSetupVersion).isEqualTo(AgentSetupVersion(2)) + // Phase 1: STARTING placeholder with no gateway id + assertThat(savedStarting.captured.status).isEqualTo(WorkspaceAgentSessionStatus.STARTING) + assertThat(savedStarting.captured.gatewayAgentId).isNull() + assertThat(savedStarting.captured.runMode).isEqualTo(StartHeadlessJobCommandHandler.HEADLESS_RUN_MODE) + assertThat(savedStarting.captured.currentSetupId).isEqualTo(AgentSetupId("gpu")) + assertThat(savedStarting.captured.currentSetupVersion).isEqualTo(AgentSetupVersion(2)) + + // Phase 2: gateway job passed through with the returned id + assertThat(savedJob.captured.id).isEqualTo("hls-abc123") + telemetry.operations.single().let { event -> assertThat(event.operation).isEqualTo(OperationLabel.START_SESSION) assertThat(event.mode).isEqualTo(ModeLabel.HEADLESS) @@ -143,6 +162,7 @@ class StartHeadlessJobCommandHandlerTest { workspace = ws, provisioning = WorkspaceRunnerLifecycleService.BootProvisioningOutcome.AlreadyReady, ) + every { contextBuilder.augment("secret prompt", any()) } returns "secret prompt" every { gateway.startHeadlessJob( AgentGatewayClient.HeadlessJobRequest( @@ -155,6 +175,10 @@ class StartHeadlessJobCommandHandlerTest { ), ) } throws RuntimeException(exceptionMessage) + // Phase 1 save (STARTING) succeeds; phase 2 never reached. + // markSessionFailed also called — stub both to return. + every { persistence.saveStartingSession(any()) } answers { firstArg() } + every { persistence.markSessionFailed(any()) } returns Unit assertThrows { handler.handle( @@ -168,6 +192,9 @@ class StartHeadlessJobCommandHandlerTest { ) } + // Phase 1 save + markSessionFailed; no phase 2 + verify(exactly = 1) { persistence.saveStartingSession(any()) } + verify(exactly = 1) { persistence.markSessionFailed(any()) } telemetry.operations.single().let { event -> assertThat(event.outcome).isEqualTo(OutcomeLabel.FAILURE) assertThat(event.reason).isEqualTo(FailureReasonLabel.UPSTREAM_UNAVAILABLE) @@ -190,6 +217,7 @@ class StartHeadlessJobCommandHandlerTest { workspace = ws, provisioning = WorkspaceRunnerLifecycleService.BootProvisioningOutcome.AlreadyReady, ) + every { contextBuilder.augment("persist me", any()) } returns "persist me" every { gateway.startHeadlessJob( AgentGatewayClient.HeadlessJobRequest( @@ -207,7 +235,9 @@ class StartHeadlessJobCommandHandlerTest { exitCode = null, output = null, ) - every { sessions.save(any()) } throws RuntimeException(exceptionMessage) + // Phase 1 (STARTING) succeeds; phase 2 (RUNNING) throws to simulate DB failure after gateway. + every { persistence.saveStartingSession(any()) } answers { firstArg() } + every { persistence.persistSession(any(), any()) } throws RuntimeException(exceptionMessage) val ex = assertThrows { @@ -260,6 +290,57 @@ class StartHeadlessJobCommandHandlerTest { } } + @Test + fun `handle persists session as STOPPED when gateway returns COMPLETED immediately`() { + val ws = workspace(WorkspaceId.parse("cccccccc-cccc-4ccc-8ccc-cccccccccccc")) + val sessionId = WorkspaceAgentSessionId.parse("dddddddd-dddd-4ddd-8ddd-dddddddddddd") + val prompt = "fast headless task" + every { + runnerLifecycle.boot(ws.id, WorkspaceAgentKind.CLAUDE, null, null) + } returns + WorkspaceRunnerLifecycleService.BootOutcome.Ready( + workspace = ws, + provisioning = WorkspaceRunnerLifecycleService.BootProvisioningOutcome.AlreadyReady, + ) + every { contextBuilder.augment(prompt, any()) } returns prompt + every { + gateway.startHeadlessJob( + AgentGatewayClient.HeadlessJobRequest( + workspace = ws, + kind = WorkspaceAgentKind.CLAUDE, + prompt = prompt, + stableSessionId = sessionId, + epoch = 1, + ), + ) + } returns + AgentGatewayClient.HeadlessJob( + id = "hls-fast", + status = AgentGatewayClient.HeadlessStatus.COMPLETED, + exitCode = 0, + output = "done", + ) + every { persistence.saveStartingSession(any()) } answers { firstArg() } + val savedJob = slot() + every { persistence.persistSession(any(), capture(savedJob)) } returns Unit + + handler.handle( + StartHeadlessJobCommand( + sessionId = sessionId, + workspaceId = ws.id, + kind = WorkspaceAgentKind.CLAUDE, + prompt = prompt, + ), + ) + + // COMPLETED gateway status must be persisted — persistSession receives the COMPLETED job + // and its gatewayStatusToSessionStatus maps it to STOPPED (not RUNNING). + assertThat(savedJob.captured.status).isEqualTo(AgentGatewayClient.HeadlessStatus.COMPLETED) + telemetry.operations.single().let { event -> + assertThat(event.outcome).isEqualTo(OutcomeLabel.SUCCESS) + } + } + private fun workspace(id: WorkspaceId = WorkspaceId.parse("bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb")) = Workspace( id = id, diff --git a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/idle/IdleScaleDownSchedulerTest.kt b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/idle/IdleScaleDownSchedulerTest.kt index 71833cc..ecebc53 100644 --- a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/idle/IdleScaleDownSchedulerTest.kt +++ b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/idle/IdleScaleDownSchedulerTest.kt @@ -414,3 +414,185 @@ class IdleScaleDownSchedulerTest { updatedAt = updatedAt, ) } + +// Tests for headless job idle branching (fix/headless-durability) +class IdleScaleDownSchedulerHeadlessTest { + private class FixedClock( + var now: Instant, + ) : Clock() { + override fun getZone(): java.time.ZoneId = java.time.ZoneOffset.UTC + + override fun withZone(zone: java.time.ZoneId): Clock = this + + override fun instant(): Instant = now + } + + private val now = Instant.parse("2026-05-19T12:00:00Z") + private val clock = FixedClock(now) + private val workspaces = mockk() + private val agentSessions = mockk() + private val orchestrator = mockk(relaxed = true) + private val gateway = mockk(relaxed = true) + private val tracker = WorkspaceActivityTracker(clock) + private val connected = ConnectedClientTracker() + private val sessionStatus = mockk(relaxed = true) + private val scheduler = + IdleScaleDownScheduler( + stores = IdleScaleDownStores(workspaces = workspaces, agentSessions = agentSessions), + dependencies = + IdleScaleDownDependencies( + orchestrator = orchestrator, + gateway = gateway, + tracker = tracker, + connected = connected, + sessionStatus = sessionStatus, + ), + runtime = + IdleScaleDownRuntime( + idleAfterSeconds = 1_800, + agentIdleAfterSeconds = 14_400, + staleRecycleQuietSeconds = 300, + ), + clock = clock, + ) + + @Test + fun `sweep recycles stale runner once headless job has terminated`() { + val ws = + Workspace( + id = WorkspaceId.random(), + name = "demo", + repoUrl = null, + branch = null, + podName = "runner-hls", + pvcName = "ws-hls", + gatewayEndpoint = "http://x:8090", + status = WorkspaceStatus.READY, + createdAt = now.minusSeconds(7_200), + updatedAt = now.minusSeconds(60), + ) + val session = + WorkspaceAgentSession( + id = WorkspaceAgentSessionId.random(), + workspaceId = ws.id, + kind = WorkspaceAgentKind.CLAUDE, + gatewayAgentId = "hls-job-xyz", + status = WorkspaceAgentSessionStatus.RUNNING, + runMode = "HEADLESS", + createdAt = now.minusSeconds(3_600), + updatedAt = now.minusSeconds(3_600), + epoch = 1, + generation = 1, + ) + every { workspaces.findAllByStatusNot(WorkspaceStatus.DESTROYED) } returns listOf(ws) + every { agentSessions.findAllByWorkspaceId(ws.id) } returns listOf(session) + every { orchestrator.isRunnerImageStale(ws) } returns true + every { + gateway.pollHeadlessJob(ws, "hls-job-xyz") + } returns + AgentGatewayClient.HeadlessJob( + id = "hls-job-xyz", + status = AgentGatewayClient.HeadlessStatus.COMPLETED, + exitCode = 0, + output = null, + ) + every { workspaces.save(any()) } answers { firstArg() } + every { + agentSessions.clearGatewayBindingIfGeneration( + id = session.id, + expectedGeneration = session.generation, + now = now, + ) + } returns true + + scheduler.sweep() + + verify { orchestrator.scaleDown(ws) } + verify(exactly = 0) { gateway.agentIdle(any(), any()) } + } + + @Test + fun `sweep does not recycle stale runner while headless job is still running`() { + val ws = + Workspace( + id = WorkspaceId.random(), + name = "demo", + repoUrl = null, + branch = null, + podName = "runner-hls-busy", + pvcName = "ws-hls-busy", + gatewayEndpoint = "http://x:8090", + status = WorkspaceStatus.READY, + createdAt = now.minusSeconds(7_200), + updatedAt = now.minusSeconds(60), + ) + val session = + WorkspaceAgentSession( + id = WorkspaceAgentSessionId.random(), + workspaceId = ws.id, + kind = WorkspaceAgentKind.CLAUDE, + gatewayAgentId = "hls-job-busy", + status = WorkspaceAgentSessionStatus.RUNNING, + runMode = "HEADLESS", + createdAt = now.minusSeconds(3_600), + updatedAt = now.minusSeconds(3_600), + epoch = 1, + generation = 1, + ) + every { workspaces.findAllByStatusNot(WorkspaceStatus.DESTROYED) } returns listOf(ws) + every { agentSessions.findAllByWorkspaceId(ws.id) } returns listOf(session) + every { orchestrator.isRunnerImageStale(ws) } returns true + every { + gateway.pollHeadlessJob(ws, "hls-job-busy") + } returns + AgentGatewayClient.HeadlessJob( + id = "hls-job-busy", + status = AgentGatewayClient.HeadlessStatus.RUNNING, + exitCode = null, + output = null, + ) + + scheduler.sweep() + + verify(exactly = 0) { orchestrator.scaleDown(any()) } + verify(exactly = 0) { gateway.agentIdle(any(), any()) } + } + + @Test + fun `sweep does not recycle stale runner when headless gateway poll fails`() { + val ws = + Workspace( + id = WorkspaceId.random(), + name = "demo", + repoUrl = null, + branch = null, + podName = "runner-hls-err", + pvcName = "ws-hls-err", + gatewayEndpoint = "http://x:8090", + status = WorkspaceStatus.READY, + createdAt = now.minusSeconds(7_200), + updatedAt = now.minusSeconds(60), + ) + val session = + WorkspaceAgentSession( + id = WorkspaceAgentSessionId.random(), + workspaceId = ws.id, + kind = WorkspaceAgentKind.CLAUDE, + gatewayAgentId = "hls-job-err", + status = WorkspaceAgentSessionStatus.RUNNING, + runMode = "HEADLESS", + createdAt = now.minusSeconds(3_600), + updatedAt = now.minusSeconds(3_600), + epoch = 1, + generation = 1, + ) + every { workspaces.findAllByStatusNot(WorkspaceStatus.DESTROYED) } returns listOf(ws) + every { agentSessions.findAllByWorkspaceId(ws.id) } returns listOf(session) + every { orchestrator.isRunnerImageStale(ws) } returns true + every { gateway.pollHeadlessJob(ws, "hls-job-err") } throws RuntimeException("gateway unreachable") + + scheduler.sweep() + + verify(exactly = 0) { orchestrator.scaleDown(any()) } + } +} diff --git a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/rag/ContextBuilderTest.kt b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/rag/ContextBuilderTest.kt index c9404b5..0218eb2 100644 --- a/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/rag/ContextBuilderTest.kt +++ b/api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/rag/ContextBuilderTest.kt @@ -41,6 +41,7 @@ class ContextBuilderTest { override fun retrieve( query: String, limit: Int, + scope: String?, ): List = snippets }