fix: headless job durability

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
+        }
+}

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)
     }
 

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<StartHeadlessJobCommand> {
     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,53 +82,56 @@ 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",
                 cause = ex,
             )
         }
 
-    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"
     }

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"
+    }
 }

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<RetrievalPort.Snippet>): List<RetrievalPort.Snippet> = snippets
 
-    private fun dedupedAndFiltered(query: String): List<RetrievalPort.Snippet> {
+    private fun dedupedAndFiltered(
+        query: String,
+        scope: String? = null,
+    ): List<RetrievalPort.Snippet> {
         val seenIds = mutableSetOf<String>()
         val seenTexts = mutableSetOf<String>()
         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 ->

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<Snippet>
 }
 

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<RetrievalPort.Snippet> =
         runCatching {
+            val args =
+                buildMap<String, Any?> {
+                    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;

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<RetrievalPort.Snippet> {
         if (!props.retrievalEnabled) return emptyList()
-        return transport.recall(query, limit)
+        return transport.recall(query, limit, scope)
     }
 }

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<RetrievalPort.Snippet> {
         if (!props.retrievalEnabled) return emptyList()
         return runCatching {

api/src/test/kotlin/com/jorisjonkers/personalstack/agents/application/command/SendUserInputCommandHandlerTest.kt

@@ -32,7 +32,7 @@ class SendUserInputCommandHandlerTest {
     private val gateway = mockk<AgentGatewayClient>(relaxed = true)
     private val contextBuilder =
         mockk<ContextBuilder> {
-            every { augment(any()) } answers { firstArg() }
+            every { augment(any(), any()) } answers { firstArg() }
         }
     private val binding = mockk<RunnerSessionBindingService>()
     private val handler = SendUserInputCommandHandler(turns, gateway, contextBuilder, binding)
@@ -61,7 +61,7 @@ class SendUserInputCommandHandlerTest {
             bound(ws, session)
         val savedTurn = slot<Turn>()
         every { turns.save(capture(savedTurn)) } answers { savedTurn.captured }
-        every { contextBuilder.augment("hello") } returns "<context>kb hit</context>\n\nhello"
+        every { contextBuilder.augment("hello", any()) } returns "<context>kb hit</context>\n\nhello"
 
         handler.handle(SendUserInputCommand(sessionId = session.id, text = "hello"))