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:
2026-05-28 14:06:58 +04:00
parent 92bea6c2c4
commit e05532e7b2
11 changed files with 316 additions and 133 deletions
@@ -56,7 +56,8 @@ fun main() {
logStoresInfo(eventStore, artifactStore)
val infraRegistry = InfrastructureModule.createProviderRegistry(listOf(buildLlamaProvider()))
val llamaProvider = buildLlamaProvider()
val infraRegistry = InfrastructureModule.createProviderRegistry(listOf(llamaProvider))
val modelId = System.getenv("CORREX_MODEL_ID") ?: "default"
val modelPath = System.getenv("CORREX_MODEL_PATH") ?: "(not set)"
val llamaUrl = System.getenv("CORREX_LLAMA_URL") ?: "http://127.0.0.1:10000"
@@ -109,6 +110,7 @@ fun main() {
engines = engines,
retryCoordinator = DefaultRetryCoordinator(eventStore),
artifactStore = artifactStore,
tokenizer = llamaProvider.tokenizer,
)
val defaultOrchestrationConfig = OrchestrationConfig(
sandboxRoot = sandboxRoot,
@@ -119,6 +121,7 @@ fun main() {
eventStore = eventStore,
inferenceRouter = inferenceRouter,
config = routerConfig,
tokenizer = llamaProvider.tokenizer,
)
val module = ServerModule(
orchestrator = orchestrator,
@@ -1,8 +1,12 @@
package com.correx.apps.server.protocol
import com.correx.core.approvals.GrantScope
import com.correx.core.approvals.Tier
import com.correx.core.events.types.ApprovalRequestId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.router.ChatMode
import kotlinx.datetime.Instant
import kotlinx.serialization.Serializable
@Serializable
@@ -11,6 +15,12 @@ enum class ApprovalDecision {
REJECT,
}
@Serializable
enum class GrantScopeDto {
SESSION,
STAGE,
}
@Serializable
sealed class ClientMessage {
@Serializable
@@ -29,6 +39,16 @@ sealed class ClientMessage {
val steeringNote: String?,
) : ClientMessage()
@Serializable
data class CreateGrant(
val sessionId: SessionId,
val scope: GrantScopeDto,
val stageId: StageId? = null,
val permittedTiers: List<Tier>,
val reason: String,
val expiresAt: Instant? = null,
) : ClientMessage()
@Serializable
data class Ping(val timestamp: Long) : ClientMessage()
@@ -4,16 +4,20 @@ import com.correx.apps.server.ServerModule
import com.correx.apps.server.bridge.DomainEventMapper
import com.correx.apps.server.bridge.SessionEventBridge
import com.correx.apps.server.protocol.ClientMessage
import com.correx.apps.server.protocol.GrantScopeDto
import com.correx.apps.server.protocol.ProtocolSerializer
import com.correx.apps.server.protocol.ProviderHealthDto
import com.correx.apps.server.protocol.ServerMessage
import com.correx.apps.server.protocol.WorkflowDto
import com.correx.apps.server.protocol.StageToolDecl
import com.correx.apps.server.protocol.ToolDecl
import com.correx.core.approvals.GrantScope
import com.correx.core.events.events.ApprovalGrantCreatedEvent
import com.correx.core.events.events.EventMetadata
import com.correx.core.events.events.NewEvent
import com.correx.core.events.events.StoredEvent
import com.correx.core.events.events.WorkflowFailedEvent
import com.correx.core.events.types.GrantId
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.EventId
import com.correx.core.events.types.SessionId
@@ -140,6 +144,7 @@ class GlobalStreamHandler(private val module: ServerModule) {
is ClientMessage.CancelSession -> module.orchestrator.cancel(msg.sessionId)
is ClientMessage.ResumeSession -> session.send(encodeError("ResumeSession not supported"))
is ClientMessage.ApprovalResponse -> handleApprovalResponse(msg, sendFrame)
is ClientMessage.CreateGrant -> handleCreateGrant(msg, sendFrame)
is ClientMessage.ChatInput -> Unit // handler deferred to subsequent task
}
}
@@ -163,6 +168,12 @@ class GlobalStreamHandler(private val module: ServerModule) {
module.approvalCoordinator.handleResponse(msg, scopeSessionId)?.let { sendFrame(it) }
}
private fun errorResponse(message: String) = ServerMessage.ProtocolError(
message = message,
sequence = null,
sessionSequence = null,
)
private fun encodeError(message: String): Frame.Text =
Frame.Text(
ProtocolSerializer.encodeServerMessage(
@@ -174,6 +185,49 @@ class GlobalStreamHandler(private val module: ServerModule) {
),
)
private suspend fun handleCreateGrant(
msg: ClientMessage.CreateGrant,
sendFrame: suspend (ServerMessage) -> Unit,
) {
val scope = when (msg.scope) {
GrantScopeDto.SESSION -> {
if (msg.stageId != null) {
sendFrame(errorResponse("CreateGrant: SESSION scope must not include stageId"))
return
}
GrantScope.SESSION
}
GrantScopeDto.STAGE -> {
val sid = msg.stageId ?: run {
sendFrame(errorResponse("CreateGrant: STAGE scope requires stageId"))
return
}
GrantScope.STAGE(sid)
}
}
val event = NewEvent(
metadata = EventMetadata(
eventId = EventId(UUID.randomUUID().toString()),
sessionId = msg.sessionId,
timestamp = Clock.System.now(),
schemaVersion = 1,
causationId = null,
correlationId = null,
),
payload = ApprovalGrantCreatedEvent(
grantId = GrantId(UUID.randomUUID().toString()),
scope = scope,
permittedTiers = msg.permittedTiers.toSet(),
reason = msg.reason,
expiresAt = msg.expiresAt,
sessionId = msg.sessionId,
stageId = (scope as? GrantScope.STAGE)?.stageId,
projectId = null,
),
)
module.eventStore.append(event)
}
private suspend fun handleStartSession(
session: DefaultWebSocketServerSession,
msg: ClientMessage.StartSession,