Add backend routing capture to vMCP audit logs (#2983)

* Add backend routing capture to vMCP audit logs

Enhance audit logging to capture which backend workload handled each
MCP request (tools/call, resources/read, prompts/get). Backend routing
information is now recorded in audit events as backend_name in the
metadata.extra field.

* changes from review

---------

Co-authored-by: taskbot <[email protected]>

Author: yrobla

pkg/audit/auditor.go

@@ -3,6 +3,7 @@ package audit
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"log/slog"
@@ -21,6 +22,30 @@ import (
 // LevelAudit is a custom audit log level - between Info and Warn
 const LevelAudit = slog.Level(2)
 
+// contextKey is an unexported type for context keys to avoid collisions
+type contextKey struct{}
+
+// backendInfoKey is the context key for storing backend routing information
+var backendInfoKey = contextKey{}
+
+// BackendInfo stores backend routing information that can be mutated by handlers.
+// This allows handlers deep in the call stack to provide backend info to the audit middleware.
+type BackendInfo struct {
+	BackendName string
+}
+
+// WithBackendInfo returns a new context with BackendInfo attached.
+func WithBackendInfo(ctx context.Context, info *BackendInfo) context.Context {
+	return context.WithValue(ctx, backendInfoKey, info)
+}
+
+// BackendInfoFromContext retrieves BackendInfo from the context.
+// Returns (nil, false) if BackendInfo is not found in the context.
+func BackendInfoFromContext(ctx context.Context) (*BackendInfo, bool) {
+	info, ok := ctx.Value(backendInfoKey).(*BackendInfo)
+	return info, ok
+}
+
 // NewAuditLogger creates a new structured audit logger that writes to the specified writer.
 func NewAuditLogger(w io.Writer) *slog.Logger {
 	if w == nil {
@@ -129,6 +154,14 @@ func (a *Auditor) Middleware(next http.Handler) http.Handler {
 
 		startTime := time.Now()
 
+		// Add BackendInfo to context if not already present
+		// (backend enrichment middleware may have already added it)
+		if _, ok := BackendInfoFromContext(r.Context()); !ok {
+			backendInfo := &BackendInfo{}
+			ctx := WithBackendInfo(r.Context(), backendInfo)
+			r = r.WithContext(ctx)
+		}
+
 		// Capture request data if configured
 		var requestData []byte
 		if a.config.IncludeRequestData && r.Body != nil {
@@ -194,7 +227,7 @@ func (a *Auditor) logAuditEvent(r *http.Request, rw *responseWriter, requestData
 	}
 
 	// Add metadata
-	a.addMetadata(event, duration, rw)
+	a.addMetadata(event, r, duration, rw)
 
 	// Add request/response data if configured
 	a.addEventData(event, r, rw, requestData)
@@ -398,7 +431,7 @@ func (*Auditor) extractTarget(r *http.Request, eventType string) map[string]stri
 }
 
 // addMetadata adds metadata to the audit event.
-func (a *Auditor) addMetadata(event *AuditEvent, duration time.Duration, rw *responseWriter) {
+func (a *Auditor) addMetadata(event *AuditEvent, r *http.Request, duration time.Duration, rw *responseWriter) {
 	if event.Metadata.Extra == nil {
 		event.Metadata.Extra = make(map[string]any)
 	}
@@ -417,6 +450,12 @@ func (a *Auditor) addMetadata(event *AuditEvent, duration time.Duration, rw *res
 	if rw.body != nil {
 		event.Metadata.Extra[MetadataExtraKeyResponseSize] = rw.body.Len()
 	}
+
+	// Add backend routing information from context if available
+	// Backend info is populated by the backend enrichment middleware
+	if backendInfo, ok := BackendInfoFromContext(r.Context()); ok && backendInfo != nil && backendInfo.BackendName != "" {
+		event.Metadata.Extra["backend_name"] = backendInfo.BackendName
+	}
 }
 
 // addEventData adds request/response data to the audit event if configured.

pkg/audit/auditor_test.go

@@ -517,8 +517,9 @@ func TestAddMetadata(t *testing.T) {
 		ResponseWriter: httptest.NewRecorder(),
 		body:           bytes.NewBufferString("test response"),
 	}
+	req := httptest.NewRequest("GET", "/test", nil)
 
-	auditor.addMetadata(event, duration, rw)
+	auditor.addMetadata(event, req, duration, rw)
 
 	require.NotNil(t, event.Metadata.Extra)
 	assert.Equal(t, int64(150), event.Metadata.Extra[MetadataExtraKeyDuration])

pkg/audit/backend_info_test.go

@@ -0,0 +1,79 @@
+package audit
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBackendInfoContext(t *testing.T) {
+	t.Parallel()
+
+	t.Run("BackendInfo can be added and retrieved from context", func(t *testing.T) {
+		t.Parallel()
+
+		// Create a BackendInfo
+		info := &BackendInfo{
+			BackendName: "test-backend",
+		}
+
+		// Add it to context
+		ctx := WithBackendInfo(context.Background(), info)
+
+		// Retrieve it
+		retrieved, ok := BackendInfoFromContext(ctx)
+		require.True(t, ok, "BackendInfo should be in context")
+		require.NotNil(t, retrieved, "BackendInfo should not be nil")
+		assert.Equal(t, "test-backend", retrieved.BackendName)
+
+		// Verify it's the same pointer
+		assert.Same(t, info, retrieved, "Should be the same BackendInfo pointer")
+	})
+
+	t.Run("BackendInfo can be mutated through context", func(t *testing.T) {
+		t.Parallel()
+
+		// Create empty BackendInfo
+		info := &BackendInfo{}
+
+		// Add to context
+		ctx := WithBackendInfo(context.Background(), info)
+
+		// Retrieve and mutate
+		retrieved, ok := BackendInfoFromContext(ctx)
+		require.True(t, ok)
+		retrieved.BackendName = "mutated-backend"
+
+		// Verify original was mutated
+		assert.Equal(t, "mutated-backend", info.BackendName)
+	})
+
+	t.Run("Missing BackendInfo returns false", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+
+		retrieved, ok := BackendInfoFromContext(ctx)
+		assert.False(t, ok, "Should return false when not in context")
+		assert.Nil(t, retrieved, "Should return nil when not in context")
+	})
+
+	t.Run("BackendInfo survives context derivation", func(t *testing.T) {
+		t.Parallel()
+
+		// Create BackendInfo and add to context
+		info := &BackendInfo{BackendName: "original"}
+		ctx := WithBackendInfo(context.Background(), info)
+
+		// Derive a new context with additional value
+		type key struct{}
+		derivedCtx := context.WithValue(ctx, key{}, "some-value")
+
+		// BackendInfo should still be accessible
+		retrieved, ok := BackendInfoFromContext(derivedCtx)
+		require.True(t, ok, "BackendInfo should survive context derivation")
+		assert.Equal(t, "original", retrieved.BackendName)
+	})
+}

pkg/vmcp/server/backend_enrichment.go

@@ -0,0 +1,85 @@
+package server
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"net/http"
+
+	"github.com/stacklok/toolhive/pkg/audit"
+	"github.com/stacklok/toolhive/pkg/logger"
+	"github.com/stacklok/toolhive/pkg/vmcp"
+	"github.com/stacklok/toolhive/pkg/vmcp/discovery"
+)
+
+// backendEnrichmentMiddleware wraps an HTTP handler to add backend routing information
+// to audit events by parsing MCP requests and looking up backends in the routing table.
+func (*Server) backendEnrichmentMiddleware(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Read and parse the request body to extract MCP method and parameters
+		var requestBody []byte
+		if r.Body != nil {
+			var err error
+			requestBody, err = io.ReadAll(r.Body)
+			// Always restore body for next handler, even on error
+			if err != nil {
+				// Log the error and restore an empty body to ensure consistent behavior
+				logger.Warnw("failed to read request body in backend enrichment middleware",
+					"error", err)
+				r.Body = io.NopCloser(bytes.NewReader([]byte{}))
+			} else {
+				// Restore body with the read content
+				r.Body = io.NopCloser(bytes.NewReader(requestBody))
+			}
+		}
+
+		// Parse MCP request to extract tool/resource name
+		var mcpRequest struct {
+			Method string         `json:"method"`
+			Params map[string]any `json:"params"`
+		}
+
+		if len(requestBody) > 0 && json.Unmarshal(requestBody, &mcpRequest) == nil {
+			// Get routing table from discovered capabilities in context
+			caps, ok := discovery.DiscoveredCapabilitiesFromContext(r.Context())
+			if ok && caps != nil && caps.RoutingTable != nil {
+				backendName := lookupBackendName(mcpRequest.Method, mcpRequest.Params, caps.RoutingTable)
+
+				// Mutate the existing BackendInfo from audit middleware
+				if backendName != "" {
+					if backendInfo, ok := audit.BackendInfoFromContext(r.Context()); ok && backendInfo != nil {
+						backendInfo.BackendName = backendName
+					}
+				}
+			}
+		}
+
+		// Call next handler
+		next.ServeHTTP(w, r)
+	})
+}
+
+// lookupBackendName looks up which backend handles a given MCP request.
+func lookupBackendName(method string, params map[string]any, routingTable *vmcp.RoutingTable) string {
+	switch method {
+	case "tools/call":
+		if toolName, ok := params["name"].(string); ok {
+			if target, exists := routingTable.Tools[toolName]; exists {
+				return target.WorkloadName
+			}
+		}
+	case "resources/read":
+		if uri, ok := params["uri"].(string); ok {
+			if target, exists := routingTable.Resources[uri]; exists {
+				return target.WorkloadName
+			}
+		}
+	case "prompts/get":
+		if promptName, ok := params["name"].(string); ok {
+			if target, exists := routingTable.Prompts[promptName]; exists {
+				return target.WorkloadName
+			}
+		}
+	}
+	return ""
+}