feat: implement CreateGrant handler and grant-aware approval engine

- GlobalStreamHandler.handleCreateGrant validates scope (SESSION/STAGE) and sends ProtocolError to client on validation failure instead of silent log.warn + drop. - Derive event stageId directly from GrantScope type, removing redundant stageIdForEvent tuple. - SessionOrchestrator integrates ApprovalEngine to check active grants before requesting user approval for T2+ tool calls. - ContextEntry gains EntryRole (SYSTEM/USER/ASSISTANT/TOOL) field for proper chat message role mapping. - RouterContextBuilder.build is now suspend; uses Tokenizer for accurate token estimation with fallback to character-based estimate. - LlamaCppInferenceProvider maps EntryRole to ChatMessage role instead of heuristic layer-based inference.
2026-05-28 14:06:58 +04:00
parent 92bea6c2c4
commit e05532e7b2
11 changed files with 316 additions and 133 deletions
@@ -3,6 +3,8 @@ package com.correx.core.context.model
 import com.correx.core.events.types.ContextEntryId
 import kotlinx.serialization.Serializable

+enum class EntryRole { SYSTEM, USER, ASSISTANT, TOOL }
+
@Serializable
 data class ContextEntry(
    val id: ContextEntryId,
@@ -10,5 +12,6 @@ data class ContextEntry(
    val content: String,
    val sourceType: String,
    val sourceId: String,
-    val tokenEstimate: Int
+    val tokenEstimate: Int,
+    val role: EntryRole = EntryRole.USER,
 )
@@ -3,6 +3,7 @@ package com.correx.core.kernel.orchestration
 import com.correx.core.approvals.ApprovalOutcome
 import com.correx.core.approvals.ApprovalStatus
 import com.correx.core.approvals.model.ApprovalDecision
+import com.correx.core.inference.Tokenizer
 import com.correx.core.artifacts.ArtifactState
 import com.correx.core.artifactstore.ArtifactStore
 import com.correx.core.events.events.ApprovalDecisionResolvedEvent
@@ -38,7 +39,9 @@ class DefaultSessionOrchestrator(
    engines: OrchestratorEngines,
    private val retryCoordinator: RetryCoordinator,
    artifactStore: ArtifactStore,
+    tokenizer: Tokenizer? = null,
 ) : SessionOrchestrator(repositories, engines, artifactStore), ApprovalGateway {
+    override val tokenizer: Tokenizer? = tokenizer
    override val cancellations: ConcurrentHashMap<SessionId, AtomicBoolean> =
        ConcurrentHashMap<SessionId, AtomicBoolean>()

@@ -3,7 +3,12 @@ package com.correx.core.kernel.orchestration
 import com.correx.core.approvals.ApprovalOutcome
 import com.correx.core.approvals.ApprovalStatus
 import com.correx.core.approvals.Tier
+import com.correx.core.approvals.DefaultApprovalRepository
+import com.correx.core.approvals.domain.ApprovalEngine
+import com.correx.core.approvals.isAtMost
+import com.correx.core.approvals.model.ApprovalContext
 import com.correx.core.approvals.model.ApprovalDecision
+import com.correx.core.approvals.model.ApprovalScopeIdentity
 import com.correx.core.approvals.model.DomainApprovalRequest
 import com.correx.core.artifacts.ArtifactState
 import com.correx.core.artifacts.kind.JsonSchema
@@ -14,6 +19,7 @@ import com.correx.core.context.model.ContextEntry
 import com.correx.core.context.model.ContextLayer
 import com.correx.core.context.model.ContextPack
 import com.correx.core.context.model.TokenBudget
+import com.correx.core.context.model.EntryRole
 import com.correx.core.events.events.ApprovalDecisionResolvedEvent
 import com.correx.core.events.events.ApprovalRequestedEvent
 import com.correx.core.events.events.ArtifactCreatedEvent
@@ -55,8 +61,10 @@ import com.correx.core.inference.InferenceRequest
 import com.correx.core.inference.InferenceResponse
 import com.correx.core.inference.InferenceRouter
 import com.correx.core.inference.ResponseFormat
+import com.correx.core.inference.Tokenizer
 import com.correx.core.inference.ToolCallRequest
 import com.correx.core.inference.ToolDefinition
+import com.correx.core.sessions.ApprovalMode
 import com.correx.core.inference.ToolFunction
 import com.correx.core.kernel.execution.WorkflowResult
 import com.correx.core.risk.RiskAssessor
@@ -115,6 +123,9 @@ abstract class SessionOrchestrator(
    private val toolRegistry: ToolRegistry? = engines.toolRegistry
    private val inferenceRepository: InferenceRepository = repositories.inferenceRepository
    internal val orchestrationRepository: OrchestrationRepository = repositories.orchestrationRepository
+    protected open val tokenizer: Tokenizer? = null
+    private val approvalEngine: ApprovalEngine = engines.approvalEngine
+    private val approvalRepository: DefaultApprovalRepository = repositories.approvalRepository
    internal abstract val cancellations: ConcurrentHashMap<SessionId, AtomicBoolean>
    internal val pendingApprovals: ConcurrentHashMap<ApprovalRequestId, CompletableDeferred<ApprovalDecision>> =
        ConcurrentHashMap()
@@ -162,7 +173,8 @@ abstract class SessionOrchestrator(
                        content = text,
                        sourceType = "systemPrompt",
                        sourceId = stageId.value,
-                        tokenEstimate = text.length / 4,
+                        tokenEstimate = estimateTokens(text),
+                        role = EntryRole.SYSTEM,
                    ),
                )
            } ?: config.defaultSystemPromptPath
@@ -182,7 +194,8 @@ abstract class SessionOrchestrator(
                        content = text,
                        sourceType = "systemPrompt",
                        sourceId = stageId.value,
-                        tokenEstimate = text.length / 4,
+                        tokenEstimate = estimateTokens(text),
+                        role = EntryRole.SYSTEM,
                    ),
                )
            } ?: emptyList()
@@ -206,7 +219,8 @@ abstract class SessionOrchestrator(
                        content = text,
                        sourceType = "agentPrompt",
                        sourceId = stageId.value,
-                        tokenEstimate = text.length / 4,
+                        tokenEstimate = estimateTokens(text),
+                        role = EntryRole.USER,
                    ),
                )
            } ?: emptyList()
@@ -333,11 +347,15 @@ abstract class SessionOrchestrator(
                    request = request,
                ),
            )
-            val requiresApproval = when (tier) {
-                Tier.T0, Tier.T1 -> false
-                Tier.T2, Tier.T3, Tier.T4 -> true
-            }
-            if (requiresApproval) {
+            if (tier.isAtMost(Tier.T1)) {
+                // no approval needed
+            } else {
+                val approvalState = approvalRepository.getApprovalState(sessionId)
+                val activeGrants = approvalState.grants.values.toList()
+                val approvalCtx = ApprovalContext(
+                    identity = ApprovalScopeIdentity(sessionId, stageId, projectId = null),
+                    mode = ApprovalMode.PROMPT,
+                )
                val requestId = ApprovalRequestId(UUID.randomUUID().toString())
                val domainRequest = DomainApprovalRequest(
                    id = requestId,
@@ -348,59 +366,101 @@ abstract class SessionOrchestrator(
                    toolName = toolCall.function.name,
                    preview = toolCall.function.arguments.take(200),
                )
-                emit(sessionId, OrchestrationPausedEvent(sessionId, stageId, "APPROVAL_PENDING"))
-                emit(
-                    sessionId,
-                    ApprovalRequestedEvent(
-                        requestId = requestId,
-                        tier = tier,
-                        validationReportId = domainRequest.validationReportId,
-                        riskSummaryId = null,
-                        sessionId = sessionId,
-                        stageId = stageId,
-                        projectId = null,
-                        toolName = toolCall.function.name,
-                        preview = toolCall.function.arguments.take(200),
-                    ),
+                val engineDecision = approvalEngine.evaluate(
+                    domainRequest, approvalCtx, activeGrants, Clock.System.now(),
                )
-                val deferred = CompletableDeferred<ApprovalDecision>()
-                pendingApprovals[requestId] = deferred
-                val decision = try {
-                    deferred.await()
-                } finally {
-                    pendingApprovals.remove(requestId)
-                }
-                emitDecisionResolved(sessionId, domainRequest, decision)
-                if (!decision.isApproved) {
-                    val rejectReason = decision.reason ?: "approval denied"
-                    emit(sessionId, ToolExecutionRejectedEvent(
-                        invocationId = invocationId,
-                        sessionId = sessionId,
-                        toolName = toolCall.function.name,
-                        tier = tier,
-                        reason = rejectReason,
-                    ))
+                if (engineDecision.state == ApprovalStatus.COMPLETED) {
+                    emitDecisionResolved(sessionId, domainRequest, engineDecision)
+                    if (!engineDecision.isApproved) {
+                        val rejectReason = engineDecision.reason ?: "denied"
+                        emit(sessionId, ToolExecutionRejectedEvent(
+                            invocationId = invocationId,
+                            sessionId = sessionId,
+                            toolName = toolCall.function.name,
+                            tier = tier,
+                            reason = rejectReason,
+                        ))
+                        val sourceId = toolCall.id ?: invocationId.value
+                        return@flatMap listOf(
+                            ContextEntry(
+                                id = ContextEntryId(UUID.randomUUID().toString()),
+                                layer = ContextLayer.L2,
+                                sourceType = "assistantToolCall",
+                                sourceId = sourceId,
+                                content = Json.encodeToString(ToolCallRequest.serializer(), toolCall),
+                                tokenEstimate = estimateTokens(toolCall.function.arguments),
+                                role = EntryRole.ASSISTANT,
+                            ),
+                            ContextEntry(
+                                id = ContextEntryId(UUID.randomUUID().toString()),
+                                layer = ContextLayer.L2,
+                                sourceType = "toolResult",
+                                sourceId = sourceId,
+                                content = "ERROR: $rejectReason",
+                                tokenEstimate = estimateTokens(rejectReason),
+                                role = EntryRole.TOOL,
+                            ),
+                        )
+                    }
+                    // grant auto-approved — fall through to execute
+                } else {
+                    emit(sessionId, OrchestrationPausedEvent(sessionId, stageId, "APPROVAL_PENDING"))
+                    emit(
+                        sessionId,
+                        ApprovalRequestedEvent(
+                            requestId = requestId,
+                            tier = tier,
+                            validationReportId = domainRequest.validationReportId,
+                            riskSummaryId = null,
+                            sessionId = sessionId,
+                            stageId = stageId,
+                            projectId = null,
+                            toolName = toolCall.function.name,
+                            preview = toolCall.function.arguments.take(200),
+                        ),
+                    )
+                    val deferred = CompletableDeferred<ApprovalDecision>()
+                    pendingApprovals[requestId] = deferred
+                    val userDecision = try {
+                        deferred.await()
+                    } finally {
+                        pendingApprovals.remove(requestId)
+                    }
+                    emitDecisionResolved(sessionId, domainRequest, userDecision)
+                    if (!userDecision.isApproved) {
+                        val rejectReason = userDecision.reason ?: "approval denied"
+                        emit(sessionId, ToolExecutionRejectedEvent(
+                            invocationId = invocationId,
+                            sessionId = sessionId,
+                            toolName = toolCall.function.name,
+                            tier = tier,
+                            reason = rejectReason,
+                        ))
+                        emit(sessionId, OrchestrationResumedEvent(sessionId, stageId))
+                        val sourceId = toolCall.id ?: invocationId.value
+                        return@flatMap listOf(
+                            ContextEntry(
+                                id = ContextEntryId(UUID.randomUUID().toString()),
+                                layer = ContextLayer.L2,
+                                sourceType = "assistantToolCall",
+                                sourceId = sourceId,
+                                content = Json.encodeToString(ToolCallRequest.serializer(), toolCall),
+                                tokenEstimate = estimateTokens(toolCall.function.arguments),
+                                role = EntryRole.ASSISTANT,
+                            ),
+                            ContextEntry(
+                                id = ContextEntryId(UUID.randomUUID().toString()),
+                                layer = ContextLayer.L2,
+                                sourceType = "toolResult",
+                                sourceId = sourceId,
+                                content = "ERROR: $rejectReason",
+                                tokenEstimate = estimateTokens(rejectReason),
+                                role = EntryRole.TOOL,
+                            ),
+                        )
+                    }
                    emit(sessionId, OrchestrationResumedEvent(sessionId, stageId))
-                    val sourceId = toolCall.id ?: invocationId.value
-                    val assistantEntry = ContextEntry(
-                        id = ContextEntryId(UUID.randomUUID().toString()),
-                        layer = ContextLayer.L2,
-                        sourceType = "assistantToolCall",
-                        sourceId = sourceId,
-                        content = Json.encodeToString(ToolCallRequest.serializer(), toolCall),
-                        tokenEstimate = toolCall.function.arguments.length / 4,
-                    )
-                    val resultEntry = ContextEntry(
-                        id = ContextEntryId(UUID.randomUUID().toString()),
-                        layer = ContextLayer.L2,
-                        sourceType = "toolResult",
-                        sourceId = sourceId,
-                        content = "ERROR: $rejectReason",
-                        tokenEstimate = rejectReason.length / 4,
-                    )
-                    return@flatMap listOf(assistantEntry, resultEntry)
                }
-                emit(sessionId, OrchestrationResumedEvent(sessionId, stageId))
            }
            val result = executor.execute(request)

@@ -436,7 +496,8 @@ abstract class SessionOrchestrator(
                sourceType = "assistantToolCall",
                sourceId = sourceId,
                content = Json.encodeToString(ToolCallRequest.serializer(), toolCall),
-                tokenEstimate = toolCall.function.arguments.length / 4,
+                tokenEstimate = estimateTokens(toolCall.function.arguments),
+                role = EntryRole.ASSISTANT,
            )
            val resultContent = when (result) {
                is ToolResult.Success -> result.output
@@ -451,13 +512,14 @@ abstract class SessionOrchestrator(
                sourceType = "toolResult",
                sourceId = sourceId,
                content = resultContent,
-                tokenEstimate = resultContent.length / 4,
+                tokenEstimate = estimateTokens(resultContent),
+                role = EntryRole.TOOL,
            )
            listOf(assistantEntry, resultEntry)
        }
    }

-    private fun buildSchemaEntries(
+    private suspend fun buildSchemaEntries(
        responseFormat: ResponseFormat,
        stageId: StageId,
    ): List<ContextEntry> {
@@ -473,7 +535,8 @@ abstract class SessionOrchestrator(
                content = instruction,
                sourceType = "schemaInstruction",
                sourceId = stageId.value,
-                tokenEstimate = instruction.length / 4,
+                tokenEstimate = estimateTokens(instruction),
+                role = EntryRole.SYSTEM,
            ),
        )
    }
@@ -742,6 +805,20 @@ abstract class SessionOrchestrator(
        )
    }

+    // --- token estimation ---
+
+    protected open suspend fun estimateTokens(content: String): Int {
+        val t = tokenizer
+        if (t != null) {
+            return runCatching { t.countTokens(content) }.getOrElse { fallbackTokenEstimate(content) }
+        }
+        return fallbackTokenEstimate(content)
+    }
+
+    private fun fallbackTokenEstimate(content: String): Int {
+        return (content.length / 4).coerceAtLeast(1)
+    }
+
    // --- private functions ---

    private suspend fun handleApproval(
@@ -4,23 +4,27 @@ import com.correx.core.context.model.CompressionMetadata
 import com.correx.core.context.model.ContextEntry
 import com.correx.core.context.model.ContextLayer
 import com.correx.core.context.model.ContextPack
+import com.correx.core.context.model.EntryRole
 import com.correx.core.context.model.TokenBudget
 import com.correx.core.events.types.ContextEntryId
 import com.correx.core.events.types.ContextPackId
 import com.correx.core.events.types.SessionId
 import com.correx.core.events.types.StageId
+import com.correx.core.inference.Tokenizer
 import com.correx.core.router.model.RouterConfig
 import com.correx.core.router.model.RouterL2Entry
 import com.correx.core.router.model.RouterState
 import com.correx.core.router.model.RouterTurn
+import com.correx.core.router.model.TurnRole
 import java.util.*

 interface RouterContextBuilder {
-    fun build(state: RouterState, budget: TokenBudget): ContextPack
+    suspend fun build(state: RouterState, budget: TokenBudget): ContextPack
 }

 class DefaultRouterContextBuilder(
    private val config: RouterConfig,
+    private val tokenizer: Tokenizer? = null,
 ) : RouterContextBuilder {

    companion object {
@@ -28,40 +32,43 @@ class DefaultRouterContextBuilder(
            "You are a routing assistant. Provide guidance based on workflow state and conversation context."
    }

-    override fun build(state: RouterState, budget: TokenBudget): ContextPack {
+    override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
        var remainingBudget = budget.limit
        val allEntries = mutableListOf<ContextEntry>()
        var droppedCount = 0

-        // === L0: System prompt (never dropped) ===
        val systemPrompt = buildContextEntry(
            sourceType = "systemPrompt",
            sourceId = "router-system",
            content = SYSTEM_PROMPT,
+            role = EntryRole.SYSTEM,
        )
        remainingBudget -= systemPrompt.tokenEstimate
        if (remainingBudget < 0) remainingBudget = 0
        allEntries += systemPrompt

-        // === L0: Workflow status (never dropped) ===
        val workflowStatusEntry = buildContextEntry(
            sourceType = "workflowStatus",
            sourceId = state.currentStageId?.value ?: "none",
            content = buildWorkflowStatusContent(state),
+            role = EntryRole.SYSTEM,
        )
        remainingBudget -= workflowStatusEntry.tokenEstimate
        if (remainingBudget < 0) remainingBudget = 0
        allEntries += workflowStatusEntry

-        // === L1: Conversation history (last N turns, capped at keepLast) ===
        val recentTurns = state.conversationHistory.takeLast(config.conversationKeepLast)
        for (turn in recentTurns) {
-            val content = buildConversationContent(turn)
+            val role = when (turn.role) {
+                TurnRole.USER -> EntryRole.USER
+                TurnRole.ROUTER -> EntryRole.ASSISTANT
+            }
            val entry = buildContextEntry(
                sourceType = "conversation",
                sourceId = "${turn.role.name}-${turn.hashCode()}",
-                content = content,
+                content = turn.content,
                layer = ContextLayer.L1,
+                role = role,
            )
            if (remainingBudget >= entry.tokenEstimate) {
                remainingBudget -= entry.tokenEstimate
@@ -72,7 +79,6 @@ class DefaultRouterContextBuilder(
            }
        }

-        // === L2: Stage summaries from L2 memory (oldest-first eviction) ===
        for (l2Entry in state.l2Memory) {
            val content = buildL2Content(l2Entry)
            val entry = buildContextEntry(
@@ -80,6 +86,7 @@ class DefaultRouterContextBuilder(
                sourceId = l2Entry.stageId.value,
                content = content,
                layer = ContextLayer.L2,
+                role = EntryRole.SYSTEM,
            )
            if (remainingBudget >= entry.tokenEstimate) {
                remainingBudget -= entry.tokenEstimate
@@ -90,7 +97,6 @@ class DefaultRouterContextBuilder(
            }
        }

-        // Assemble layers and pack
        val layers = allEntries.groupBy { it.layer }
        val budgetUsed = allEntries.sumOf { it.tokenEstimate }

@@ -119,19 +125,16 @@ class DefaultRouterContextBuilder(
        return s.lowercase().let { it[0].uppercaseChar() + it.drop(1) }
    }

-    private fun buildConversationContent(turn: RouterTurn): String {
-        return "[${turn.role.name}] ${turn.content}"
-    }
-
    private fun buildL2Content(l2Entry: RouterL2Entry): String {
        return "Stage ${l2Entry.stageId.value} (${l2Entry.outcome}): ${l2Entry.summary}"
    }

-    private fun buildContextEntry(
+    private suspend fun buildContextEntry(
        sourceType: String,
        sourceId: String,
        content: String,
        layer: ContextLayer = ContextLayer.L0,
+        role: EntryRole = EntryRole.USER,
    ): ContextEntry {
        return ContextEntry(
            id = ContextEntryId(UUID.randomUUID().toString()),
@@ -140,10 +143,19 @@ class DefaultRouterContextBuilder(
            sourceType = sourceType,
            sourceId = sourceId,
            tokenEstimate = estimateTokens(content),
+            role = role,
        )
    }

-    private fun estimateTokens(content: String): Int {
-        return (content.length / 2).coerceAtLeast(1)
+    private suspend fun estimateTokens(content: String): Int {
+        val t = tokenizer
+        if (t != null) {
+            return runCatching { t.countTokens(content) }.getOrElse { fallbackEstimate(content) }
+        }
+        return fallbackEstimate(content)
    }
-}
+
+    private fun fallbackEstimate(content: String): Int {
+        return (content.length / 4).coerceAtLeast(1)
+    }
+}