From 127cc1b64e505c878cc0ea3b181c37c632bf6cd5 Mon Sep 17 00:00:00 2001
From: "Jakub A. W" <jakubwasek@gmail.com>
Date: Thu, 6 Aug 2026 00:44:58 +0200
Subject: [PATCH] feat(route): expose affinity context on attempt outcomes

---
 ext/route.go             |  5 +++++
 internal/app/app.go      | 14 +++++++++++++-
 internal/app/app_test.go | 30 ++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/ext/route.go b/ext/route.go
index 24ce2cf3..1650d3d9 100644
--- a/ext/route.go
+++ b/ext/route.go
@@ -54,6 +54,11 @@ func (t RouteTarget) Qualified() string { return t.Provider + "/" + t.Model }
 // scores slower than a target that succeeds at once.
 type RouteOutcome struct {
 	RouteTarget
+	// Source is the virtual model originally addressed when this attempt was
+	// selected through one. SessionID is the detected client session. Together
+	// they let selectors update cache affinity only after a successful attempt.
+	Source    string
+	SessionID string
 	// Endpoint is the upstream API endpoint (e.g. "/chat/completions").
 	Endpoint string
 	// StatusCode is the final upstream HTTP status; 0 on a network error.
diff --git a/internal/app/app.go b/internal/app/app.go
index 5e6af7b2..608c611f 100644
--- a/internal/app/app.go
+++ b/internal/app/app.go
@@ -146,10 +146,13 @@ func routeSelectorHooks(selector ext.RouteSelector) llmclient.Hooks {
 			})
 			return ctx
 		},
-		OnRequestEnd: func(_ context.Context, info llmclient.ResponseInfo) {
+		OnRequestEnd: func(ctx context.Context, info llmclient.ResponseInfo) {
 			observe("attempt_end", func() {
+				source, sessionID := routeAffinityContext(ctx)
 				selector.OnAttemptEnd(ext.RouteOutcome{
 					RouteTarget: ext.RouteTarget{Provider: info.Provider, Model: info.Model},
+					Source:      source,
+					SessionID:   sessionID,
 					Endpoint:    info.Endpoint,
 					StatusCode:  info.StatusCode,
 					Duration:    info.Duration,
@@ -161,6 +164,15 @@ func routeSelectorHooks(selector ext.RouteSelector) llmclient.Hooks {
 	}
 }
 
+func routeAffinityContext(ctx context.Context) (source, sessionID string) {
+	sessionID = core.SessionIDFromContext(ctx)
+	workflow := core.GetWorkflow(ctx)
+	if workflow == nil || workflow.Resolution == nil || !workflow.Resolution.AliasApplied {
+		return "", sessionID
+	}
+	return workflow.Resolution.RequestedQualifiedModel(), sessionID
+}
+
 // selectorLabel returns the selector's name for logs, tolerating a panicking
 // Name implementation, so recovery paths never re-enter extension code.
 func selectorLabel(selector ext.RouteSelector) (name string) {
diff --git a/internal/app/app_test.go b/internal/app/app_test.go
index 2b315b38..b944b819 100644
--- a/internal/app/app_test.go
+++ b/internal/app/app_test.go
@@ -17,10 +17,40 @@ import (
 	"github.com/enterpilot/gomodel/internal/core"
 	"github.com/enterpilot/gomodel/internal/guardrails"
 	"github.com/enterpilot/gomodel/internal/live"
+	"github.com/enterpilot/gomodel/internal/llmclient"
 	"github.com/enterpilot/gomodel/internal/providers"
 	"github.com/enterpilot/gomodel/internal/server"
 )
 
+type routeObservationSelector struct {
+	outcome ext.RouteOutcome
+}
+
+func (*routeObservationSelector) Name() string                           { return "observer" }
+func (*routeObservationSelector) Select(ext.RouteRequest) (string, bool) { return "", false }
+func (*routeObservationSelector) OnAttemptStart(ext.RouteTarget)         {}
+func (s *routeObservationSelector) OnAttemptEnd(outcome ext.RouteOutcome) {
+	s.outcome = outcome
+}
+
+func TestRouteSelectorHooksExposeSuccessfulRouteAffinityContext(t *testing.T) {
+	selector := &routeObservationSelector{}
+	hooks := routeSelectorHooks(selector)
+	ctx := core.WithSessionID(context.Background(), "session-a")
+	ctx = core.WithWorkflow(ctx, &core.Workflow{Resolution: &core.RequestModelResolution{
+		Requested:        core.NewRequestedModelSelector("smart", ""),
+		ResolvedSelector: core.ModelSelector{Provider: "openai", Model: "gpt"},
+		AliasApplied:     true,
+	}})
+	ctx = hooks.OnRequestStart(ctx, llmclient.RequestInfo{Provider: "openai", Model: "gpt"})
+	hooks.OnRequestEnd(ctx, llmclient.ResponseInfo{Provider: "openai", Model: "gpt", StatusCode: http.StatusOK})
+
+	if selector.outcome.Source != "smart" || selector.outcome.SessionID != "session-a" {
+		t.Fatalf("route affinity context = %q/%q, want smart/session-a",
+			selector.outcome.Source, selector.outcome.SessionID)
+	}
+}
+
 type runtimeRefreshMockProvider struct {
 	models *core.ModelsResponse
 	err    error