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.
This commit is contained in:
+3
@@ -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>()
|
||||
|
||||
|
||||
+139
-62
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user