kami/correx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: implement CreateGrant handler and grant-aware approval engine

Browse Source 
- 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:
kami
2026-05-28 14:06:58 +04:00
parent 92bea6c2c4
commit e05532e7b2
11 changed files with 316 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/server/src/main/kotlin/com/correx/apps/server/Main.kt
+4 -1
View File
@@ -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,
apps/server/src/main/kotlin/com/correx/apps/server/protocol/ClientMessage.kt
+20
View File
@@ -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()
apps/server/src/main/kotlin/com/correx/apps/server/ws/GlobalStreamHandler.kt
+54
View File
@@ -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,
core/context/src/main/kotlin/com/correx/core/context/model/ContextEntry.kt
+4 -1
View File
@@ -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,
)
core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/DefaultSessionOrchestrator.kt
+3
View File
@@ -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>()
core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt
+139 -62
View File
@@ -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(
core/router/src/main/kotlin/com/correx/core/router/RouterContextBuilder.kt
+29 -17
View File
@@ -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)
}
}
infrastructure/inference/llama_cpp/src/main/kotlin/com/correx/infrastructure/inference/llama/cpp/LlamaCppInferenceProvider.kt
+13 -12
View File
@@ -2,6 +2,7 @@ package com.correx.infrastructure.inference.llama.cpp
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.events.types.ProviderId
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.FinishReason
@@ -149,17 +150,17 @@ class LlamaCppInferenceProvider(
.flatMap { it.value }
.joinToString("\n\n") { it.content }
.takeIf { it.isNotBlank() }
val conversationMessages = sorted.filter { it.key != ContextLayer.L0 }.flatMap { (layer, entries) ->
when (layer) {
ContextLayer.L1 -> entries.map { ChatMessage("user", it.content) }
ContextLayer.L2 -> entries.map {
ChatMessage(
role = if (it.sourceType == "assistant") "assistant" else "user",
content = it.content,
)
}
else -> emptyList()
val conversationMessages = sorted.filter { it.key != ContextLayer.L0 }.flatMap { (_, entries) ->
entries.map { entry ->
ChatMessage(
role = when (entry.role) {
EntryRole.SYSTEM -> "system"
EntryRole.ASSISTANT -> "assistant"
EntryRole.TOOL -> "tool"
EntryRole.USER -> "user"
},
content = entry.content,
)
}
}
val messages = buildList {
@@ -168,4 +169,4 @@ class LlamaCppInferenceProvider(
}
return messages.ifEmpty { listOf(ChatMessage("user", "")) }
}
}
}
infrastructure/src/main/kotlin/com/correx/infrastructure/InfrastructureModule.kt
+3 -1
View File
@@ -10,6 +10,7 @@ import com.correx.core.events.EventDispatcher
import com.correx.core.events.stores.EventStore
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.Tokenizer
import com.correx.core.router.DefaultRouterContextBuilder
import com.correx.core.router.DefaultRouterFacade
import com.correx.core.router.DefaultRouterReducer
@@ -120,12 +121,13 @@ object InfrastructureModule {
eventStore: EventStore,
inferenceRouter: InferenceRouter,
config: RouterConfig = RouterConfig(),
tokenizer: Tokenizer? = null,
): RouterFacade {
val reducer = DefaultRouterReducer()
val projector = RouterProjector(reducer)
val replayer = DefaultEventReplayer(eventStore, projector)
val repository = DefaultRouterRepository(replayer)
val contextBuilder = DefaultRouterContextBuilder(config)
val contextBuilder = DefaultRouterContextBuilder(config, tokenizer)
return DefaultRouterFacade(repository, contextBuilder, inferenceRouter, eventStore, config)
}
}
testing/deterministic/src/test/kotlin/RouterContextBuilderTest.kt
+35 -27
View File
@@ -1,5 +1,8 @@
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 kotlinx.coroutines.runBlocking
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.router.DefaultRouterContextBuilder
@@ -23,6 +26,10 @@ class RouterContextBuilderTest {
private val config = RouterConfig(conversationKeepLast = 3, tokenBudget = TokenBudget(limit = 200))
private val builder = DefaultRouterContextBuilder(config)
private fun buildPack(state: RouterState, budget: TokenBudget): ContextPack = runBlocking {
buildPack(state, budget)
}
private val sessionId = SessionId("test-session")
private val stageId = StageId("stage-1")
private val clock = Clock.System
@@ -41,7 +48,7 @@ class RouterContextBuilderTest {
RouterTurn(TurnRole.USER, "hi", clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
assertEquals(10000, pack.budgetLimit)
assertTrue(pack.budgetUsed <= pack.budgetLimit)
}
@@ -63,7 +70,7 @@ class RouterContextBuilderTest {
RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 100))
val pack = buildPack(state, TokenBudget(limit = 100))
// L0 entries (system prompt + workflow status) always fit; L1/L2 should be dropped
assertTrue(pack.compressionMetadata.entriesDropped > 0)
}
@@ -83,7 +90,7 @@ class RouterContextBuilderTest {
)
// L0 consumes ~60 tokens; budget 93 leaves ~33 for L2.
// Each L2 entry is ~12 tokens; 2 fit (s1, s2), s3 is dropped
val pack = builder.build(state, TokenBudget(limit = 93))
val pack = buildPack(state, TokenBudget(limit = 93))
val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList()
// Oldest-first eviction: s1 and s2 fit, s3 is the one dropped
val remainingStageIds = l2Entries.map { it.sourceId }.toSet()
@@ -127,7 +134,7 @@ class RouterContextBuilderTest {
workflowStatus = WorkflowStatus.IDLE,
currentStageId = null,
)
val pack = builder.build(state, TokenBudget(limit = 0))
val pack = buildPack(state, TokenBudget(limit = 0))
val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList()
assertTrue(l0Entries.any { it.sourceType == "systemPrompt" })
}
@@ -139,7 +146,7 @@ class RouterContextBuilderTest {
workflowStatus = WorkflowStatus.FAILED,
currentStageId = StageId("failed-stage"),
)
val pack = builder.build(state, TokenBudget(limit = 0))
val pack = buildPack(state, TokenBudget(limit = 0))
val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList()
assertTrue(l0Entries.any { it.sourceType == "workflowStatus" })
val workflowEntry = l0Entries.find { it.sourceType == "workflowStatus" }
@@ -156,7 +163,7 @@ class RouterContextBuilderTest {
conversationHistory = emptyList(),
l2Memory = emptyList(),
)
val pack = builder.build(state, TokenBudget(limit = 0))
val pack = buildPack(state, TokenBudget(limit = 0))
val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList()
assertEquals(2, l0Entries.size) // systemPrompt + workflowStatus
}
@@ -174,7 +181,7 @@ class RouterContextBuilderTest {
RouterL2Entry(StageId("s1"), "x".repeat(500), StageOutcomeKind.SUCCESS, clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 50))
val pack = buildPack(state, TokenBudget(limit = 50))
val l0Entries = pack.layers[ContextLayer.L0] ?: emptyList()
val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList()
val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList()
@@ -225,7 +232,7 @@ class RouterContextBuilderTest {
RouterL2Entry(StageId("s2"), "summary", StageOutcomeKind.FAILURE, clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList()
assertEquals(2, l2Entries.size)
assertEquals(0, pack.compressionMetadata.entriesDropped)
@@ -241,7 +248,7 @@ class RouterContextBuilderTest {
sessionId = sessionId,
workflowStatus = WorkflowStatus.IDLE,
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l0Entries = pack.layers[ContextLayer.L0]
assertNotNull(l0Entries)
assertTrue(l0Entries!!.any { it.sourceType == "systemPrompt" && it.layer == ContextLayer.L0 })
@@ -254,7 +261,7 @@ class RouterContextBuilderTest {
workflowStatus = WorkflowStatus.RUNNING,
currentStageId = stageId,
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l0Entries = pack.layers[ContextLayer.L0]
assertTrue(l0Entries!!.any { it.sourceType == "workflowStatus" && it.layer == ContextLayer.L0 })
}
@@ -270,7 +277,7 @@ class RouterContextBuilderTest {
RouterTurn(TurnRole.ROUTER, "hi", clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l1Entries = pack.layers[ContextLayer.L1]
assertNotNull(l1Entries)
assertEquals(2, l1Entries!!.size)
@@ -288,7 +295,7 @@ class RouterContextBuilderTest {
RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l2Entries = pack.layers[ContextLayer.L2]
assertNotNull(l2Entries)
assertEquals(1, l2Entries!!.size)
@@ -309,7 +316,7 @@ class RouterContextBuilderTest {
RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val presentLayers = pack.layers.keys
assertTrue(presentLayers.containsAll(listOf(ContextLayer.L0, ContextLayer.L1, ContextLayer.L2)))
assertFalse(presentLayers.contains(ContextLayer.L3))
@@ -336,7 +343,7 @@ class RouterContextBuilderTest {
RouterTurn(TurnRole.USER, "five", Instant.parse("2026-01-05T00:00:00Z")),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList()
// conversationKeepLast=2, so only the last 2 turns should appear
assertEquals(2, l1Entries.size)
@@ -357,7 +364,7 @@ class RouterContextBuilderTest {
RouterTurn(TurnRole.ROUTER, "fourth", Instant.parse("2026-01-04T00:00:00Z")),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList()
assertEquals(3, l1Entries.size)
// The last 3: second, third, fourth
@@ -373,7 +380,7 @@ class RouterContextBuilderTest {
@Test
fun `build with empty state produces L0 only`() {
val state = RouterState()
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l0Entries = pack.layers[ContextLayer.L0]
assertNotNull(l0Entries)
assertEquals(2, l0Entries!!.size)
@@ -384,7 +391,7 @@ class RouterContextBuilderTest {
@Test
fun `pack contains correct metadata on empty state`() {
val state = RouterState()
val pack = builder.build(state, TokenBudget(limit = 5000))
val pack = buildPack(state, TokenBudget(limit = 5000))
assertEquals(5000, pack.budgetLimit)
assertEquals(0, pack.compressionMetadata.entriesDropped)
assertEquals(listOf("L0Immutable", "Conversation"), pack.compressionMetadata.appliedStrategies)
@@ -398,7 +405,7 @@ class RouterContextBuilderTest {
workflowStatus = WorkflowStatus.RUNNING,
currentStageId = stageId,
)
val pack = builder.build(state, TokenBudget(limit = 1000))
val pack = buildPack(state, TokenBudget(limit = 1000))
assertEquals("test-session-router-pack", pack.id.value)
assertEquals(sessionId, pack.sessionId)
assertEquals(stageId, pack.stageId)
@@ -411,7 +418,7 @@ class RouterContextBuilderTest {
@Test
fun `system prompt has expected content`() {
val state = RouterState()
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val systemEntry = pack.layers[ContextLayer.L0]?.find { it.sourceType == "systemPrompt" }
assertNotNull(systemEntry)
assertEquals(
@@ -427,7 +434,7 @@ class RouterContextBuilderTest {
workflowStatus = WorkflowStatus.COMPLETED,
currentStageId = null,
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val workflowEntry = pack.layers[ContextLayer.L0]?.find { it.sourceType == "workflowStatus" }
assertNotNull(workflowEntry)
assertTrue(workflowEntry!!.content.contains("Status: Completed"))
@@ -444,10 +451,11 @@ class RouterContextBuilderTest {
RouterTurn(TurnRole.USER, "user message", clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l1Entries = pack.layers[ContextLayer.L1] ?: emptyList()
assertEquals(1, l1Entries.size)
assertEquals("[USER] user message", l1Entries[0].content)
assertEquals("user message", l1Entries[0].content)
assertEquals(EntryRole.USER, l1Entries[0].role)
}
@Test
@@ -460,7 +468,7 @@ class RouterContextBuilderTest {
RouterL2Entry(StageId("stage-x"), "completed with 3 items", StageOutcomeKind.SUCCESS, clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val l2Entries = pack.layers[ContextLayer.L2] ?: emptyList()
assertEquals(1, l2Entries.size)
val entry = l2Entries[0]
@@ -487,7 +495,7 @@ class RouterContextBuilderTest {
RouterL2Entry(StageId("s1"), "summary", StageOutcomeKind.SUCCESS, clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
val allEntries = pack.layers.values.flatten()
val computedSum = allEntries.sumOf { it.tokenEstimate }
assertEquals(computedSum, pack.budgetUsed)
@@ -509,7 +517,7 @@ class RouterContextBuilderTest {
RouterL2Entry(StageId("s2"), "x".repeat(500), StageOutcomeKind.SUCCESS, clock.now()),
),
)
val pack = builder.build(state, TokenBudget(limit = 65))
val pack = buildPack(state, TokenBudget(limit = 65))
// L0 consumes ~60 tokens, leaving 5 — both L2 entries (each ~259 tokens) dropped
// conversationKeepLast=0 means conversation entry is not included
assertEquals(2, pack.compressionMetadata.entriesDropped)
@@ -548,7 +556,7 @@ class RouterContextBuilderTest {
),
),
)
val pack = builder.build(state, TokenBudget(limit = 10000))
val pack = buildPack(state, TokenBudget(limit = 10000))
// L0: system prompt + workflow status
val l0 = pack.layers[ContextLayer.L0]!!
@@ -577,7 +585,7 @@ class RouterContextBuilderTest {
val state = RouterState(
workflowStatus = WorkflowStatus.IDLE,
)
val pack = builder.build(state, TokenBudget(limit = 1000))
val pack = buildPack(state, TokenBudget(limit = 1000))
assertTrue(pack.id.value.contains("unknown"))
}
}
testing/deterministic/src/test/kotlin/RouterFacadeTest.kt
+12 -12
View File
@@ -109,7 +109,7 @@ class RouterFacadeTest {
)
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter("response"),
eventStore = mockStore,
@@ -129,7 +129,7 @@ class RouterFacadeTest {
RouterState(sessionId = sessionId, workflowStatus = WorkflowStatus.IDLE, currentStageId = null)
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter("response"),
eventStore = mockStore,
@@ -157,7 +157,7 @@ class RouterFacadeTest {
)
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
capturedStates.add(state)
return emptyContextPack()
}
@@ -190,7 +190,7 @@ class RouterFacadeTest {
RouterState(sessionId = sessionId, workflowStatus = WorkflowStatus.RUNNING)
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
capturedStates.add(state)
return emptyContextPack()
}
@@ -218,7 +218,7 @@ class RouterFacadeTest {
fun `state is passed through to context builder`(): Unit = runBlocking {
val capturedState = mutableListOf<RouterState>()
val mockContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
capturedState.add(state)
return emptyContextPack()
}
@@ -288,7 +288,7 @@ class RouterFacadeTest {
)
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter,
eventStore = mockEventStore(),
@@ -319,7 +319,7 @@ class RouterFacadeTest {
)
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter,
eventStore = mockEventStore(),
@@ -345,7 +345,7 @@ class RouterFacadeTest {
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter,
eventStore = mockEventStore(),
@@ -373,7 +373,7 @@ class RouterFacadeTest {
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter,
eventStore = mockEventStore(),
@@ -394,7 +394,7 @@ class RouterFacadeTest {
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack {
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack {
val pack = ContextPack(
id = ContextPackId("test-pack"),
sessionId = state.sessionId ?: SessionId("unknown"),
@@ -449,7 +449,7 @@ class RouterFacadeTest {
override suspend fun getRouterState(sessionId: SessionId): RouterState = RouterState()
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = object : InferenceRouter {
override suspend fun route(
@@ -497,7 +497,7 @@ class RouterFacadeTest {
)
},
routerContextBuilder = object : RouterContextBuilder {
override fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
override suspend fun build(state: RouterState, budget: TokenBudget): ContextPack = emptyContextPack()
},
inferenceRouter = mockInferenceRouter("inference response"),
eventStore = eventStore,