From 3881f69aa049fdb4d7cd283045ad5fd8004e095d Mon Sep 17 00:00:00 2001
From: "Tim Z." <91899951+aimbotterz@users.noreply.github.com>
Date: Tue, 7 Apr 2026 13:56:19 +0200
Subject: [PATCH] feat(routing): inject native ML complexity Router scoring
 into telegram wrapper header

---
 pkg/agent/loop.go    | 25 +++++++++++++++----------
 pkg/agent/metrics.go |  4 ----
 pkg/agent/turn.go    | 15 +++++++++++++++
 3 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go
index 4247cd516..b3fffd64d 100644
--- a/pkg/agent/loop.go
+++ b/pkg/agent/loop.go
@@ -1624,11 +1624,9 @@ func (al *AgentLoop) runAgentLoop(
 	}
 
 	if result.finalContent != "" {
-		var totalTokens, promptTokens, compTokens int
+		var totalTokens int
 		if usage := ts.GetLastUsage(); usage != nil {
 			totalTokens = usage.TotalTokens
-			promptTokens = usage.PromptTokens
-			compTokens = usage.CompletionTokens
 		}
 
 		version := al.GetConfig().BuildInfo.Version
@@ -1644,7 +1642,7 @@ func (al *AgentLoop) runAgentLoop(
 			Agent:      agent.ID,
 			Route:      providerName,
 			Model:      modelName,
-			Complexity: calculateComplexity(promptTokens, compTokens),
+			Complexity: int(ts.getComplexityScore() * 100),
 			Tokens:     totalTokens,
 			ToolCalls:  ts.currentIteration(),
 			Processing: time.Since(ts.startedAt),
@@ -1850,7 +1848,8 @@ func (al *AgentLoop) runTurn(ctx context.Context, ts *turnState) (turnResult, er
 		ts.ingestMessage(turnCtx, al, rootMsg)
 	}
 
-	activeCandidates, activeModel, usedLight := al.selectCandidates(ts.agent, ts.userMessage, messages)
+	activeCandidates, activeModel, usedLight, score := al.selectCandidates(ts.agent, ts.userMessage, messages)
+	ts.setComplexityScore(score)
 	activeProvider := ts.agent.Provider
 	if usedLight && ts.agent.LightProvider != nil {
 		activeProvider = ts.agent.LightProvider
@@ -2977,12 +2976,18 @@ func (al *AgentLoop) selectCandidates(
 	agent *AgentInstance,
 	userMsg string,
 	history []providers.Message,
-) (candidates []providers.FallbackCandidate, model string, usedLight bool) {
+) (candidates []providers.FallbackCandidate, model string, usedLight bool, score float64) {
 	if agent.Router == nil || len(agent.LightCandidates) == 0 {
-		return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false
+		return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false, 0
 	}
 
-	_, usedLight, score := agent.Router.SelectModel(userMsg, history, agent.Model)
+	_, usedLight, score = agent.Router.SelectModel(userMsg, history, agent.Model)
+	
+	// Cache the native rule-based complexity score to the primary turnState
+	if ts := turnStateFromContext(context.Background()); ts != nil {
+		// Note: Since we don't have direct access to ts here, we actually pass it up natively in runTurn
+	}
+	
 	if !usedLight {
 		logger.DebugCF("agent", "Model routing: primary model selected",
 			map[string]any{
@@ -2990,7 +2995,7 @@ func (al *AgentLoop) selectCandidates(
 				"score":     score,
 				"threshold": agent.Router.Threshold(),
 			})
-		return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false
+		return agent.Candidates, resolvedCandidateModel(agent.Candidates, agent.Model), false, score
 	}
 
 	logger.InfoCF("agent", "Model routing: light model selected",
@@ -3000,7 +3005,7 @@ func (al *AgentLoop) selectCandidates(
 			"score":       score,
 			"threshold":   agent.Router.Threshold(),
 		})
-	return agent.LightCandidates, resolvedCandidateModel(agent.LightCandidates, agent.Router.LightModel()), true
+	return agent.LightCandidates, resolvedCandidateModel(agent.LightCandidates, agent.Router.LightModel()), true, score
 }
 
 // resolveContextManager selects the ContextManager implementation based on config.
diff --git a/pkg/agent/metrics.go b/pkg/agent/metrics.go
index a30c83a5f..fe5e60ee0 100644
--- a/pkg/agent/metrics.go
+++ b/pkg/agent/metrics.go
@@ -23,7 +23,3 @@ func WrapResponse(rawOutput string, m Metrics) string {
 	)
 	return header + rawOutput
 }
-
-func calculateComplexity(promptTokens, completionTokens int) int {
-	return int(float64(promptTokens)*0.1 + float64(completionTokens)*0.5)
-}
diff --git a/pkg/agent/turn.go b/pkg/agent/turn.go
index 8f099ed1d..c5099252f 100644
--- a/pkg/agent/turn.go
+++ b/pkg/agent/turn.go
@@ -103,6 +103,9 @@ type turnState struct {
 	tokenBudget      *atomic.Int64        // Shared token budget counter
 	lastFinishReason string               // Last LLM finish_reason
 	lastUsage        *providers.UsageInfo // Last LLM usage info
+	
+	// Telemetry tracking
+	complexityScore float64 // Cached heuristic difficulty score generated by the router
 
 	// Back-reference to the owning AgentLoop (set for SubTurns only, used for hard abort cascade)
 	al *AgentLoop
@@ -214,6 +217,18 @@ func (ts *turnState) setIteration(iteration int) {
 	ts.iteration = iteration
 }
 
+func (ts *turnState) setComplexityScore(score float64) {
+	ts.mu.Lock()
+	defer ts.mu.Unlock()
+	ts.complexityScore = score
+}
+
+func (ts *turnState) getComplexityScore() float64 {
+	ts.mu.RLock()
+	defer ts.mu.RUnlock()
+	return ts.complexityScore
+}
+
 func (ts *turnState) currentIteration() int {
 	ts.mu.RLock()
 	defer ts.mu.RUnlock()