feat(auth): track unique tools in access metrics
Some checks failed
CI / build-and-test (push) Failing after -31m49s
Some checks failed
CI / build-and-test (push) Failing after -31m49s
* Add tool tracking to AccessTracker and metrics * Update tests to validate tool tracking functionality * Modify middleware to record tool usage * Enhance observability with tool context * Update UI to display unique tools in metrics
This commit is contained in:
@@ -1,6 +1,8 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
@@ -8,6 +10,7 @@ import (
|
||||
"testing"
|
||||
|
||||
"git.warky.dev/wdevs/amcs/internal/config"
|
||||
"git.warky.dev/wdevs/amcs/internal/observability"
|
||||
)
|
||||
|
||||
func testLogger() *slog.Logger {
|
||||
@@ -188,3 +191,50 @@ func TestMiddlewareRecordsForwardedRemoteAddr(t *testing.T) {
|
||||
t.Fatalf("snapshot remote_addr = %q, want %q", snap[0].RemoteAddr, "203.0.113.99")
|
||||
}
|
||||
}
|
||||
|
||||
func TestMiddlewareRecordsMCPToolUsage(t *testing.T) {
|
||||
keyring, err := NewKeyring([]config.APIKey{{ID: "client-a", Value: "secret"}})
|
||||
if err != nil {
|
||||
t.Fatalf("NewKeyring() error = %v", err)
|
||||
}
|
||||
tracker := NewAccessTracker()
|
||||
logger := testLogger()
|
||||
|
||||
authenticated := Middleware(config.AuthConfig{HeaderName: "x-brain-key"}, keyring, nil, nil, tracker, logger)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
}))
|
||||
handler := observability.AccessLog(logger)(authenticated)
|
||||
|
||||
payload := map[string]any{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "1",
|
||||
"method": "tools/call",
|
||||
"params": map[string]any{
|
||||
"name": "list_projects",
|
||||
},
|
||||
}
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
t.Fatalf("json.Marshal() error = %v", err)
|
||||
}
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/mcp", bytes.NewReader(body))
|
||||
req.Header.Set("x-brain-key", "secret")
|
||||
rec := httptest.NewRecorder()
|
||||
handler.ServeHTTP(rec, req)
|
||||
|
||||
if rec.Code != http.StatusNoContent {
|
||||
t.Fatalf("status = %d, want %d", rec.Code, http.StatusNoContent)
|
||||
}
|
||||
|
||||
metrics := tracker.Metrics(10)
|
||||
if metrics.UniqueTools != 1 {
|
||||
t.Fatalf("UniqueTools = %d, want 1", metrics.UniqueTools)
|
||||
}
|
||||
if len(metrics.TopTools) != 1 {
|
||||
t.Fatalf("len(TopTools) = %d, want 1", len(metrics.TopTools))
|
||||
}
|
||||
if metrics.TopTools[0].Key != "list_projects" || metrics.TopTools[0].RequestCount != 1 {
|
||||
t.Fatalf("TopTools[0] = %+v, want list_projects with count 1", metrics.TopTools[0])
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user