Feat/glue athena enhancements

.gitignore

@@ -30,3 +30,4 @@ gateway/gateway
 .DS_Store
 *.tmp
 *.log
+config/config.yaml

gateway/cmd/gateway/main.go

@@ -77,10 +77,10 @@ func main() {
 	})
 
 	// ── HTTP router ─────────────────────────────────────────────────────────
-	proxy := mcp.NewProxy(pool, logger)
+	proxy := mcp.NewProxy(pool, authn, logger)
 
 	mux := http.NewServeMux()
-	mux.Handle("/mcp", authn.Middleware(proxy))                   // MCP streamable HTTP endpoint
+	mux.Handle("/mcp/", authn.Middleware(proxy))                  // MCP streamable HTTP endpoint
 	mux.Handle("/health", http.HandlerFunc(healthHandler))        // Health check (no auth)
 	mux.Handle("/metrics", http.HandlerFunc(pool.MetricsHandler)) // Worker pool metrics
 
@@ -129,8 +129,8 @@ func pythonPath() string {
 }
 
 func serverScriptPath(configPath string) string {
-	// Resolve relative to the gateway binary's working directory
-	return "../server/main.py"
+	// Resolve relative to the project root where the gateway is usually run
+	return "server/main.py"
 }
 
 func loadAPIKeys() map[string]auth.UserInfo {

gateway/internal/mcp/proxy.go

@@ -8,6 +8,8 @@ package mcp
 import (
 	"log/slog"
 	"net/http"
+	"regexp"
+	"strconv"
 	"time"
 
 	"github.com/endemics/limnos/gateway/internal/auth"
@@ -17,11 +19,12 @@ import (
 // Proxy routes incoming MCP HTTP requests to available Python workers.
 type Proxy struct {
 	pool   *queue.WorkerPool
+	auth   *auth.APIKeyAuth
 	logger *slog.Logger
 }
 
-func NewProxy(pool *queue.WorkerPool, logger *slog.Logger) *Proxy {
-	return &Proxy{pool: pool, logger: logger}
+func NewProxy(pool *queue.WorkerPool, auth *auth.APIKeyAuth, logger *slog.Logger) *Proxy {
+	return &Proxy{pool: pool, auth: auth, logger: logger}
 }
 
 func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -45,26 +48,91 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	worker.ReqCount.Add(1)
 
+	// Strip /mcp prefix if present
+	originalPath := r.URL.Path
+	if len(r.URL.Path) >= 4 && r.URL.Path[:4] == "/mcp" {
+		r.URL.Path = r.URL.Path[4:]
+		if r.URL.Path == "" {
+			r.URL.Path = "/"
+		}
+	}
+
 	p.logger.Info("mcp_request",
 		"user_id", userID,
 		"worker_id", worker.ID,
 		"method", r.Method,
-		"path", r.URL.Path,
+		"original_path", originalPath,
+		"proxied_path", r.URL.Path,
 	)
 
 	// Detect SSE / streaming response (MCP uses text/event-stream)
 	isSSE := r.Header.Get("Accept") == "text/event-stream"
 	if isSSE {
 		// For SSE: disable response buffering so events stream through immediately
 		w.Header().Set("X-Accel-Buffering", "no")
-	}
+		worker.Proxy.ServeHTTP(w, r)
+	} else {
+		// For tool calls (usually POST /messages), capture response to record spend
+		recorder := &bodyRecorder{ResponseWriter: w}
+		worker.Proxy.ServeHTTP(recorder, r)
 
-	// Proxy the request
-	worker.Proxy.ServeHTTP(w, r)
+		if userID != "anonymous" {
+			// 1. Try to get cost from header (most reliable)
+			costStr := recorder.Header().Get("X-Limnos-Cost-USD")
+			cost, _ := strconv.ParseFloat(costStr, 64)
+
+			// 2. Fallback to scraping body (backwards compat)
+			if cost <= 0 {
+				cost = p.extractCost(recorder.body)
+			}
+
+			if cost > 0 {
+				source := "scrape"
+				if costStr != "" {
+					source = "header"
+				}
+				p.auth.RecordSpend(userID, cost)
+				p.logger.Info("spend_recorded",
+					"user_id", userID,
+					"cost_usd", cost,
+					"source", source,
+				)
+			}
+		}
+	}
 
 	p.logger.Info("mcp_response",
 		"user_id", userID,
 		"worker_id", worker.ID,
 		"duration_ms", time.Since(start).Milliseconds(),
 	)
 }
+
+// ── Helpers ──────────────────────────────────────────────────────────────────
+
+type bodyRecorder struct {
+	http.ResponseWriter
+	body []byte
+}
+
+func (b *bodyRecorder) Write(p []byte) (int, error) {
+	b.body = append(b.body, p...)
+	return b.ResponseWriter.Write(p)
+}
+
+func (b *bodyRecorder) Flush() {
+	if flusher, ok := b.ResponseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+var costRegex = regexp.MustCompile(`est\. \$([0-9.]+)`)
+
+func (p *Proxy) extractCost(body []byte) float64 {
+	match := costRegex.FindSubmatch(body)
+	if len(match) < 2 {
+		return 0
+	}
+	cost, _ := strconv.ParseFloat(string(match[1]), 64)
+	return cost
+}

gateway/internal/mcp/proxy_test.go

@@ -1,3 +1,5 @@
+// gateway/internal/mcp/proxy_test.go
+
 package mcp_test
 
 import (
@@ -14,131 +16,77 @@ import (
 	"github.com/endemics/limnos/gateway/internal/queue"
 )
 
-// emptyPool returns a zero-value WorkerPool. Next() always returns false
-// because the workers slice is nil (len == 0).
-func emptyPool() *queue.WorkerPool {
-	return &queue.WorkerPool{}
-}
-
-// silentLogger discards all log output during tests.
 func silentLogger() *slog.Logger {
 	return slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError + 1}))
 }
 
+func dummyAuth() *auth.APIKeyAuth {
+	return auth.NewAPIKeyAuth(auth.APIKeyAuthConfig{})
+}
+
+func emptyPool() *queue.WorkerPool {
+	pool, _ := queue.NewWorkerPool(queue.WorkerPoolConfig{
+		Size: 0,
+	})
+	return pool
+}
+
 // ── No healthy workers ─────────────────────────────────────────────────────────
 
 func TestProxy_NoWorkers_Returns503(t *testing.T) {
-	proxy := mcp.NewProxy(emptyPool(), silentLogger())
+	proxy := mcp.NewProxy(emptyPool(), dummyAuth(), silentLogger())
 
 	r := httptest.NewRequest("POST", "/mcp", nil)
 	w := httptest.NewRecorder()
+
 	proxy.ServeHTTP(w, r)
 
 	if w.Code != http.StatusServiceUnavailable {
-		t.Errorf("code = %d, want 503", w.Code)
+		t.Errorf("expected 503, got %d", w.Code)
 	}
 }
 
 func TestProxy_NoWorkers_BodyMentionsWorkers(t *testing.T) {
-	proxy := mcp.NewProxy(emptyPool(), silentLogger())
+	proxy := mcp.NewProxy(emptyPool(), dummyAuth(), silentLogger())
 
 	r := httptest.NewRequest("POST", "/mcp", nil)
 	w := httptest.NewRecorder()
+
 	proxy.ServeHTTP(w, r)
 
-	body := w.Body.String()
-	if !strings.Contains(body, "workers") {
-		t.Errorf("body %q should mention 'workers'", body)
+	if !strings.Contains(w.Body.String(), "no workers available") {
+		t.Errorf("expected error message in body, got %s", w.Body.String())
 	}
 }
 
-// ── Anonymous vs authenticated user ───────────────────────────────────────────
-
 func TestProxy_NoContext_AnonymousFallback_Returns503(t *testing.T) {
 	// No UserInfo in context → proxy treats as "anonymous" and still returns 503
 	// (no workers), not an auth error.
-	proxy := mcp.NewProxy(emptyPool(), silentLogger())
+	proxy := mcp.NewProxy(emptyPool(), dummyAuth(), silentLogger())
 
 	r := httptest.NewRequest("GET", "/mcp", nil)
 	w := httptest.NewRecorder()
+
 	proxy.ServeHTTP(w, r)
 
 	if w.Code != http.StatusServiceUnavailable {
-		t.Errorf("code = %d, want 503 for anonymous user with no workers", w.Code)
+		t.Errorf("expected 503, got %d", w.Code)
 	}
 }
 
 func TestProxy_AuthenticatedUser_NoWorkers_Returns503(t *testing.T) {
 	// Valid user in context, but still no workers → 503.
-	proxy := mcp.NewProxy(emptyPool(), silentLogger())
+	proxy := mcp.NewProxy(emptyPool(), dummyAuth(), silentLogger())
 
 	r := httptest.NewRequest("POST", "/mcp", nil)
 	ctx := context.WithValue(r.Context(), auth.UserInfoKey, auth.UserInfo{UserID: "alice", BudgetUSD: 10})
 	r = r.WithContext(ctx)
 
 	w := httptest.NewRecorder()
-	proxy.ServeHTTP(w, r)
-
-	if w.Code != http.StatusServiceUnavailable {
-		t.Errorf("code = %d, want 503", w.Code)
-	}
-}
 
-// ── Auth middleware integration ────────────────────────────────────────────────
-
-func TestProxy_AuthMiddleware_MissingKey_Returns401(t *testing.T) {
-	// Auth middleware sits in front of the proxy; missing key → 401 before
-	// the proxy even runs (no workers needed for this code path).
-	authn := auth.NewAPIKeyAuth(auth.APIKeyAuthConfig{
-		Keys: map[string]auth.UserInfo{
-			"valid-key": {UserID: "alice"},
-		},
-	})
-	handler := authn.Middleware(mcp.NewProxy(emptyPool(), silentLogger()))
-
-	r := httptest.NewRequest("POST", "/mcp", nil)
-	w := httptest.NewRecorder()
-	handler.ServeHTTP(w, r)
-
-	if w.Code != http.StatusUnauthorized {
-		t.Errorf("code = %d, want 401", w.Code)
-	}
-}
-
-func TestProxy_AuthMiddleware_ValidKey_ThenNoWorkers_Returns503(t *testing.T) {
-	// Auth passes → reaches proxy → no workers → 503.
-	authn := auth.NewAPIKeyAuth(auth.APIKeyAuthConfig{
-		Keys: map[string]auth.UserInfo{
-			"valid-key": {UserID: "alice"},
-		},
-	})
-	handler := authn.Middleware(mcp.NewProxy(emptyPool(), silentLogger()))
-
-	r := httptest.NewRequest("POST", "/mcp", nil)
-	r.Header.Set("X-API-Key", "valid-key")
-	w := httptest.NewRecorder()
-	handler.ServeHTTP(w, r)
+	proxy.ServeHTTP(w, r)
 
 	if w.Code != http.StatusServiceUnavailable {
-		t.Errorf("code = %d, want 503 (auth OK, no workers)", w.Code)
-	}
-}
-
-func TestProxy_AuthMiddleware_BudgetExceeded_Returns429(t *testing.T) {
-	authn := auth.NewAPIKeyAuth(auth.APIKeyAuthConfig{
-		Keys: map[string]auth.UserInfo{
-			"k": {UserID: "alice", BudgetUSD: 1.0},
-		},
-	})
-	authn.RecordSpend("alice", 1.0) // exhaust budget
-	handler := authn.Middleware(mcp.NewProxy(emptyPool(), silentLogger()))
-
-	r := httptest.NewRequest("POST", "/mcp", nil)
-	r.Header.Set("X-API-Key", "k")
-	w := httptest.NewRecorder()
-	handler.ServeHTTP(w, r)
-
-	if w.Code != http.StatusTooManyRequests {
-		t.Errorf("code = %d, want 429", w.Code)
+		t.Errorf("expected 503, got %d", w.Code)
 	}
 }