diff --git a/cmd/proxy/main.go b/cmd/proxy/main.go index 24bc3c65..2e9ffa36 100644 --- a/cmd/proxy/main.go +++ b/cmd/proxy/main.go @@ -75,6 +75,16 @@ func proxyLatencyFromProxyTimingHeaders(headers http.Header) (time.Duration, boo return finalAt.Sub(receivedAt), true } +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + + return "" +} + func newBenchmarkHTTPClient(timeout time.Duration) *http.Client { // NOTE: We intentionally reuse a single client+transport for all requests so we // get connection pooling/keep-alives. Creating a new client per request @@ -354,7 +364,7 @@ var benchmarkCmd *cobra.Command var dispatcherCmd = &cobra.Command{ Use: "dispatcher", Short: "Run the event dispatcher service", - Long: `Run the event dispatcher service with pluggable backends. Supports file, lunary, and helicone services.`, + Long: `Run the event dispatcher service with pluggable backends. Supports file, lunary, helicone, and cloudwatch services.`, Run: runDispatcher, } @@ -501,6 +511,17 @@ func runDispatcher(cmd *cobra.Command, args []string) { config["api-key"] = envKey } } + if dispatcherService == "cloudwatch" { + if logGroup := os.Getenv("DISPATCHER_CLOUDWATCH_LOG_GROUP"); logGroup != "" { + config["log-group"] = logGroup + } + if logStream := os.Getenv("DISPATCHER_CLOUDWATCH_LOG_STREAM"); logStream != "" { + config["log-stream"] = logStream + } + if region := firstNonEmpty(os.Getenv("DISPATCHER_CLOUDWATCH_REGION"), os.Getenv("AWS_REGION"), os.Getenv("AWS_DEFAULT_REGION")); region != "" { + config["region"] = region + } + } if err := plugin.Init(config); err != nil { logger.Fatal("Failed to initialize plugin", zap.Error(err)) @@ -1590,7 +1611,7 @@ Latency breakdown: cobraRoot.PersistentFlags().StringVar(&manageAPIBaseURL, "manage-api-base-url", "http://localhost:8080", "Base URL for management API (default: http://localhost:8080)") // Add dispatcher command flags - dispatcherCmd.Flags().StringVar(&dispatcherService, "service", config.EnvOrDefault("DISPATCHER_SERVICE", "file"), "Dispatcher service type (file, lunary, helicone)") + dispatcherCmd.Flags().StringVar(&dispatcherService, "service", config.EnvOrDefault("DISPATCHER_SERVICE", "file"), "Dispatcher service type (file, lunary, helicone, cloudwatch)") dispatcherCmd.Flags().StringVar(&dispatcherEndpoint, "endpoint", config.EnvOrDefault("DISPATCHER_ENDPOINT", ""), "Dispatcher endpoint URL (file: path, lunary/helicone: API endpoint)") dispatcherCmd.Flags().StringVar(&dispatcherAPIKey, "api-key", config.EnvOrDefault("LLM_PROXY_API_KEY", ""), "API key for external services (lunary, helicone)") dispatcherCmd.Flags().IntVar(&dispatcherBuffer, "buffer", config.EnvIntOrDefault("DISPATCHER_BUFFER", 1000), "Event bus buffer size") diff --git a/docs/guides/api-configuration.md b/docs/guides/api-configuration.md index 56ce28f3..ae297788 100644 --- a/docs/guides/api-configuration.md +++ b/docs/guides/api-configuration.md @@ -154,6 +154,7 @@ The caching system follows HTTP standards: - **Authentication**: Cached responses for authenticated requests are only served if marked as publicly cacheable (`Cache-Control: public` or `s-maxage` present) - **Streaming responses**: Captured during streaming and stored after completion - **TTL precedence**: `s-maxage` (shared cache) takes precedence over `max-age` +- **Optional namespace header**: Clients may send `X-LLM-Proxy-Cache-Namespace` to partition cache keys for logical invalidation; the value is sanitized, included in the proxy cache key, and not forwarded upstream - **Headers**: Responses include `X-PROXY-CACHE`, `X-PROXY-CACHE-KEY`, and `Cache-Status` for observability ### Cache Stats Aggregation diff --git a/go.mod b/go.mod index 3797aed5..d176cdaf 100644 --- a/go.mod +++ b/go.mod @@ -5,6 +5,8 @@ go 1.23.9 require ( github.com/alicebob/miniredis/v2 v2.31.0 github.com/andybalholm/brotli v1.2.0 + github.com/aws/aws-sdk-go-v2/config v1.31.0 + github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs v1.58.0 github.com/chzyer/readline v1.5.1 github.com/gin-contrib/sessions v1.0.4 github.com/gin-gonic/gin v1.10.1 @@ -28,6 +30,19 @@ require ( require ( filippo.io/edwards25519 v1.1.0 // indirect github.com/alicebob/gopher-json v0.0.0-20200520072559-a9ecdc9d1d3a // indirect + github.com/aws/aws-sdk-go-v2 v1.39.0 // indirect + github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.1 // indirect + github.com/aws/aws-sdk-go-v2/credentials v1.18.13 // indirect + github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.7 // indirect + github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.7 // indirect + github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.7 // indirect + github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.1 // indirect + github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.7 // indirect + github.com/aws/aws-sdk-go-v2/service/sso v1.29.3 // indirect + github.com/aws/aws-sdk-go-v2/service/ssooidc v1.34.5 // indirect + github.com/aws/aws-sdk-go-v2/service/sts v1.38.4 // indirect + github.com/aws/smithy-go v1.23.0 // indirect github.com/bytedance/sonic v1.13.2 // indirect github.com/bytedance/sonic/loader v0.2.4 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect diff --git a/go.sum b/go.sum index c9f9da95..832c0bec 100644 --- a/go.sum +++ b/go.sum @@ -7,6 +7,36 @@ github.com/alicebob/miniredis/v2 v2.31.0 h1:ObEFUNlJwoIiyjxdrYF0QIDE7qXcLc7D3WpS github.com/alicebob/miniredis/v2 v2.31.0/go.mod h1:UB/T2Uztp7MlFSDakaX1sTXUv5CASoprx0wulRT6HBg= github.com/andybalholm/brotli v1.2.0 h1:ukwgCxwYrmACq68yiUqwIWnGY0cTPox/M94sVwToPjQ= github.com/andybalholm/brotli v1.2.0/go.mod h1:rzTDkvFWvIrjDXZHkuS16NPggd91W3kUSvPlQ1pLaKY= +github.com/aws/aws-sdk-go-v2 v1.39.0 h1:xm5WV/2L4emMRmMjHFykqiA4M/ra0DJVSWUkDyBjbg4= +github.com/aws/aws-sdk-go-v2 v1.39.0/go.mod h1:sDioUELIUO9Znk23YVmIk86/9DOpkbyyVb1i/gUNFXY= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.1 h1:i8p8P4diljCr60PpJp6qZXNlgX4m2yQFpYk+9ZT+J4E= +github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.1/go.mod h1:ddqbooRZYNoJ2dsTwOty16rM+/Aqmk/GOXrK8cg7V00= +github.com/aws/aws-sdk-go-v2/config v1.31.0 h1:9yH0xiY5fUnVNLRWO0AtayqwU1ndriZdN78LlhruJR4= +github.com/aws/aws-sdk-go-v2/config v1.31.0/go.mod h1:VeV3K72nXnhbe4EuxxhzsDc/ByrCSlZwUnWH52Nde/I= +github.com/aws/aws-sdk-go-v2/credentials v1.18.13 h1:gkpEm65/ZfrGJ3wbFH++Ki7DyaWtsWbK9idX6OXCo2E= +github.com/aws/aws-sdk-go-v2/credentials v1.18.13/go.mod h1:eVTHz1yI2/WIlXTE8f70mcrSxNafXD5sJpTIM9f+kmo= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.7 h1:Is2tPmieqGS2edBnmOJIbdvOA6Op+rRpaYR60iBAwXM= +github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.7/go.mod h1:F1i5V5421EGci570yABvpIXgRIBPb5JM+lSkHF6Dq5w= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.7 h1:UCxq0X9O3xrlENdKf1r9eRJoKz/b0AfGkpp3a7FPlhg= +github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.7/go.mod h1:rHRoJUNUASj5Z/0eqI4w32vKvC7atoWR0jC+IkmVH8k= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.7 h1:Y6DTZUn7ZUC4th9FMBbo8LVE+1fyq3ofw+tRwkUd3PY= +github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.7/go.mod h1:x3XE6vMnU9QvHN/Wrx2s44kwzV2o2g5x/siw4ZUJ9g8= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3 h1:bIqFDwgGXXN1Kpp99pDOdKMTTb5d2KyU5X/BZxjOkRo= +github.com/aws/aws-sdk-go-v2/internal/ini v1.8.3/go.mod h1:H5O/EsxDWyU+LP/V8i5sm8cxoZgc2fdNR9bxlOFrQTo= +github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs v1.58.0 h1:XH0kj0KcoKd+BAadpiS83/Wf+25q4FmH3gDei4u+PzA= +github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs v1.58.0/go.mod h1:ptJgRWK9opQK1foOTBKUg3PokkKA0/xcTXWIxwliaIY= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.1 h1:oegbebPEMA/1Jny7kvwejowCaHz1FWZAQ94WXFNCyTM= +github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.1/go.mod h1:kemo5Myr9ac0U9JfSjMo9yHLtw+pECEHsFtJ9tqCEI8= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.7 h1:mLgc5QIgOy26qyh5bvW+nDoAppxgn3J2WV3m9ewq7+8= +github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.7/go.mod h1:wXb/eQnqt8mDQIQTTmcw58B5mYGxzLGZGK8PWNFZ0BA= +github.com/aws/aws-sdk-go-v2/service/sso v1.29.3 h1:7PKX3VYsZ8LUWceVRuv0+PU+E7OtQb1lgmi5vmUE9CM= +github.com/aws/aws-sdk-go-v2/service/sso v1.29.3/go.mod h1:Ql6jE9kyyWI5JHn+61UT/Y5Z0oyVJGmgmJbZD5g4unY= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.34.5 h1:gBBZmSuIySGqDLtXdZiYpwyzbJKXQD2jjT0oDY6ywbo= +github.com/aws/aws-sdk-go-v2/service/ssooidc v1.34.5/go.mod h1:XclEty74bsGBCr1s0VSaA11hQ4ZidK4viWK7rRfO88I= +github.com/aws/aws-sdk-go-v2/service/sts v1.38.4 h1:PR00NXRYgY4FWHqOGx3fC3lhVKjsp1GdloDv2ynMSd8= +github.com/aws/aws-sdk-go-v2/service/sts v1.38.4/go.mod h1:Z+Gd23v97pX9zK97+tX4ppAgqCt3Z2dIXB02CtBncK8= +github.com/aws/smithy-go v1.23.0 h1:8n6I3gXzWJB2DxBDnfxgBaSX6oe0d/t10qGz7OKqMCE= +github.com/aws/smithy-go v1.23.0/go.mod h1:t1ufH5HMublsJYulve2RKmHDC15xu1f26kHCp/HgceI= github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= diff --git a/internal/database/migrations/sql/mysql/00006_add_token_metadata.sql b/internal/database/migrations/sql/mysql/00006_add_token_metadata.sql new file mode 100644 index 00000000..13cdee6d --- /dev/null +++ b/internal/database/migrations/sql/mysql/00006_add_token_metadata.sql @@ -0,0 +1,5 @@ +-- +goose Up +ALTER TABLE tokens ADD COLUMN metadata TEXT; + +-- +goose Down +ALTER TABLE tokens DROP COLUMN metadata; \ No newline at end of file diff --git a/internal/database/migrations/sql/postgres/00007_add_token_metadata.sql b/internal/database/migrations/sql/postgres/00007_add_token_metadata.sql new file mode 100644 index 00000000..9d0775c8 --- /dev/null +++ b/internal/database/migrations/sql/postgres/00007_add_token_metadata.sql @@ -0,0 +1,5 @@ +-- +goose Up +ALTER TABLE tokens ADD COLUMN IF NOT EXISTS metadata TEXT; + +-- +goose Down +ALTER TABLE tokens DROP COLUMN IF EXISTS metadata; diff --git a/internal/database/models.go b/internal/database/models.go index a8fa8d8d..31003cea 100644 --- a/internal/database/models.go +++ b/internal/database/models.go @@ -20,6 +20,7 @@ type Token struct { ID string `json:"id"` Token string `json:"token"` ProjectID string `json:"project_id"` + Metadata *string `json:"metadata,omitempty"` ExpiresAt *time.Time `json:"expires_at,omitempty"` IsActive bool `json:"is_active"` DeactivatedAt *time.Time `json:"deactivated_at,omitempty"` diff --git a/internal/database/token.go b/internal/database/token.go index 352ff83a..f2fd7e66 100644 --- a/internal/database/token.go +++ b/internal/database/token.go @@ -3,6 +3,7 @@ package database import ( "context" "database/sql" + "encoding/json" "errors" "fmt" "strings" @@ -29,8 +30,8 @@ func (d *DB) CreateToken(ctx context.Context, token Token) error { } query := ` - INSERT INTO tokens (id, token, project_id, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + INSERT INTO tokens (id, token, project_id, metadata, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ` _, err := d.ExecContextRebound( @@ -39,6 +40,7 @@ func (d *DB) CreateToken(ctx context.Context, token Token) error { token.ID, token.Token, token.ProjectID, + token.Metadata, token.ExpiresAt, token.IsActive, nil, @@ -57,12 +59,13 @@ func (d *DB) CreateToken(ctx context.Context, token Token) error { // GetTokenByID retrieves a token by its UUID. func (d *DB) GetTokenByID(ctx context.Context, id string) (Token, error) { query := ` - SELECT id, token, project_id, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at, cache_hit_count + SELECT id, token, project_id, metadata, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at, cache_hit_count FROM tokens WHERE id = ? ` var token Token + var metadata sql.NullString var expiresAt, lastUsedAt, deactivatedAt sql.NullTime var maxRequests sql.NullInt32 @@ -70,6 +73,7 @@ func (d *DB) GetTokenByID(ctx context.Context, id string) (Token, error) { &token.ID, &token.Token, &token.ProjectID, + &metadata, &expiresAt, &token.IsActive, &deactivatedAt, @@ -89,6 +93,9 @@ func (d *DB) GetTokenByID(ctx context.Context, id string) (Token, error) { if expiresAt.Valid { token.ExpiresAt = &expiresAt.Time } + if metadata.Valid { + token.Metadata = &metadata.String + } if lastUsedAt.Valid { token.LastUsedAt = &lastUsedAt.Time } @@ -106,12 +113,13 @@ func (d *DB) GetTokenByID(ctx context.Context, id string) (Token, error) { // GetTokenByToken retrieves a token by its token string (for authentication). func (d *DB) GetTokenByToken(ctx context.Context, tokenString string) (Token, error) { query := ` - SELECT id, token, project_id, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at, cache_hit_count + SELECT id, token, project_id, metadata, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at, cache_hit_count FROM tokens WHERE token = ? ` var token Token + var metadata sql.NullString var expiresAt, lastUsedAt, deactivatedAt sql.NullTime var maxRequests sql.NullInt32 @@ -119,6 +127,7 @@ func (d *DB) GetTokenByToken(ctx context.Context, tokenString string) (Token, er &token.ID, &token.Token, &token.ProjectID, + &metadata, &expiresAt, &token.IsActive, &deactivatedAt, @@ -138,6 +147,9 @@ func (d *DB) GetTokenByToken(ctx context.Context, tokenString string) (Token, er if expiresAt.Valid { token.ExpiresAt = &expiresAt.Time } + if metadata.Valid { + token.Metadata = &metadata.String + } if lastUsedAt.Valid { token.LastUsedAt = &lastUsedAt.Time } @@ -162,12 +174,12 @@ func (d *DB) UpdateToken(ctx context.Context, token Token) error { queryByID := ` UPDATE tokens - SET project_id = ?, expires_at = ?, is_active = ?, request_count = ?, max_requests = ?, last_used_at = ? + SET project_id = ?, metadata = ?, expires_at = ?, is_active = ?, request_count = ?, max_requests = ?, last_used_at = ? WHERE id = ? ` queryByToken := ` UPDATE tokens - SET project_id = ?, expires_at = ?, is_active = ?, request_count = ?, max_requests = ?, last_used_at = ? + SET project_id = ?, metadata = ?, expires_at = ?, is_active = ?, request_count = ?, max_requests = ?, last_used_at = ? WHERE token = ? ` @@ -182,6 +194,7 @@ func (d *DB) UpdateToken(ctx context.Context, token Token) error { ctx, query, token.ProjectID, + token.Metadata, token.ExpiresAt, token.IsActive, token.RequestCount, @@ -232,7 +245,7 @@ func (d *DB) DeleteToken(ctx context.Context, tokenID string) error { // ListTokens retrieves all tokens from the database. func (d *DB) ListTokens(ctx context.Context) ([]Token, error) { query := ` - SELECT id, token, project_id, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at, cache_hit_count + SELECT id, token, project_id, metadata, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at, cache_hit_count FROM tokens ORDER BY created_at DESC ` @@ -243,7 +256,7 @@ func (d *DB) ListTokens(ctx context.Context) ([]Token, error) { // GetTokensByProjectID retrieves all tokens for a project. func (d *DB) GetTokensByProjectID(ctx context.Context, projectID string) ([]Token, error) { query := ` - SELECT id, token, project_id, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at, cache_hit_count + SELECT id, token, project_id, metadata, expires_at, is_active, deactivated_at, request_count, max_requests, created_at, last_used_at, cache_hit_count FROM tokens WHERE project_id = ? ORDER BY created_at DESC @@ -419,6 +432,7 @@ func (d *DB) queryTokens(ctx context.Context, query string, args ...interface{}) var tokens []Token for rows.Next() { var token Token + var metadata sql.NullString var expiresAt, lastUsedAt, deactivatedAt sql.NullTime var maxRequests sql.NullInt32 @@ -426,6 +440,7 @@ func (d *DB) queryTokens(ctx context.Context, query string, args ...interface{}) &token.ID, &token.Token, &token.ProjectID, + &metadata, &expiresAt, &token.IsActive, &deactivatedAt, @@ -441,6 +456,9 @@ func (d *DB) queryTokens(ctx context.Context, query string, args ...interface{}) if expiresAt.Valid { token.ExpiresAt = &expiresAt.Time } + if metadata.Valid { + token.Metadata = &metadata.String + } if lastUsedAt.Valid { token.LastUsedAt = &lastUsedAt.Time } @@ -533,10 +551,12 @@ func (a *DBTokenStoreAdapter) GetTokensByProjectID(ctx context.Context, projectI // ImportTokenData and ExportTokenData helpers func ImportTokenData(td token.TokenData) Token { + metadata := marshalTokenMetadata(td.Metadata) return Token{ ID: td.ID, Token: td.Token, ProjectID: td.ProjectID, + Metadata: metadata, ExpiresAt: td.ExpiresAt, IsActive: td.IsActive, DeactivatedAt: td.DeactivatedAt, @@ -553,6 +573,7 @@ func ExportTokenData(t Token) token.TokenData { ID: t.ID, Token: t.Token, ProjectID: t.ProjectID, + Metadata: unmarshalTokenMetadata(t.Metadata), ExpiresAt: t.ExpiresAt, IsActive: t.IsActive, DeactivatedAt: t.DeactivatedAt, @@ -564,6 +585,37 @@ func ExportTokenData(t Token) token.TokenData { } } +func marshalTokenMetadata(metadata map[string]string) *string { + if len(metadata) == 0 { + return nil + } + + encoded, err := json.Marshal(metadata) + if err != nil { + return nil + } + + result := string(encoded) + return &result +} + +func unmarshalTokenMetadata(metadataJSON *string) map[string]string { + if metadataJSON == nil || *metadataJSON == "" { + return nil + } + + var metadata map[string]string + if err := json.Unmarshal([]byte(*metadataJSON), &metadata); err != nil { + return nil + } + + if len(metadata) == 0 { + return nil + } + + return metadata +} + // --- RevocationStore interface implementation --- // RevokeToken disables a token by setting is_active to false and deactivated_at to current time diff --git a/internal/database/token_test.go b/internal/database/token_test.go index e6df7cfc..4fb53282 100644 --- a/internal/database/token_test.go +++ b/internal/database/token_test.go @@ -41,6 +41,7 @@ func TestTokenCRUD(t *testing.T) { token1 := Token{ Token: "test-token-1", ProjectID: project.ID, + Metadata: stringPtr(`{"feature":"sofabuddy","user_id":"42"}`), ExpiresAt: &expiresAt, IsActive: true, RequestCount: 0, @@ -83,6 +84,9 @@ func TestTokenCRUD(t *testing.T) { if retrievedToken1.ProjectID != token1.ProjectID { t.Fatalf("Expected project ID %s, got %s", token1.ProjectID, retrievedToken1.ProjectID) } + if retrievedToken1.Metadata == nil || *retrievedToken1.Metadata != *token1.Metadata { + t.Fatalf("Expected metadata %v, got %v", token1.Metadata, retrievedToken1.Metadata) + } if retrievedToken1.ExpiresAt == nil { t.Fatalf("Expected expiration time, got nil") } @@ -239,6 +243,10 @@ func TestTokenCRUD(t *testing.T) { } } +func stringPtr(value string) *string { + return &value +} + func TestIncrementTokenUsageBatch(t *testing.T) { db, cleanup := testDB(t) defer cleanup() diff --git a/internal/dispatcher/plugin.go b/internal/dispatcher/plugin.go index b38163e3..28206229 100644 --- a/internal/dispatcher/plugin.go +++ b/internal/dispatcher/plugin.go @@ -33,8 +33,11 @@ type EventPayload struct { // TokensUsage represents token usage statistics type TokensUsage struct { - Completion int `json:"completion"` - Prompt int `json:"prompt"` + Input int `json:"input_tokens"` + InputDetails map[string]any `json:"input_tokens_details,omitempty"` + Output int `json:"output_tokens"` + OutputDetails map[string]any `json:"output_tokens_details,omitempty"` + Total int `json:"total_tokens"` } // BackendPlugin defines the interface for dispatcher backends diff --git a/internal/dispatcher/plugins/cloudwatch.go b/internal/dispatcher/plugins/cloudwatch.go new file mode 100644 index 00000000..95b90759 --- /dev/null +++ b/internal/dispatcher/plugins/cloudwatch.go @@ -0,0 +1,319 @@ +package plugins + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "os" + "sort" + "strconv" + "sync" + "time" + + awsconfig "github.com/aws/aws-sdk-go-v2/config" + cloudwatchlogs "github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs" + cloudwatchtypes "github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs/types" + "github.com/sofatutor/llm-proxy/internal/dispatcher" +) + +var errCloudWatchRetriesExhausted = errors.New("cloudwatch put log events retries exhausted") + +type cloudWatchLogsClient interface { + CreateLogStream(context.Context, *cloudwatchlogs.CreateLogStreamInput, ...func(*cloudwatchlogs.Options)) (*cloudwatchlogs.CreateLogStreamOutput, error) + PutLogEvents(context.Context, *cloudwatchlogs.PutLogEventsInput, ...func(*cloudwatchlogs.Options)) (*cloudwatchlogs.PutLogEventsOutput, error) +} + +// CloudWatchPlugin writes sanitized observability records to CloudWatch Logs. +type CloudWatchPlugin struct { + client cloudWatchLogsClient + logGroup string + logStream string + region string + + mu sync.Mutex + sequenceToken *string +} + +// NewCloudWatchPlugin creates a CloudWatch plugin. +func NewCloudWatchPlugin() *CloudWatchPlugin { + return &CloudWatchPlugin{} +} + +// Init initializes the plugin with CloudWatch configuration. +func (p *CloudWatchPlugin) Init(cfg map[string]string) error { + p.logGroup = firstNonEmpty(cfg["log-group"], os.Getenv("DISPATCHER_CLOUDWATCH_LOG_GROUP")) + p.logStream = firstNonEmpty(cfg["log-stream"], os.Getenv("DISPATCHER_CLOUDWATCH_LOG_STREAM")) + p.region = firstNonEmpty(cfg["region"], os.Getenv("DISPATCHER_CLOUDWATCH_REGION"), os.Getenv("AWS_REGION"), os.Getenv("AWS_DEFAULT_REGION")) + + if p.logGroup == "" { + return fmt.Errorf("cloudwatch plugin requires 'log-group' configuration") + } + if p.logStream == "" { + hostname, _ := os.Hostname() + if hostname == "" { + hostname = "unknown" + } + p.logStream = fmt.Sprintf("llm-proxy-dispatcher-%s-%d", hostname, os.Getpid()) + } + + if p.client != nil { + return nil + } + + loadOptions := []func(*awsconfig.LoadOptions) error{} + if p.region != "" { + loadOptions = append(loadOptions, awsconfig.WithRegion(p.region)) + } + + awsCfg, err := awsconfig.LoadDefaultConfig(context.Background(), loadOptions...) + if err != nil { + return fmt.Errorf("load aws config: %w", err) + } + + p.client = cloudwatchlogs.NewFromConfig(awsCfg) + return nil +} + +// SendEvents writes a batch of sanitized events to CloudWatch Logs. +func (p *CloudWatchPlugin) SendEvents(ctx context.Context, events []dispatcher.EventPayload) error { + if len(events) == 0 { + return nil + } + if p.client == nil { + return fmt.Errorf("cloudwatch plugin not initialized") + } + + sortedEvents := append([]dispatcher.EventPayload(nil), events...) + sort.SliceStable(sortedEvents, func(i, j int) bool { + return sortedEvents[i].Timestamp.Before(sortedEvents[j].Timestamp) + }) + + logEvents := make([]cloudwatchtypes.InputLogEvent, 0, len(sortedEvents)) + for _, event := range sortedEvents { + message, err := cloudWatchLogMessage(event) + if err != nil { + return err + } + logEvents = append(logEvents, cloudwatchtypes.InputLogEvent{ + Message: &message, + Timestamp: awsInt64(event.Timestamp.UnixMilli()), + }) + } + + for attempt := 0; attempt < 3; attempt++ { + output, err := p.client.PutLogEvents(ctx, p.newPutLogEventsInput(logEvents)) + if err == nil { + p.setSequenceToken(output.NextSequenceToken) + return nil + } + + var notFound *cloudwatchtypes.ResourceNotFoundException + if errorAs(err, ¬Found) { + if createErr := p.ensureLogStream(ctx); createErr != nil { + return createErr + } + p.setSequenceToken(nil) + continue + } + + var invalidSequence *cloudwatchtypes.InvalidSequenceTokenException + if errorAs(err, &invalidSequence) { + p.setSequenceToken(invalidSequence.ExpectedSequenceToken) + continue + } + + var alreadyAccepted *cloudwatchtypes.DataAlreadyAcceptedException + if errorAs(err, &alreadyAccepted) { + p.setSequenceToken(alreadyAccepted.ExpectedSequenceToken) + return nil + } + + return err + } + + return fmt.Errorf("put cloudwatch log events: %w", errCloudWatchRetriesExhausted) +} + +// Close cleans up plugin resources. +func (p *CloudWatchPlugin) Close() error { + return nil +} + +func (p *CloudWatchPlugin) ensureLogStream(ctx context.Context) error { + _, createErr := p.client.CreateLogStream(ctx, &cloudwatchlogs.CreateLogStreamInput{ + LogGroupName: &p.logGroup, + LogStreamName: &p.logStream, + }) + if createErr == nil { + return nil + } + + var alreadyExists *cloudwatchtypes.ResourceAlreadyExistsException + if errorAs(createErr, &alreadyExists) { + return nil + } + + return createErr +} + +func (p *CloudWatchPlugin) newPutLogEventsInput(logEvents []cloudwatchtypes.InputLogEvent) *cloudwatchlogs.PutLogEventsInput { + input := &cloudwatchlogs.PutLogEventsInput{ + LogEvents: logEvents, + LogGroupName: &p.logGroup, + LogStreamName: &p.logStream, + } + + if sequenceToken := p.getSequenceToken(); sequenceToken != nil { + input.SequenceToken = sequenceToken + } + + return input +} + +func (p *CloudWatchPlugin) getSequenceToken() *string { + p.mu.Lock() + defer p.mu.Unlock() + + if p.sequenceToken == nil { + return nil + } + + value := *p.sequenceToken + return &value +} + +func (p *CloudWatchPlugin) setSequenceToken(sequenceToken *string) { + p.mu.Lock() + defer p.mu.Unlock() + + if sequenceToken == nil || *sequenceToken == "" { + p.sequenceToken = nil + return + } + + value := *sequenceToken + p.sequenceToken = &value +} + +func cloudWatchLogMessage(event dispatcher.EventPayload) (string, error) { + message := map[string]any{ + "type": "llm_proxy_usage", + "timestamp": event.Timestamp.UTC().Format(time.RFC3339Nano), + "run_id": event.RunID, + "event": event.Event, + "request_id": metadataString(event.Metadata, "request_id"), + "path": metadataString(event.Metadata, "path"), + "model": metadataString(event.Metadata, "model"), + "project_id": metadataString(event.Metadata, "project_id"), + "token_id": metadataString(event.Metadata, "token_id"), + "status": metadataInt(event.Metadata, "status"), + "duration_ms": metadataInt64(event.Metadata, "duration_ms"), + } + + if event.UserID != nil || len(event.UserProps) > 0 { + user := make(map[string]any, len(event.UserProps)+1) + for key, value := range event.UserProps { + user[key] = value + } + if event.UserID != nil && *event.UserID != "" { + user["id"] = *event.UserID + } + if len(user) > 0 { + message["user"] = user + } + } + if event.TokensUsage != nil { + usage := map[string]any{ + "input_tokens": event.TokensUsage.Input, + "output_tokens": event.TokensUsage.Output, + "total_tokens": event.TokensUsage.Total, + } + if len(event.TokensUsage.InputDetails) > 0 { + usage["input_tokens_details"] = event.TokensUsage.InputDetails + } + if len(event.TokensUsage.OutputDetails) > 0 { + usage["output_tokens_details"] = event.TokensUsage.OutputDetails + } + message["usage"] = usage + } + if tokenMetadata, ok := event.Metadata["token_metadata"].(map[string]string); ok && len(tokenMetadata) > 0 { + message["token_metadata"] = tokenMetadata + } + + encoded, err := json.Marshal(message) + if err != nil { + return "", fmt.Errorf("marshal cloudwatch event: %w", err) + } + + return string(encoded), nil +} + +func metadataString(metadata map[string]any, key string) string { + if metadata == nil { + return "" + } + if value, ok := metadata[key].(string); ok { + return value + } + return "" +} + +func metadataInt(metadata map[string]any, key string) int { + if metadata == nil { + return 0 + } + switch value := metadata[key].(type) { + case int: + return value + case int8: + return int(value) + case int16: + return int(value) + case int32: + return int(value) + case int64: + return int(value) + case uint: + return int(value) + case uint8: + return int(value) + case uint16: + return int(value) + case uint32: + return int(value) + case uint64: + return int(value) + case float64: + return int(value) + case float32: + return int(value) + case string: + parsed, _ := strconv.Atoi(value) + return parsed + default: + return 0 + } +} + +func metadataInt64(metadata map[string]any, key string) int64 { + return int64(metadataInt(metadata, key)) +} + +func awsInt64(value int64) *int64 { + return &value +} + +func errorAs(err error, target interface{}) bool { + return errors.As(err, target) +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + + return "" +} diff --git a/internal/dispatcher/plugins/cloudwatch_test.go b/internal/dispatcher/plugins/cloudwatch_test.go new file mode 100644 index 00000000..34bf607a --- /dev/null +++ b/internal/dispatcher/plugins/cloudwatch_test.go @@ -0,0 +1,377 @@ +package plugins + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + "time" + + cloudwatchlogs "github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs" + cloudwatchtypes "github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs/types" + "github.com/sofatutor/llm-proxy/internal/dispatcher" +) + +type fakeCloudWatchLogsClient struct { + putInputs []*cloudwatchlogs.PutLogEventsInput + createInputs []*cloudwatchlogs.CreateLogStreamInput + putResponses []fakePutLogEventsResponse + createErr error + createdLogStream bool +} + +type fakePutLogEventsResponse struct { + output *cloudwatchlogs.PutLogEventsOutput + err error +} + +func (f *fakeCloudWatchLogsClient) CreateLogStream(_ context.Context, input *cloudwatchlogs.CreateLogStreamInput, _ ...func(*cloudwatchlogs.Options)) (*cloudwatchlogs.CreateLogStreamOutput, error) { + f.createInputs = append(f.createInputs, input) + f.createdLogStream = true + if f.createErr != nil { + return nil, f.createErr + } + return &cloudwatchlogs.CreateLogStreamOutput{}, nil +} + +func (f *fakeCloudWatchLogsClient) PutLogEvents(_ context.Context, input *cloudwatchlogs.PutLogEventsInput, _ ...func(*cloudwatchlogs.Options)) (*cloudwatchlogs.PutLogEventsOutput, error) { + f.putInputs = append(f.putInputs, input) + if len(f.putResponses) > 0 { + response := f.putResponses[0] + f.putResponses = f.putResponses[1:] + if response.err != nil { + return nil, response.err + } + if response.output != nil { + return response.output, nil + } + } + return &cloudwatchlogs.PutLogEventsOutput{}, nil +} + +func TestCloudWatchLogMessage_SanitizesPayload(t *testing.T) { + message, err := cloudWatchLogMessage(dispatcher.EventPayload{ + RunID: "run-1", + Event: "start", + Timestamp: time.Unix(1700000000, 0), + UserID: stringPtrCW("42"), + UserProps: map[string]any{"plan": "pro"}, + TokensUsage: &dispatcher.TokensUsage{ + Input: 3, + InputDetails: map[string]any{"cached_tokens": 1}, + Output: 2, + OutputDetails: map[string]any{"reasoning_tokens": 0}, + Total: 5, + }, + Input: json.RawMessage(`{"sensitive":"input"}`), + Output: json.RawMessage(`{"sensitive":"output"}`), + Metadata: map[string]any{ + "request_id": "req-1", + "path": "/v1/responses", + "model": "gpt-4.1-mini", + "project_id": "project-1", + "token_id": "token-uuid-1", + "status": 200, + "duration_ms": int64(123), + "token_metadata": map[string]string{"feature": "sofabuddy", "user_id": "42"}, + }, + }) + if err != nil { + t.Fatalf("cloudWatchLogMessage() error = %v", err) + } + var decoded map[string]any + if err := json.Unmarshal([]byte(message), &decoded); err != nil { + t.Fatalf("json.Unmarshal() error = %v", err) + } + if _, exists := decoded["input"]; exists { + t.Fatalf("expected sanitized payload to omit input") + } + if _, exists := decoded["output"]; exists { + t.Fatalf("expected sanitized payload to omit output") + } + user, ok := decoded["user"].(map[string]any) + if !ok { + t.Fatalf("expected nested user object, got %T", decoded["user"]) + } + if user["id"] != "42" { + t.Fatalf("expected user.id 42, got %v", user["id"]) + } + if user["plan"] != "pro" { + t.Fatalf("expected user.plan pro, got %v", user["plan"]) + } + if decoded["duration_ms"].(float64) != 123 { + t.Fatalf("expected duration_ms 123, got %v", decoded["duration_ms"]) + } + usage, ok := decoded["usage"].(map[string]any) + if !ok { + t.Fatalf("expected canonical usage object, got %T", decoded["usage"]) + } + if usage["input_tokens"].(float64) != 3 { + t.Fatalf("expected input_tokens 3, got %v", usage["input_tokens"]) + } + if usage["output_tokens"].(float64) != 2 { + t.Fatalf("expected output_tokens 2, got %v", usage["output_tokens"]) + } + if usage["total_tokens"].(float64) != 5 { + t.Fatalf("expected usage.total_tokens 5, got %v", usage["total_tokens"]) + } + if _, exists := decoded["input_tokens"]; exists { + t.Fatalf("expected top-level input_tokens to be omitted") + } + if _, exists := decoded["output_tokens"]; exists { + t.Fatalf("expected top-level output_tokens to be omitted") + } + if _, exists := decoded["total_tokens"]; exists { + t.Fatalf("expected top-level total_tokens to be omitted") + } +} + +func TestCloudWatchPlugin_SendEvents_CreatesStreamOnMissing(t *testing.T) { + client := &fakeCloudWatchLogsClient{ + putResponses: []fakePutLogEventsResponse{ + {err: &cloudwatchtypes.ResourceNotFoundException{}}, + {output: &cloudwatchlogs.PutLogEventsOutput{}}, + }, + } + plugin := NewCloudWatchPlugin() + plugin.client = client + if err := plugin.Init(map[string]string{"log-group": "/llm-proxy", "log-stream": "dispatcher-1", "region": "eu-central-1"}); err != nil { + t.Fatalf("Init() error = %v", err) + } + + err := plugin.SendEvents(context.Background(), []dispatcher.EventPayload{{ + RunID: "run-1", + Event: "start", + Timestamp: time.Now(), + Metadata: map[string]any{ + "request_id": "req-1", + "path": "/v1/responses", + "status": 200, + }, + }}) + if err != nil { + t.Fatalf("SendEvents() error = %v", err) + } + if len(client.createInputs) != 1 { + t.Fatalf("expected CreateLogStream to be called once, got %d", len(client.createInputs)) + } + if len(client.putInputs) != 2 { + t.Fatalf("expected PutLogEvents to be called twice, got %d", len(client.putInputs)) + } +} + +func TestCloudWatchPlugin_SendEvents_SortsByTimestamp(t *testing.T) { + client := &fakeCloudWatchLogsClient{} + plugin := NewCloudWatchPlugin() + plugin.client = client + if err := plugin.Init(map[string]string{"log-group": "/llm-proxy", "log-stream": "dispatcher-1", "region": "eu-central-1"}); err != nil { + t.Fatalf("Init() error = %v", err) + } + + later := time.Unix(1700000010, 0) + earlier := later.Add(-time.Minute) + err := plugin.SendEvents(context.Background(), []dispatcher.EventPayload{ + {RunID: "run-late", Event: "finish", Timestamp: later, Metadata: map[string]any{"status": 200}}, + {RunID: "run-early", Event: "start", Timestamp: earlier, Metadata: map[string]any{"status": 200}}, + }) + if err != nil { + t.Fatalf("SendEvents() error = %v", err) + } + if len(client.putInputs) != 1 { + t.Fatalf("expected one PutLogEvents call, got %d", len(client.putInputs)) + } + if len(client.putInputs[0].LogEvents) != 2 { + t.Fatalf("expected two log events, got %d", len(client.putInputs[0].LogEvents)) + } + if *client.putInputs[0].LogEvents[0].Timestamp > *client.putInputs[0].LogEvents[1].Timestamp { + t.Fatalf("expected log events to be sorted chronologically") + } +} + +func TestCloudWatchPlugin_SendEvents_RetriesWithExpectedSequenceToken(t *testing.T) { + expectedToken := "token-2" + nextToken := "token-3" + client := &fakeCloudWatchLogsClient{ + putResponses: []fakePutLogEventsResponse{ + {err: &cloudwatchtypes.InvalidSequenceTokenException{ExpectedSequenceToken: &expectedToken}}, + {output: &cloudwatchlogs.PutLogEventsOutput{NextSequenceToken: &nextToken}}, + }, + } + plugin := NewCloudWatchPlugin() + plugin.client = client + plugin.setSequenceToken(stringPtrCW("token-1")) + if err := plugin.Init(map[string]string{"log-group": "/llm-proxy", "log-stream": "dispatcher-1", "region": "eu-central-1"}); err != nil { + t.Fatalf("Init() error = %v", err) + } + + err := plugin.SendEvents(context.Background(), []dispatcher.EventPayload{{ + RunID: "run-1", + Event: "start", + Timestamp: time.Now(), + Metadata: map[string]any{"status": 200}, + }}) + if err != nil { + t.Fatalf("SendEvents() error = %v", err) + } + if len(client.putInputs) != 2 { + t.Fatalf("expected PutLogEvents to be retried once, got %d calls", len(client.putInputs)) + } + if client.putInputs[0].SequenceToken == nil || *client.putInputs[0].SequenceToken != "token-1" { + t.Fatalf("expected first request to use existing sequence token, got %v", client.putInputs[0].SequenceToken) + } + if client.putInputs[1].SequenceToken == nil || *client.putInputs[1].SequenceToken != expectedToken { + t.Fatalf("expected retry to use updated sequence token, got %v", client.putInputs[1].SequenceToken) + } + if token := plugin.getSequenceToken(); token == nil || *token != nextToken { + t.Fatalf("expected plugin to persist next sequence token %q, got %v", nextToken, token) + } +} + +func TestCloudWatchPlugin_SendEvents_DataAlreadyAcceptedIsTreatedAsSuccess(t *testing.T) { + expectedToken := "token-4" + client := &fakeCloudWatchLogsClient{ + putResponses: []fakePutLogEventsResponse{{ + err: &cloudwatchtypes.DataAlreadyAcceptedException{ExpectedSequenceToken: &expectedToken}, + }}, + } + plugin := NewCloudWatchPlugin() + plugin.client = client + if err := plugin.Init(map[string]string{"log-group": "/llm-proxy", "log-stream": "dispatcher-1", "region": "eu-central-1"}); err != nil { + t.Fatalf("Init() error = %v", err) + } + + err := plugin.SendEvents(context.Background(), []dispatcher.EventPayload{{ + RunID: "run-1", + Event: "start", + Timestamp: time.Now(), + Metadata: map[string]any{"status": 200}, + }}) + if err != nil { + t.Fatalf("SendEvents() error = %v", err) + } + if token := plugin.getSequenceToken(); token == nil || *token != expectedToken { + t.Fatalf("expected plugin to persist expected sequence token %q, got %v", expectedToken, token) + } +} + +func TestCloudWatchPlugin_SendEvents_ExhaustedSequenceRetries(t *testing.T) { + expectedToken := "still-wrong" + client := &fakeCloudWatchLogsClient{ + putResponses: []fakePutLogEventsResponse{ + {err: &cloudwatchtypes.InvalidSequenceTokenException{ExpectedSequenceToken: &expectedToken}}, + {err: &cloudwatchtypes.InvalidSequenceTokenException{ExpectedSequenceToken: &expectedToken}}, + {err: &cloudwatchtypes.InvalidSequenceTokenException{ExpectedSequenceToken: &expectedToken}}, + }, + } + plugin := NewCloudWatchPlugin() + plugin.client = client + if err := plugin.Init(map[string]string{"log-group": "/llm-proxy", "log-stream": "dispatcher-1", "region": "eu-central-1"}); err != nil { + t.Fatalf("Init() error = %v", err) + } + + err := plugin.SendEvents(context.Background(), []dispatcher.EventPayload{{ + RunID: "run-1", + Event: "start", + Timestamp: time.Now(), + Metadata: map[string]any{"status": 200}, + }}) + if err == nil || !errors.Is(err, errCloudWatchRetriesExhausted) { + t.Fatalf("SendEvents() error = %v, want errCloudWatchRetriesExhausted", err) + } +} + +func TestCloudWatchPlugin_Init_RequiresLogGroup(t *testing.T) { + plugin := NewCloudWatchPlugin() + plugin.client = &fakeCloudWatchLogsClient{} + + err := plugin.Init(map[string]string{"log-stream": "dispatcher-1"}) + if err == nil || !strings.Contains(err.Error(), "log-group") { + t.Fatalf("Init() error = %v, want missing log-group error", err) + } +} + +func TestCloudWatchPlugin_Init_DefaultsLogStreamWhenMissing(t *testing.T) { + plugin := NewCloudWatchPlugin() + plugin.client = &fakeCloudWatchLogsClient{} + + err := plugin.Init(map[string]string{"log-group": "/llm-proxy", "region": "eu-central-1"}) + if err != nil { + t.Fatalf("Init() error = %v", err) + } + if plugin.logStream == "" || !strings.Contains(plugin.logStream, "llm-proxy-dispatcher-") { + t.Fatalf("expected generated log stream, got %q", plugin.logStream) + } +} + +func TestCloudWatchPlugin_EnsureLogStream_AlreadyExists(t *testing.T) { + plugin := NewCloudWatchPlugin() + plugin.client = &fakeCloudWatchLogsClient{createErr: &cloudwatchtypes.ResourceAlreadyExistsException{}} + plugin.logGroup = "/llm-proxy" + plugin.logStream = "dispatcher-1" + + if err := plugin.ensureLogStream(context.Background()); err != nil { + t.Fatalf("ensureLogStream() error = %v", err) + } +} + +func TestCloudWatchPlugin_EnsureLogStream_PropagatesCreateError(t *testing.T) { + plugin := NewCloudWatchPlugin() + plugin.client = &fakeCloudWatchLogsClient{createErr: errors.New("boom")} + plugin.logGroup = "/llm-proxy" + plugin.logStream = "dispatcher-1" + + err := plugin.ensureLogStream(context.Background()) + if err == nil || err.Error() != "boom" { + t.Fatalf("ensureLogStream() error = %v, want boom", err) + } +} + +func TestCloudWatchPlugin_Close_NoOp(t *testing.T) { + plugin := NewCloudWatchPlugin() + if err := plugin.Close(); err != nil { + t.Fatalf("Close() error = %v", err) + } +} + +func TestCloudWatchMetadataHelpers(t *testing.T) { + metadata := map[string]any{ + "as-string": "value", + "as-int": 12, + "as-float": float64(34), + "as-text": "56", + "invalid": true, + } + + if got := metadataString(metadata, "as-string"); got != "value" { + t.Fatalf("metadataString() = %q, want value", got) + } + if got := metadataString(metadata, "missing"); got != "" { + t.Fatalf("metadataString() missing = %q, want empty", got) + } + if got := metadataInt(metadata, "as-int"); got != 12 { + t.Fatalf("metadataInt() int = %d, want 12", got) + } + if got := metadataInt(metadata, "as-float"); got != 34 { + t.Fatalf("metadataInt() float = %d, want 34", got) + } + if got := metadataInt(metadata, "as-text"); got != 56 { + t.Fatalf("metadataInt() string = %d, want 56", got) + } + if got := metadataInt(metadata, "invalid"); got != 0 { + t.Fatalf("metadataInt() invalid = %d, want 0", got) + } + if got := metadataInt(nil, "as-int"); got != 0 { + t.Fatalf("metadataInt() nil metadata = %d, want 0", got) + } + if got := firstNonEmpty("", "fallback", "last"); got != "fallback" { + t.Fatalf("firstNonEmpty() = %q, want fallback", got) + } + if got := firstNonEmpty("", ""); got != "" { + t.Fatalf("firstNonEmpty() empty = %q, want empty", got) + } +} + +func stringPtrCW(value string) *string { + return &value +} diff --git a/internal/dispatcher/plugins/helicone.go b/internal/dispatcher/plugins/helicone.go index 5c4294e9..6803f988 100644 --- a/internal/dispatcher/plugins/helicone.go +++ b/internal/dispatcher/plugins/helicone.go @@ -169,11 +169,18 @@ func heliconePayloadFromEvent(event dispatcher.EventPayload) (map[string]interfa // Inject usage if we computed it so Helicone can compute cost if event.TokensUsage != nil { if m, ok := providerResponse["json"].(map[string]interface{}); ok { - m["usage"] = map[string]int{ - "prompt_tokens": event.TokensUsage.Prompt, - "completion_tokens": event.TokensUsage.Completion, - "total_tokens": event.TokensUsage.Prompt + event.TokensUsage.Completion, + usage := map[string]any{ + "prompt_tokens": event.TokensUsage.Input, + "completion_tokens": event.TokensUsage.Output, + "total_tokens": event.TokensUsage.Total, } + if len(event.TokensUsage.InputDetails) > 0 { + usage["prompt_tokens_details"] = event.TokensUsage.InputDetails + } + if len(event.TokensUsage.OutputDetails) > 0 { + usage["completion_tokens_details"] = event.TokensUsage.OutputDetails + } + m["usage"] = usage } } // If OutputBase64 is set, include as base64 diff --git a/internal/dispatcher/plugins/registry.go b/internal/dispatcher/plugins/registry.go index 46660dd6..29179f1d 100644 --- a/internal/dispatcher/plugins/registry.go +++ b/internal/dispatcher/plugins/registry.go @@ -25,6 +25,10 @@ func init() { Registry["helicone"] = func() dispatcher.BackendPlugin { return NewHeliconePlugin() } + + Registry["cloudwatch"] = func() dispatcher.BackendPlugin { + return NewCloudWatchPlugin() + } } // NewPlugin creates a new plugin instance by name diff --git a/internal/dispatcher/transformer.go b/internal/dispatcher/transformer.go index 8b6f7e1e..d8bb8b9f 100644 --- a/internal/dispatcher/transformer.go +++ b/internal/dispatcher/transformer.go @@ -6,6 +6,7 @@ import ( "encoding/json" "io" "log" + "mime/multipart" "strconv" "strings" "time" @@ -177,7 +178,7 @@ func (t *DefaultEventTransformer) Transform(evt eventbus.Event) (*EventPayload, // Basic transformation payload := &EventPayload{ Type: "llm", - Event: "start", // For now, all events are considered "start" events + Event: eventNameForEvent(evt), RunID: runID, Timestamp: time.Now(), LogID: evt.LogID, @@ -185,10 +186,25 @@ func (t *DefaultEventTransformer) Transform(evt eventbus.Event) (*EventPayload, "method": evt.Method, "path": evt.Path, "status": evt.Status, - "duration_ms": evt.Duration.Milliseconds(), + "duration_ms": durationMilliseconds(evt.Duration), "request_id": evt.RequestID, }, } + if model := modelFromRequestBody(evt.RequestBody); model != "" { + payload.Metadata["model"] = model + } + if evt.ProjectID != "" { + payload.Metadata["project_id"] = evt.ProjectID + } + if evt.TokenID != "" { + payload.Metadata["token_id"] = evt.TokenID + } + if len(evt.TokenMetadata) > 0 { + payload.Metadata["token_metadata"] = evt.TokenMetadata + if userID := firstNonEmpty(evt.TokenMetadata["user_id"], evt.TokenMetadata["app_user_id"]); userID != "" { + payload.UserID = &userID + } + } // Add request body as input (JSON or base64) if len(evt.RequestBody) > 0 { @@ -199,30 +215,29 @@ func (t *DefaultEventTransformer) Transform(evt eventbus.Event) (*EventPayload, } } + decodedResponseBody := decodeResponseBodyForProcessing(evt.ResponseBody, evt.ResponseHeaders) + // --- OpenAI-specific output transformation --- - isOpenAI := strings.HasPrefix(evt.Path, "/v1/completions") || - strings.HasPrefix(evt.Path, "/v1/chat/completions") || - strings.HasPrefix(evt.Path, "/v1/threads/") + isOpenAI := usesOpenAIResponseTransformer(evt.Path) - if isOpenAI && len(evt.ResponseBody) > 0 { + if isOpenAI && len(decodedResponseBody) > 0 { // Only use OpenAI transformer if response is valid JSON - if js := json.Valid(evt.ResponseBody); js { + if js := json.Valid(decodedResponseBody); js { openaiTransformer := &eventtransformer.OpenAITransformer{} parsed, err := openaiTransformer.TransformEvent(map[string]any{ - "response_body": string(evt.ResponseBody), - "path": evt.Path, + "Method": evt.Method, + "Path": evt.Path, + "RequestBody": string(evt.RequestBody), + "ResponseBody": string(decodedResponseBody), + "ResponseHeaders": headerToAnyMap(evt.ResponseHeaders), }) if err == nil && parsed != nil { if js, err := json.Marshal(parsed); err == nil { payload.Output = js - // Optionally extract token usage if present - if usage, ok := parsed["usage"].(map[string]any); ok { - payload.TokensUsage = &TokensUsage{ - Prompt: int(usage["prompt_tokens"].(float64)), - Completion: int(usage["completion_tokens"].(float64)), - } + if model := firstNonEmpty(modelFromParsedOpenAIEvent(parsed), metadataStringValue(payload.Metadata, "model")); model != "" { + payload.Metadata["model"] = model } - return payload, nil + payload.TokensUsage = firstTokensUsage(parsed["token_usage"], parsed["usage"]) } } // If parsing fails, fall through to generic logic @@ -238,6 +253,10 @@ func (t *DefaultEventTransformer) Transform(evt eventbus.Event) (*EventPayload, } } + if payload.TokensUsage == nil { + payload.TokensUsage = fallbackTokensUsage(evt.RequestBody, decodedResponseBody, metadataStringValue(payload.Metadata, "model")) + } + // Add response headers to metadata only if Verbose is true if t.Verbose && evt.ResponseHeaders != nil { headers := make(map[string]any) @@ -253,3 +272,524 @@ func (t *DefaultEventTransformer) Transform(evt eventbus.Event) (*EventPayload, return payload, nil } + +func decodeResponseBodyForProcessing(responseBody []byte, headers map[string][]string) []byte { + if len(responseBody) == 0 { + return nil + } + + js, _ := safeRawMessageOrBase64(responseBody, headers) + if js != nil { + if len(js) >= 2 && js[0] == '"' && js[len(js)-1] == '"' { + var decoded string + if err := json.Unmarshal(js, &decoded); err == nil { + return []byte(decoded) + } + } + + return []byte(js) + } + + return responseBody +} + +func headerToAnyMap(header map[string][]string) map[string]any { + if len(header) == 0 { + return nil + } + + converted := make(map[string]any, len(header)) + for key, values := range header { + if len(values) == 1 { + converted[key] = values[0] + continue + } + + items := make([]any, len(values)) + for index, value := range values { + items[index] = value + } + converted[key] = items + } + + return converted +} + +func tokensUsageFromUsageMap(usage map[string]any) *TokensUsage { + input, okInput := firstUsageValue(usage, "input_tokens", "prompt_tokens") + output, okOutput := firstUsageValue(usage, "output_tokens", "completion_tokens") + total, okTotal := firstUsageValue(usage, "total_tokens") + if !okInput && !okOutput && !okTotal { + return nil + } + if !okTotal { + total = input + output + } + + return &TokensUsage{ + Input: input, + InputDetails: firstUsageDetailsMap(usage, "input_tokens_details", "prompt_tokens_details"), + Output: output, + OutputDetails: firstUsageDetailsMap(usage, "output_tokens_details", "completion_tokens_details"), + Total: total, + } +} + +func firstTokensUsage(values ...any) *TokensUsage { + for _, value := range values { + if usage := tokensUsageFromValue(value); usage != nil { + return usage + } + } + + return nil +} + +func tokensUsageFromValue(value any) *TokensUsage { + switch typed := value.(type) { + case map[string]any: + return tokensUsageFromUsageMap(typed) + case map[string]int: + return &TokensUsage{ + Input: firstIntMapValue(typed, "input_tokens", "prompt_tokens"), + Output: firstIntMapValue(typed, "output_tokens", "completion_tokens"), + Total: firstIntMapValue(typed, "total_tokens"), + } + case map[string]float64: + return &TokensUsage{ + Input: firstFloatMapValue(typed, "input_tokens", "prompt_tokens"), + Output: firstFloatMapValue(typed, "output_tokens", "completion_tokens"), + Total: firstFloatMapValue(typed, "total_tokens"), + } + default: + return nil + } +} + +func firstUsageDetailsMap(usage map[string]any, keys ...string) map[string]any { + for _, key := range keys { + value, ok := usage[key] + if !ok { + continue + } + if typed, ok := value.(map[string]any); ok && len(typed) > 0 { + return cloneStringAnyMap(typed) + } + } + + return nil +} + +func cloneStringAnyMap(source map[string]any) map[string]any { + if len(source) == 0 { + return nil + } + cloned := make(map[string]any, len(source)) + for key, value := range source { + cloned[key] = cloneAnyValue(value) + } + return cloned +} + +func cloneAnyValue(value any) any { + switch typed := value.(type) { + case map[string]any: + return cloneStringAnyMap(typed) + case []any: + cloned := make([]any, len(typed)) + for index, item := range typed { + cloned[index] = cloneAnyValue(item) + } + return cloned + default: + return value + } +} + +func modelFromParsedOpenAIEvent(parsed map[string]any) string { + if model, ok := parsed["model"].(string); ok && model != "" { + return model + } + + responseBody, ok := parsed["response_body"].(string) + if !ok || responseBody == "" || !json.Valid([]byte(responseBody)) { + return "" + } + + var responseObject map[string]any + if err := json.Unmarshal([]byte(responseBody), &responseObject); err != nil { + return "" + } + + if model, ok := responseObject["model"].(string); ok { + return model + } + + return "" +} + +func modelFromRequestBody(requestBody []byte) string { + if len(requestBody) == 0 { + return "" + } + + if json.Valid(requestBody) { + var requestObject map[string]any + if err := json.Unmarshal(requestBody, &requestObject); err != nil { + return "" + } + + if model, ok := requestObject["model"].(string); ok { + return model + } + } + + if model := multipartFieldValue(requestBody, "model"); model != "" { + return model + } + + return "" +} + +func multipartFieldValue(requestBody []byte, fieldName string) string { + boundary := multipartBoundaryFromBody(requestBody) + if boundary == "" { + return "" + } + + reader := multipart.NewReader(bytes.NewReader(requestBody), boundary) + for { + part, err := reader.NextPart() + if err == io.EOF { + return "" + } + if err != nil { + return "" + } + + if part.FormName() != fieldName { + _ = part.Close() + continue + } + + value, err := io.ReadAll(io.LimitReader(part, 4096)) + _ = part.Close() + if err != nil { + return "" + } + + return strings.TrimSpace(string(value)) + } +} + +func multipartBoundaryFromBody(requestBody []byte) string { + line := requestBody + if newlineIndex := bytes.IndexByte(line, '\n'); newlineIndex >= 0 { + line = line[:newlineIndex] + } + + line = bytes.TrimSpace(bytes.TrimSuffix(line, []byte("\r"))) + if len(line) <= 2 || !bytes.HasPrefix(line, []byte("--")) { + return "" + } + + return string(line[2:]) +} + +func floatToInt(value any) (int, bool) { + switch typed := value.(type) { + case int: + return typed, true + case int32: + return int(typed), true + case int64: + return int(typed), true + case float64: + return int(typed), true + case float32: + return int(typed), true + default: + return 0, false + } +} + +func firstUsageValue(usage map[string]any, keys ...string) (int, bool) { + for _, key := range keys { + if value, ok := floatToInt(usage[key]); ok { + return value, true + } + } + + return 0, false +} + +func firstIntMapValue(values map[string]int, keys ...string) int { + for _, key := range keys { + if value, ok := values[key]; ok { + return value + } + } + + return 0 +} + +func firstFloatMapValue(values map[string]float64, keys ...string) int { + for _, key := range keys { + if value, ok := values[key]; ok { + return int(value) + } + } + + return 0 +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + + return "" +} + +func durationMilliseconds(duration time.Duration) int64 { + if duration <= 0 { + return 0 + } + + if milliseconds := duration.Milliseconds(); milliseconds > 0 { + return milliseconds + } + + return 1 +} + +func metadataStringValue(metadata map[string]any, key string) string { + if metadata == nil { + return "" + } + + value, _ := metadata[key].(string) + return value +} + +func eventNameForEvent(evt eventbus.Event) string { + if evt.Status != 0 || evt.Duration > 0 || len(evt.ResponseBody) > 0 || len(evt.ResponseHeaders) > 0 { + return "finish" + } + + return "start" +} + +func usesOpenAIResponseTransformer(path string) bool { + return strings.HasPrefix(path, "/v1/completions") || + strings.HasPrefix(path, "/v1/chat/completions") || + strings.HasPrefix(path, "/v1/embeddings") || + strings.HasPrefix(path, "/v1/responses") || + strings.HasPrefix(path, "/v1/threads/") || + strings.HasPrefix(path, "/v1/audio/transcriptions") || + strings.HasPrefix(path, "/v1/audio/translations") +} + +func fallbackTokensUsage(requestBody, responseBody []byte, model string) *TokensUsage { + promptTokens := promptTokensFromRequestBody(requestBody, model) + completionTokens := completionTokensFromResponseBody(responseBody, model) + if promptTokens == 0 && completionTokens == 0 { + return nil + } + + return &TokensUsage{ + Input: promptTokens, + Output: completionTokens, + Total: promptTokens + completionTokens, + } +} + +func promptTokensFromRequestBody(requestBody []byte, model string) int { + promptSource := "" + if json.Valid(requestBody) { + var requestObject map[string]any + if err := json.Unmarshal(requestBody, &requestObject); err == nil { + if messages, ok := requestObject["messages"]; ok { + if encodedMessages, err := json.Marshal(messages); err == nil { + promptSource = string(encodedMessages) + } + } + if instructions, ok := requestObject["instructions"].(string); ok && instructions != "" { + promptSource += instructions + } + if inputSource := promptInputSource(requestObject["input"]); inputSource != "" { + promptSource += inputSource + } + } + } + if promptSource == "" { + promptSource = multipartFieldValue(requestBody, "prompt") + } + if promptSource == "" { + return 0 + } + + tokens, err := eventtransformer.CountOpenAITokensForModel(promptSource, model) + if err != nil { + return 0 + } + + return tokens +} + +func promptInputSource(input any) string { + switch typed := input.(type) { + case string: + return typed + case nil: + return "" + default: + encodedInput, err := json.Marshal(typed) + if err != nil { + return "" + } + return string(encodedInput) + } +} + +func completionTokensFromResponseBody(responseBody []byte, model string) int { + if len(responseBody) == 0 { + return 0 + } + + content, ok := assistantReplyContentFromResponseBody(responseBody) + if !ok || content == "" { + return 0 + } + + tokens, err := eventtransformer.CountOpenAITokensForModel(content, model) + if err != nil { + return 0 + } + + return tokens +} + +func assistantReplyContentFromResponseBody(responseBody []byte) (string, bool) { + responseObject, ok := responseObjectFromBody(responseBody) + if !ok { + return "", false + } + + if choices, ok := responseObject["choices"].([]any); ok && len(choices) > 0 { + if choice, ok := choices[0].(map[string]any); ok { + if message, ok := choice["message"].(map[string]any); ok { + if content, ok := message["content"].(string); ok { + return content, true + } + } + } + } + + if choices, ok := responseObject["choices"].([]map[string]any); ok && len(choices) > 0 { + if message, ok := choices[0]["message"].(map[string]any); ok { + if content, ok := message["content"].(string); ok { + return content, true + } + } + } + + if content, ok := responseOutputText(responseObject); ok { + return content, true + } + + if text, ok := responseObject["text"].(string); ok && text != "" { + return text, true + } + + return "", false +} + +func responseOutputText(responseObject map[string]any) (string, bool) { + outputItems, ok := responseObject["output"].([]any) + if !ok { + return "", false + } + + var content strings.Builder + for _, outputItem := range outputItems { + itemMap, ok := outputItem.(map[string]any) + if !ok { + continue + } + + segments, ok := itemMap["content"].([]any) + if !ok { + continue + } + + for _, segment := range segments { + segmentMap, ok := segment.(map[string]any) + if !ok { + continue + } + + if text, ok := responseOutputSegmentText(segmentMap); ok { + content.WriteString(text) + } + } + } + + if content.Len() == 0 { + return "", false + } + + return content.String(), true +} + +func responseOutputSegmentText(segment map[string]any) (string, bool) { + if text, ok := segment["text"].(string); ok && text != "" { + return text, true + } + + if textMap, ok := segment["text"].(map[string]any); ok { + if value, ok := textMap["value"].(string); ok && value != "" { + return value, true + } + } + + if text, ok := segment["output_text"].(string); ok && text != "" { + return text, true + } + + return "", false +} + +func responseObjectFromBody(responseBody []byte) (map[string]any, bool) { + if len(responseBody) == 0 { + return nil, false + } + + if json.Valid(responseBody) { + var responseObject map[string]any + if err := json.Unmarshal(responseBody, &responseObject); err == nil { + return responseObject, true + } + } + + if !utf8.Valid(responseBody) { + return nil, false + } + + responseText := string(responseBody) + if strings.Contains(responseText, "event: ") && strings.Contains(responseText, "data: ") { + if merged, err := eventtransformer.MergeThreadStreamingChunks(responseText); err == nil { + return merged, true + } + } + + if eventtransformer.IsOpenAIStreaming(responseText) { + if merged, err := eventtransformer.MergeOpenAIStreamingChunks(responseText); err == nil { + return merged, true + } + } + + return nil, false +} diff --git a/internal/dispatcher/transformer_additional_test.go b/internal/dispatcher/transformer_additional_test.go index 917f0fef..713701f5 100644 --- a/internal/dispatcher/transformer_additional_test.go +++ b/internal/dispatcher/transformer_additional_test.go @@ -2,8 +2,11 @@ package dispatcher import ( "bytes" + "compress/gzip" "encoding/json" + "mime/multipart" "net/http" + "reflect" "testing" "time" @@ -12,6 +15,249 @@ import ( "github.com/sofatutor/llm-proxy/internal/eventbus" ) +func TestCloneStringAnyMap_DeepClone(t *testing.T) { + original := map[string]any{ + "message": "hello", + "nested": map[string]any{ + "count": 3, + }, + "items": []any{ + map[string]any{"kind": "text", "value": "alpha"}, + "tail", + }, + } + + cloned := cloneStringAnyMap(original) + if cloned == nil { + t.Fatal("expected cloned map") + } + if reflect.DeepEqual(cloned, original) == false { + t.Fatalf("expected equal content, got %#v want %#v", cloned, original) + } + + cloned["message"] = "changed" + clonedNested := cloned["nested"].(map[string]any) + clonedNested["count"] = 9 + clonedItems := cloned["items"].([]any) + clonedItems[0].(map[string]any)["value"] = "beta" + + if original["message"] != "hello" { + t.Fatalf("expected top-level value to remain unchanged, got %v", original["message"]) + } + if original["nested"].(map[string]any)["count"] != 3 { + t.Fatalf("expected nested map to remain unchanged, got %v", original["nested"]) + } + if original["items"].([]any)[0].(map[string]any)["value"] != "alpha" { + t.Fatalf("expected nested slice map to remain unchanged, got %v", original["items"]) + } + + if cloneStringAnyMap(map[string]any{}) != nil { + t.Fatal("expected empty source to return nil") + } +} + +func TestFirstMapValues(t *testing.T) { + intValues := map[string]int{"prompt_tokens": 7} + if got := firstIntMapValue(intValues, "input_tokens", "prompt_tokens"); got != 7 { + t.Fatalf("expected int fallback value 7, got %d", got) + } + if got := firstIntMapValue(intValues, "missing"); got != 0 { + t.Fatalf("expected missing int value 0, got %d", got) + } + + floatValues := map[string]float64{"completion_tokens": 4.9} + if got := firstFloatMapValue(floatValues, "output_tokens", "completion_tokens"); got != 4 { + t.Fatalf("expected float fallback value 4, got %d", got) + } + if got := firstFloatMapValue(floatValues, "missing"); got != 0 { + t.Fatalf("expected missing float value 0, got %d", got) + } +} + +func TestFloatToIntAndFirstUsageValue(t *testing.T) { + tests := []struct { + name string + value any + want int + wantOK bool + }{ + {name: "int", value: int(3), want: 3, wantOK: true}, + {name: "int32", value: int32(4), want: 4, wantOK: true}, + {name: "int64", value: int64(5), want: 5, wantOK: true}, + {name: "float32", value: float32(6.9), want: 6, wantOK: true}, + {name: "float64", value: float64(7.2), want: 7, wantOK: true}, + {name: "unsupported", value: "8", want: 0, wantOK: false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := floatToInt(tt.value) + if ok != tt.wantOK { + t.Fatalf("expected ok=%v, got %v", tt.wantOK, ok) + } + if got != tt.want { + t.Fatalf("expected %d, got %d", tt.want, got) + } + }) + } + + usage := map[string]any{"completion_tokens": float64(9), "total_tokens": int64(12)} + if got, ok := firstUsageValue(usage, "output_tokens", "completion_tokens"); !ok || got != 9 { + t.Fatalf("expected completion token fallback 9, got %d ok=%v", got, ok) + } + if got, ok := firstUsageValue(usage, "total_tokens"); !ok || got != 12 { + t.Fatalf("expected total token value 12, got %d ok=%v", got, ok) + } + if _, ok := firstUsageValue(usage, "missing"); ok { + t.Fatal("expected missing usage keys to return ok=false") + } +} + +func TestPromptInputSource(t *testing.T) { + if got := promptInputSource("hello"); got != "hello" { + t.Fatalf("expected string input source, got %q", got) + } + if got := promptInputSource(nil); got != "" { + t.Fatalf("expected nil input source to be empty, got %q", got) + } + if got := promptInputSource([]any{"alpha", map[string]any{"beta": true}}); got != `["alpha",{"beta":true}]` { + t.Fatalf("expected structured input to be marshaled, got %q", got) + } +} + +func TestAssistantReplyContentFromResponseBody(t *testing.T) { + if content, ok := assistantReplyContentFromResponseBody([]byte(`{"choices":[{"message":{"content":"hello"}}]}`)); !ok || content != "hello" { + t.Fatalf("expected chat choices content, got %q ok=%v", content, ok) + } + + responseObject := map[string]any{ + "choices": []map[string]any{{ + "message": map[string]any{"content": "typed"}, + }}, + } + responseBytes, err := json.Marshal(responseObject) + if err != nil { + t.Fatalf("marshal response err: %v", err) + } + if content, ok := assistantReplyContentFromResponseBody(responseBytes); !ok || content != "typed" { + t.Fatalf("expected typed choices content, got %q ok=%v", content, ok) + } + + if content, ok := assistantReplyContentFromResponseBody([]byte(`{"output":[{"content":[{"output_text":"fallback"}]}]}`)); !ok || content != "fallback" { + t.Fatalf("expected responses output fallback, got %q ok=%v", content, ok) + } + + if _, ok := assistantReplyContentFromResponseBody([]byte(`{"choices":[]}`)); ok { + t.Fatal("expected false when no assistant content is present") + } +} + +func TestResponseObjectFromBody(t *testing.T) { + if responseObject, ok := responseObjectFromBody([]byte(`{"id":"resp_1"}`)); !ok || responseObject["id"] != "resp_1" { + t.Fatalf("expected JSON response body to decode, got %#v ok=%v", responseObject, ok) + } + + chatStream := "data: {\"id\":\"chatcmpl_1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n" + + "data: {\"choices\":[{\"delta\":{\"content\":\" world\"},\"finish_reason\":\"stop\"}]}\n\n" + + "data: [DONE]\n\n" + if responseObject, ok := responseObjectFromBody([]byte(chatStream)); !ok || responseObject["id"] != "chatcmpl_1" { + t.Fatalf("expected streaming response body to merge, got %#v ok=%v", responseObject, ok) + } + + threadStream := "event: thread.message.delta\n" + + "data: {\"delta\":{\"content\":[{\"type\":\"text\",\"text\":{\"value\":\"hello\"}}]}}\n\n" + + "event: thread.run.completed\n" + + "data: {\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n" + if responseObject, ok := responseObjectFromBody([]byte(threadStream)); !ok || responseObject["usage"] == nil { + t.Fatalf("expected thread stream body to merge, got %#v ok=%v", responseObject, ok) + } + + if _, ok := responseObjectFromBody([]byte{0xff, 0xfe, 0xfd}); ok { + t.Fatal("expected invalid UTF-8 response body to fail") + } + if _, ok := responseObjectFromBody(nil); ok { + t.Fatal("expected empty response body to fail") + } +} + +func TestResponseOutputText(t *testing.T) { + response := map[string]any{ + "output": []any{ + map[string]any{ + "content": []any{ + map[string]any{"text": "Hello"}, + map[string]any{"text": map[string]any{"value": " there"}}, + map[string]any{"output_text": "!"}, + map[string]any{"ignored": true}, + }, + }, + map[string]any{"content": []any{"skip-non-map"}}, + "skip-non-map", + }, + } + + content, ok := responseOutputText(response) + if !ok { + t.Fatal("expected response output text") + } + if content != "Hello there!" { + t.Fatalf("expected concatenated content, got %q", content) + } + + if _, ok := responseOutputText(map[string]any{"output": []any{map[string]any{"content": []any{map[string]any{"ignored": true}}}}}); ok { + t.Fatal("expected false when no text segments are present") + } + if _, ok := responseOutputText(map[string]any{"output": "invalid"}); ok { + t.Fatal("expected false for invalid output shape") + } +} + +func TestResponseOutputSegmentText(t *testing.T) { + tests := []struct { + name string + segment map[string]any + want string + wantOK bool + }{ + { + name: "plain text", + segment: map[string]any{"text": "hello"}, + want: "hello", + wantOK: true, + }, + { + name: "nested value", + segment: map[string]any{"text": map[string]any{"value": "nested"}}, + want: "nested", + wantOK: true, + }, + { + name: "output text fallback", + segment: map[string]any{"output_text": "fallback"}, + want: "fallback", + wantOK: true, + }, + { + name: "missing content", + segment: map[string]any{"text": map[string]any{"value": ""}}, + want: "", + wantOK: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := responseOutputSegmentText(tt.segment) + if ok != tt.wantOK { + t.Fatalf("expected ok=%v, got %v", tt.wantOK, ok) + } + if got != tt.want { + t.Fatalf("expected %q, got %q", tt.want, got) + } + }) + } +} + // Covers gzip Content-Encoding with non-gzipped but valid JSON payload: should fall back to raw JSON func TestSafeRawMessageOrBase64_GzipHeaderButPlainJSON(t *testing.T) { headers := map[string][]string{ @@ -93,3 +339,516 @@ func TestDefaultEventTransformer_Transform_VerboseHeaders(t *testing.T) { t.Fatalf("expected multi-value header preserved, got %T %v", headers["X-Custom"], headers["X-Custom"]) } } + +func TestDefaultEventTransformer_Transform_ResponsesUsageAndTokenMetadata(t *testing.T) { + tr := NewDefaultEventTransformer(false) + evt := eventbus.Event{ + RequestID: "req-123", + Method: "POST", + Path: "/v1/responses", + ProjectID: "project-1", + TokenID: "token-uuid-1", + TokenMetadata: map[string]string{"feature": "sofabuddy", "user_id": "42"}, + Status: 200, + Duration: 5 * time.Millisecond, + RequestBody: []byte(`{"model":"gpt-4.1-mini","input":"hello"}`), + ResponseBody: []byte(`{"id":"resp_1","model":"gpt-4.1-mini","usage":{"input_tokens":11,"output_tokens":7,"total_tokens":18},"output":[]}`), + ResponseHeaders: http.Header{ + "Content-Type": {"application/json"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.UserID == nil || *payload.UserID != "42" { + t.Fatalf("expected user id 42, got %v", payload.UserID) + } + if payload.TokensUsage == nil || payload.TokensUsage.Input != 11 || payload.TokensUsage.Output != 7 || payload.TokensUsage.Total != 18 { + t.Fatalf("unexpected token usage: %#v", payload.TokensUsage) + } + if payload.Metadata["project_id"] != "project-1" { + t.Fatalf("expected project_id metadata, got %v", payload.Metadata["project_id"]) + } + if payload.Metadata["token_id"] != "token-uuid-1" { + t.Fatalf("expected token_id metadata, got %v", payload.Metadata["token_id"]) + } + if payload.Metadata["model"] != "gpt-4.1-mini" { + t.Fatalf("expected model metadata, got %v", payload.Metadata["model"]) + } + tokenMetadata, ok := payload.Metadata["token_metadata"].(map[string]string) + if !ok { + t.Fatalf("expected token metadata map, got %T", payload.Metadata["token_metadata"]) + } + if tokenMetadata["feature"] != "sofabuddy" { + t.Fatalf("expected feature metadata, got %v", tokenMetadata) + } +} + +func TestDefaultEventTransformer_Transform_EmbeddingsUsageNormalization(t *testing.T) { + tr := NewDefaultEventTransformer(false) + evt := eventbus.Event{ + RequestID: "req-embeddings-1", + Method: "POST", + Path: "/v1/embeddings", + Status: 200, + Duration: 8 * time.Millisecond, + RequestBody: []byte(`{"model":"text-embedding-3-small","input":"hello embeddings"}`), + ResponseBody: []byte(`{"object":"list","data":[{"object":"embedding","embedding":[0.1,0.2],"index":0}],"model":"text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}`), + ResponseHeaders: http.Header{ + "Content-Type": {"application/json"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.Metadata["model"] != "text-embedding-3-small" { + t.Fatalf("expected embeddings model metadata, got %v", payload.Metadata["model"]) + } + if payload.TokensUsage == nil { + t.Fatalf("expected embeddings token usage") + } + if payload.TokensUsage.Input != 2 || payload.TokensUsage.Output != 0 || payload.TokensUsage.Total != 2 { + t.Fatalf("unexpected embeddings token usage: %#v", payload.TokensUsage) + } + if len(payload.Output) == 0 { + t.Fatalf("expected embeddings output to be preserved") + } +} + +func TestDefaultEventTransformer_Transform_ResponsesTokenUsageFallbackFromOutputText(t *testing.T) { + tr := NewDefaultEventTransformer(false) + evt := eventbus.Event{ + RequestID: "req-responses-1", + Method: "POST", + Path: "/v1/responses", + Status: 200, + Duration: 12 * time.Millisecond, + RequestBody: []byte(`{"model":"gpt-4.1-mini","input":"hello"}`), + ResponseBody: []byte(`{"id":"resp_1","model":"gpt-4.1-mini","output":[{"type":"message","content":[{"type":"output_text","text":"hi there"}]}]}`), + ResponseHeaders: http.Header{ + "Content-Type": {"application/json"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.TokensUsage == nil { + t.Fatalf("expected computed token usage") + } + if payload.TokensUsage.Output <= 0 { + t.Fatalf("expected responses completion tokens > 0, got %#v", payload.TokensUsage) + } +} + +func TestDefaultEventTransformer_Transform_ResponsesStreamUsesCompletedEvent(t *testing.T) { + tr := NewDefaultEventTransformer(false) + stream := "event: response.created\n" + + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":123,\"status\":\"in_progress\",\"model\":\"gpt-4.1-mini\",\"output\":[]}}\n\n" + + "event: response.output_text.delta\n" + + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"Hallo Welt\"}\n\n" + + "event: response.completed\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":123,\"status\":\"completed\",\"model\":\"gpt-4.1-mini\",\"usage\":{\"input_tokens\":40,\"output_tokens\":12,\"total_tokens\":52},\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Hallo Welt\"}]}]}}\n\n" + + "data: [DONE]\n\n" + vt := eventbus.Event{ + RequestID: "req-responses-stream-1", + Method: "POST", + Path: "/v1/responses", + Status: 200, + Duration: 12 * time.Millisecond, + RequestBody: []byte(`{"model":"gpt-4.1-mini","input":"hello"}`), + ResponseBody: []byte(stream), + ResponseHeaders: http.Header{ + "Content-Type": {"text/event-stream"}, + }, + } + + payload, err := tr.Transform(vt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.TokensUsage == nil { + t.Fatalf("expected token usage from response.completed") + } + if payload.TokensUsage.Input != 40 || payload.TokensUsage.Output != 12 || payload.TokensUsage.Total != 52 { + t.Fatalf("expected final stream usage 40/12, got %#v", payload.TokensUsage) + } + if payload.Metadata["model"] != "gpt-4.1-mini" { + t.Fatalf("expected model metadata, got %v", payload.Metadata["model"]) + } + var output map[string]any + if err := json.Unmarshal(payload.Output, &output); err != nil { + t.Fatalf("failed to unmarshal payload output: %v", err) + } + if output["status"] != "completed" { + t.Fatalf("expected completed merged response status, got %v", output["status"]) + } +} + +func TestDefaultEventTransformer_Transform_ModelFallbackFromRequestBody(t *testing.T) { + tr := NewDefaultEventTransformer(false) + evt := eventbus.Event{ + RequestID: "req-chat-1", + Method: "POST", + Path: "/v1/chat/completions", + Status: 200, + Duration: 12 * time.Millisecond, + RequestBody: []byte(`{"model":"gpt-4.1-mini","messages":[{"role":"user","content":"hello"}]}`), + ResponseBody: []byte(`{"id":"chatcmpl_1","choices":[{"index":0,"message":{"role":"assistant","content":"hi there"},"finish_reason":"stop"}]}`), + ResponseHeaders: http.Header{ + "Content-Type": {"application/json"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.Metadata["model"] != "gpt-4.1-mini" { + t.Fatalf("expected request model fallback, got %v", payload.Metadata["model"]) + } +} + +func TestDefaultEventTransformer_Transform_AudioTranscriptionExtractsMultipartModelAndUsage(t *testing.T) { + tr := NewDefaultEventTransformer(false) + + var requestBody bytes.Buffer + writer := multipart.NewWriter(&requestBody) + if err := writer.WriteField("model", "gpt-4o-transcribe"); err != nil { + t.Fatalf("write model field: %v", err) + } + if err := writer.WriteField("language", "de"); err != nil { + t.Fatalf("write language field: %v", err) + } + filePart, err := writer.CreateFormFile("file", "recording.wav") + if err != nil { + t.Fatalf("create file part: %v", err) + } + if _, err := filePart.Write([]byte("fake-audio-payload")); err != nil { + t.Fatalf("write file part: %v", err) + } + if err := writer.Close(); err != nil { + t.Fatalf("close multipart writer: %v", err) + } + + evt := eventbus.Event{ + RequestID: "req-audio-transcription-1", + Method: "POST", + Path: "/v1/audio/transcriptions", + Status: 200, + Duration: 25 * time.Millisecond, + RequestBody: requestBody.Bytes(), + ResponseBody: []byte(`{"text":"Hallo Welt","usage":{"input_tokens":12,"output_tokens":3,"total_tokens":15}}`), + ResponseHeaders: http.Header{ + "Content-Type": {"application/json"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.Metadata["model"] != "gpt-4o-transcribe" { + t.Fatalf("expected multipart model metadata, got %v", payload.Metadata["model"]) + } + if payload.TokensUsage == nil { + t.Fatalf("expected transcription usage") + } + if payload.TokensUsage.Input != 12 || payload.TokensUsage.Output != 3 || payload.TokensUsage.Total != 15 { + t.Fatalf("unexpected transcription usage: %#v", payload.TokensUsage) + } +} + +func TestDefaultEventTransformer_Transform_AudioTranscriptionFallsBackToTranscriptTextUsage(t *testing.T) { + tr := NewDefaultEventTransformer(false) + + var requestBody bytes.Buffer + writer := multipart.NewWriter(&requestBody) + if err := writer.WriteField("model", "whisper-1"); err != nil { + t.Fatalf("write model field: %v", err) + } + filePart, err := writer.CreateFormFile("file", "recording.wav") + if err != nil { + t.Fatalf("create file part: %v", err) + } + if _, err := filePart.Write([]byte("fake-audio-payload")); err != nil { + t.Fatalf("write file part: %v", err) + } + if err := writer.Close(); err != nil { + t.Fatalf("close multipart writer: %v", err) + } + + evt := eventbus.Event{ + RequestID: "req-audio-transcription-2", + Method: "POST", + Path: "/v1/audio/transcriptions", + Status: 200, + Duration: 25 * time.Millisecond, + RequestBody: requestBody.Bytes(), + ResponseBody: []byte(`{"text":"Hallo Welt"}`), + ResponseHeaders: http.Header{ + "Content-Type": {"application/json"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.Metadata["model"] != "whisper-1" { + t.Fatalf("expected multipart model metadata, got %v", payload.Metadata["model"]) + } + if payload.TokensUsage == nil { + t.Fatalf("expected fallback token usage") + } + if payload.TokensUsage.Input != 0 { + t.Fatalf("expected zero input tokens for multipart transcription fallback, got %#v", payload.TokensUsage) + } + if payload.TokensUsage.Output <= 0 { + t.Fatalf("expected transcription text fallback tokens > 0, got %#v", payload.TokensUsage) + } + if payload.TokensUsage.Total != payload.TokensUsage.Output { + t.Fatalf("expected total tokens to match output-only fallback, got %#v", payload.TokensUsage) + } +} + +func TestDefaultEventTransformer_Transform_TokenUsageFallbackFromComputedUsage(t *testing.T) { + tr := NewDefaultEventTransformer(false) + evt := eventbus.Event{ + RequestID: "req-chat-2", + Method: "POST", + Path: "/v1/chat/completions", + Status: 200, + Duration: 12 * time.Millisecond, + RequestBody: []byte(`{"model":"gpt-4.1-mini","messages":[{"role":"user","content":"hello"}]}`), + ResponseBody: []byte(`{"id":"chatcmpl_1","choices":[{"index":0,"message":{"role":"assistant","content":"hi there"},"finish_reason":"stop"}]}`), + ResponseHeaders: http.Header{ + "Content-Type": {"application/json"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.TokensUsage == nil { + t.Fatalf("expected computed token usage") + } + if payload.TokensUsage.Input <= 0 { + t.Fatalf("expected prompt tokens > 0, got %#v", payload.TokensUsage) + } + if payload.TokensUsage.Output <= 0 { + t.Fatalf("expected completion tokens > 0, got %#v", payload.TokensUsage) + } +} + +func TestDefaultEventTransformer_Transform_TokenUsageFallbackFromGzippedChatResponse(t *testing.T) { + tr := NewDefaultEventTransformer(false) + + var compressed bytes.Buffer + gzipWriter := gzip.NewWriter(&compressed) + _, err := gzipWriter.Write([]byte(`{"id":"chatcmpl_1","choices":[{"index":0,"message":{"role":"assistant","content":"{\"suggestions\":[\"A\",\"B\",\"C\"]}"},"finish_reason":"stop"}]}`)) + if err != nil { + t.Fatalf("gzip write err: %v", err) + } + if err := gzipWriter.Close(); err != nil { + t.Fatalf("gzip close err: %v", err) + } + + evt := eventbus.Event{ + RequestID: "req-chat-gzip-1", + Method: "POST", + Path: "/v1/chat/completions", + Status: 200, + Duration: 12 * time.Millisecond, + RequestBody: []byte(`{"model":"gpt-4.1-mini","messages":[{"role":"user","content":"hello"}]}`), + ResponseBody: compressed.Bytes(), + ResponseHeaders: http.Header{ + "Content-Type": {"application/json"}, + "Content-Encoding": {"gzip"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.TokensUsage == nil { + t.Fatalf("expected token usage") + } + if payload.TokensUsage.Output <= 0 { + t.Fatalf("expected completion tokens > 0 for gzipped chat response, got %#v", payload.TokensUsage) + } + if payload.Metadata["model"] != "gpt-4.1-mini" { + t.Fatalf("expected model metadata, got %v", payload.Metadata["model"]) + } +} + +func TestDefaultEventTransformer_Transform_TokenUsageFallbackFromChatStream(t *testing.T) { + tr := NewDefaultEventTransformer(false) + stream := "data: {\"id\":\"chatcmpl_1\",\"object\":\"chat.completion.chunk\",\"created\":123,\"model\":\"gpt-4.1-mini\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n" + + "data: {\"choices\":[{\"delta\":{\"content\":\" there\"},\"finish_reason\":\"stop\"}]}\n\n" + + "data: [DONE]\n\n" + evt := eventbus.Event{ + RequestID: "req-chat-stream-1", + Method: "POST", + Path: "/v1/chat/completions", + Status: 200, + Duration: 12 * time.Millisecond, + RequestBody: []byte(`{"model":"gpt-4.1-mini","messages":[{"role":"user","content":"hello"}]}`), + ResponseBody: []byte(stream), + ResponseHeaders: http.Header{ + "Content-Type": {"text/event-stream"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.TokensUsage == nil { + t.Fatalf("expected computed token usage from stream") + } + if payload.TokensUsage.Output <= 0 { + t.Fatalf("expected stream completion tokens > 0, got %#v", payload.TokensUsage) + } + if payload.Metadata["model"] != "gpt-4.1-mini" { + t.Fatalf("expected model metadata from request body, got %v", payload.Metadata["model"]) + } +} + +func TestDefaultEventTransformer_Transform_TokenUsageFallbackFromThreadStream(t *testing.T) { + tr := NewDefaultEventTransformer(false) + stream := "event: thread.run.created\n" + + "data: {\"id\":\"run_1\",\"assistant_id\":\"asst_1\",\"thread_id\":\"thread_1\",\"status\":\"queued\",\"created_at\":123}\n\n" + + "event: thread.message.delta\n" + + "data: {\"delta\":{\"content\":[{\"type\":\"text\",\"text\":{\"value\":\"hello\"}}]}}\n\n" + + "event: thread.message.completed\n" + + "data: {\"content\":[{\"type\":\"text\",\"text\":{\"value\":\"hello there\"}}]}\n\n" + + "event: thread.run.completed\n" + + "data: {\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n" + evt := eventbus.Event{ + RequestID: "req-thread-stream-1", + Method: "POST", + Path: "/v1/threads/runs", + Status: 200, + Duration: 12 * time.Millisecond, + RequestBody: []byte(`{"model":"gpt-4.1-mini","input":"hello"}`), + ResponseBody: []byte(stream), + ResponseHeaders: http.Header{ + "Content-Type": {"text/event-stream"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.TokensUsage == nil { + t.Fatalf("expected token usage from thread stream") + } + if payload.TokensUsage.Output <= 0 { + t.Fatalf("expected thread stream completion tokens > 0, got %#v", payload.TokensUsage) + } +} + +func TestDefaultEventTransformer_Transform_ChatErrorStillKeepsModelAndPromptTokens(t *testing.T) { + tr := NewDefaultEventTransformer(false) + evt := eventbus.Event{ + RequestID: "req-chat-404", + Method: "POST", + Path: "/v1/chat/completions", + Status: 404, + Duration: 600 * time.Microsecond, + RequestBody: []byte(`{"model":"gpt-4.1-mini","messages":[{"role":"user","content":"hello"}]}`), + ResponseBody: []byte(`not found`), + ResponseHeaders: http.Header{ + "Content-Type": {"text/plain"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.Metadata["model"] != "gpt-4.1-mini" { + t.Fatalf("expected request model on error response, got %v", payload.Metadata["model"]) + } + if payload.TokensUsage == nil || payload.TokensUsage.Input <= 0 || payload.TokensUsage.Output != 0 { + t.Fatalf("expected prompt-only token usage on error response, got %#v", payload.TokensUsage) + } + if payload.Metadata["duration_ms"] != int64(1) { + t.Fatalf("expected sub-millisecond duration to round up to 1ms, got %v", payload.Metadata["duration_ms"]) + } +} + +func TestDefaultEventTransformer_Transform_ModerationsModelFallback(t *testing.T) { + tr := NewDefaultEventTransformer(false) + evt := eventbus.Event{ + RequestID: "req-mod-1", + Method: "POST", + Path: "/v1/moderations", + Status: 200, + Duration: 2 * time.Millisecond, + RequestBody: []byte(`{"model":"omni-moderation-latest","input":"hello"}`), + ResponseBody: []byte(`{"id":"modr-1","results":[{"flagged":false}]}`), + ResponseHeaders: http.Header{ + "Content-Type": {"application/json"}, + }, + } + + payload, err := tr.Transform(evt) + if err != nil { + t.Fatalf("Transform err: %v", err) + } + if payload == nil { + t.Fatalf("expected payload") + } + if payload.Metadata["model"] != "omni-moderation-latest" { + t.Fatalf("expected moderation model from request body, got %v", payload.Metadata["model"]) + } + if payload.TokensUsage == nil || payload.TokensUsage.Input <= 0 { + t.Fatalf("expected moderation prompt tokens from input fallback, got %#v", payload.TokensUsage) + } +} diff --git a/internal/dispatcher/transformer_test.go b/internal/dispatcher/transformer_test.go index c15adfe6..26c6ae38 100644 --- a/internal/dispatcher/transformer_test.go +++ b/internal/dispatcher/transformer_test.go @@ -97,8 +97,8 @@ func TestDefaultEventTransformer_Transform(t *testing.T) { t.Error("Expected non-empty RunID") } - if result.Event != "start" { - t.Errorf("Expected Event 'start', got %s", result.Event) + if result.Event != "finish" { + t.Errorf("Expected Event 'finish', got %s", result.Event) } // Check metadata @@ -155,3 +155,22 @@ func TestDefaultEventTransformer_Transform(t *testing.T) { }) } } + +func TestDefaultEventTransformer_Transform_StartEventWithoutResponseState(t *testing.T) { + transformer := &DefaultEventTransformer{} + result, err := transformer.Transform(eventbus.Event{ + RequestID: "req-start", + Method: "POST", + Path: "/v1/chat/completions", + RequestBody: []byte(`{"model":"gpt-4.1-mini","messages":[]}`), + }) + if err != nil { + t.Fatalf("Transform failed: %v", err) + } + if result == nil { + t.Fatal("Expected non-nil result") + } + if result.Event != "start" { + t.Fatalf("Expected Event 'start', got %s", result.Event) + } +} diff --git a/internal/eventbus/eventbus.go b/internal/eventbus/eventbus.go index 0102724e..834fefa2 100644 --- a/internal/eventbus/eventbus.go +++ b/internal/eventbus/eventbus.go @@ -14,6 +14,9 @@ type Event struct { RequestID string Method string Path string + ProjectID string + TokenID string + TokenMetadata map[string]string Status int Duration time.Duration ResponseHeaders http.Header diff --git a/internal/eventbus/redis_streams.go b/internal/eventbus/redis_streams.go index f1187c87..67a9f67f 100644 --- a/internal/eventbus/redis_streams.go +++ b/internal/eventbus/redis_streams.go @@ -87,34 +87,42 @@ func (a *RedisStreamsClientAdapter) XInfoGroups(ctx context.Context, stream stri // RedisStreamsConfig holds configuration for Redis Streams event bus. type RedisStreamsConfig struct { - StreamKey string // Redis stream key name - ConsumerGroup string // Consumer group name - ConsumerName string // Unique consumer name within the group - MaxLen int64 // Max stream length (0 = unlimited, uses MAXLEN ~ approximation) - BlockTimeout time.Duration // Block timeout for XREADGROUP (0 = non-blocking) - ClaimMinIdleTime time.Duration // Minimum idle time before claiming pending messages - BatchSize int64 // Number of messages to read at once + StreamKey string // Redis stream key name + ConsumerGroup string // Consumer group name + ConsumerName string // Unique consumer name within the group + PublishBufferSize int // Bounded async producer queue size + MaxLen int64 // Max stream length (0 = unlimited, uses MAXLEN ~ approximation) + BlockTimeout time.Duration // Block timeout for XREADGROUP (0 = non-blocking) + ClaimMinIdleTime time.Duration // Minimum idle time before claiming pending messages + BatchSize int64 // Number of messages to read at once } // DefaultRedisStreamsConfig returns default configuration. func DefaultRedisStreamsConfig() RedisStreamsConfig { return RedisStreamsConfig{ - StreamKey: "llm-proxy-events", - ConsumerGroup: "llm-proxy-dispatchers", - ConsumerName: "dispatcher-1", - MaxLen: 10000, - BlockTimeout: 5 * time.Second, - ClaimMinIdleTime: 30 * time.Second, - BatchSize: 100, + StreamKey: "llm-proxy-events", + ConsumerGroup: "llm-proxy-dispatchers", + ConsumerName: "dispatcher-1", + PublishBufferSize: 1000, + MaxLen: 10000, + BlockTimeout: 5 * time.Second, + ClaimMinIdleTime: 30 * time.Second, + BatchSize: 100, } } +type publishRequest struct { + ctx context.Context + evt Event +} + // RedisStreamsEventBus implements EventBus using Redis Streams. // It provides durable, distributed event delivery with consumer groups, // acknowledgment, and at-least-once delivery semantics. type RedisStreamsEventBus struct { client RedisStreamsClient config RedisStreamsConfig + publishCh chan publishRequest stats busStats stopCh chan struct{} stopOnce sync.Once @@ -126,11 +134,18 @@ type RedisStreamsEventBus struct { // NewRedisStreamsEventBus creates a new Redis Streams event bus. func NewRedisStreamsEventBus(client RedisStreamsClient, config RedisStreamsConfig) *RedisStreamsEventBus { - return &RedisStreamsEventBus{ - client: client, - config: config, - stopCh: make(chan struct{}), + if config.PublishBufferSize <= 0 { + config.PublishBufferSize = DefaultRedisStreamsConfig().PublishBufferSize } + b := &RedisStreamsEventBus{ + client: client, + config: config, + publishCh: make(chan publishRequest, config.PublishBufferSize), + stopCh: make(chan struct{}), + } + b.wg.Add(1) + go b.publishLoop() + return b } // EnsureConsumerGroup creates the consumer group if it doesn't exist. @@ -165,6 +180,42 @@ func isGroupExistsError(err error) bool { // Publish adds an event to the Redis stream using XADD. func (b *RedisStreamsEventBus) Publish(ctx context.Context, evt Event) { + select { + case <-b.stopCh: + b.stats.dropped.Add(1) + return + default: + } + + request := publishRequest{ctx: ctx, evt: evt} + select { + case b.publishCh <- request: + default: + b.stats.dropped.Add(1) + } +} + +func (b *RedisStreamsEventBus) publishLoop() { + defer b.wg.Done() + + for { + select { + case request := <-b.publishCh: + b.publishToStream(request.ctx, request.evt) + case <-b.stopCh: + for { + select { + case request := <-b.publishCh: + b.publishToStream(request.ctx, request.evt) + default: + return + } + } + } + } +} + +func (b *RedisStreamsEventBus) publishToStream(ctx context.Context, evt Event) { data, err := json.Marshal(evt) if err != nil { log.Printf("[eventbus] Failed to marshal event: %v", err) diff --git a/internal/eventbus/redis_streams_test.go b/internal/eventbus/redis_streams_test.go index b41bd49b..f1600fc4 100644 --- a/internal/eventbus/redis_streams_test.go +++ b/internal/eventbus/redis_streams_test.go @@ -338,26 +338,19 @@ func TestRedisStreamsEventBus_Publish(t *testing.T) { client := newMockRedisStreamsClient() config := DefaultRedisStreamsConfig() bus := NewRedisStreamsEventBus(client, config) + defer bus.Stop() ctx := context.Background() evt := Event{RequestID: "test-123", Method: "POST", Path: "/v1/chat/completions", Status: 200} bus.Publish(ctx, evt) - // Verify message was added + waitForRedisStats(t, bus, 1, 0) + if len(client.stream) != 1 { t.Fatalf("expected 1 message in stream, got %d", len(client.stream)) } - // Verify stats - pub, drop := bus.Stats() - if pub != 1 { - t.Errorf("expected 1 published, got %d", pub) - } - if drop != 0 { - t.Errorf("expected 0 dropped, got %d", drop) - } - // Verify message content data := client.stream[0].Values["data"].(string) var parsedEvt Event @@ -375,17 +368,12 @@ func TestRedisStreamsEventBus_Publish_Error(t *testing.T) { config := DefaultRedisStreamsConfig() bus := NewRedisStreamsEventBus(client, config) + defer bus.Stop() ctx := context.Background() bus.Publish(ctx, Event{RequestID: "test"}) - pub, drop := bus.Stats() - if pub != 0 { - t.Errorf("expected 0 published on error, got %d", pub) - } - if drop != 1 { - t.Errorf("expected 1 dropped on error, got %d", drop) - } + waitForRedisStats(t, bus, 0, 1) } func TestRedisStreamsEventBus_Publish_MaxLen(t *testing.T) { @@ -393,6 +381,7 @@ func TestRedisStreamsEventBus_Publish_MaxLen(t *testing.T) { config := DefaultRedisStreamsConfig() config.MaxLen = 3 bus := NewRedisStreamsEventBus(client, config) + defer bus.Stop() ctx := context.Background() @@ -402,11 +391,51 @@ func TestRedisStreamsEventBus_Publish_MaxLen(t *testing.T) { } // Stream should be trimmed to MaxLen + waitForRedisStreamLength(t, client, 3) if len(client.stream) != 3 { t.Errorf("expected stream length 3 after trim, got %d", len(client.stream)) } } +func TestRedisStreamsEventBus_Publish_DropsWhenProducerQueueFull(t *testing.T) { + client := &blockingRedisStreamsClient{ + RedisStreamsClient: newMockRedisStreamsClient(), + release: make(chan struct{}), + started: make(chan struct{}, 1), + } + config := DefaultRedisStreamsConfig() + config.PublishBufferSize = 1 + bus := NewRedisStreamsEventBus(client, config) + + bus.Publish(context.Background(), Event{RequestID: "first"}) + <-client.started + bus.Publish(context.Background(), Event{RequestID: "second"}) + bus.Publish(context.Background(), Event{RequestID: "third"}) + + waitForRedisStats(t, bus, 0, 1) + + close(client.release) + waitForRedisStats(t, bus, 2, 1) + bus.Stop() +} + +func TestRedisStreamsEventBus_Publish_ReturnsWithoutWaitingForRedis(t *testing.T) { + client := &blockingRedisStreamsClient{RedisStreamsClient: newMockRedisStreamsClient(), release: make(chan struct{})} + config := DefaultRedisStreamsConfig() + config.PublishBufferSize = 1 + bus := NewRedisStreamsEventBus(client, config) + defer func() { + close(client.release) + bus.Stop() + }() + + start := time.Now() + bus.Publish(context.Background(), Event{RequestID: "non-blocking"}) + if elapsed := time.Since(start); elapsed > 50*time.Millisecond { + t.Fatalf("expected non-blocking publish, took %s", elapsed) + } +} + func TestRedisStreamsEventBus_EnsureConsumerGroup(t *testing.T) { client := newMockRedisStreamsClient() config := DefaultRedisStreamsConfig() @@ -615,6 +644,7 @@ func TestRedisStreamsEventBus_StreamLength(t *testing.T) { bus.Publish(ctx, Event{RequestID: fmt.Sprintf("r%d", i)}) } + waitForBusStreamLength(t, bus, 5) length, err = bus.StreamLength(ctx) if err != nil { t.Fatalf("unexpected error: %v", err) @@ -769,6 +799,7 @@ func TestRedisStreamsEventBus_Integration_Miniredis(t *testing.T) { } // Verify stream length + waitForBusStreamLength(t, bus, 5) length, err := bus.StreamLength(ctx) if err != nil { t.Fatalf("StreamLength error: %v", err) @@ -799,13 +830,7 @@ func TestRedisStreamsEventBus_Integration_Miniredis(t *testing.T) { } // Verify stats - pub, drop := bus.Stats() - if pub != 5 { - t.Errorf("expected 5 published, got %d", pub) - } - if drop != 0 { - t.Errorf("expected 0 dropped, got %d", drop) - } + waitForRedisStats(t, bus, 5, 0) } // Test RedisStreamsClientAdapter methods @@ -1047,10 +1072,73 @@ func TestRedisStreamsEventBus_Publish_JSONMarshalCoverage(t *testing.T) { } bus.Publish(ctx, evt) - pub, drop := bus.Stats() - if pub != 1 || drop != 0 { - t.Errorf("expected 1 published, 0 dropped, got pub=%d, drop=%d", pub, drop) + waitForRedisStats(t, bus, 1, 0) +} + +type blockingRedisStreamsClient struct { + RedisStreamsClient + release chan struct{} + started chan struct{} +} + +func (b *blockingRedisStreamsClient) XAdd(ctx context.Context, args *redis.XAddArgs) (string, error) { + if b.started != nil { + select { + case b.started <- struct{}{}: + default: + } + } + <-b.release + return b.RedisStreamsClient.XAdd(ctx, args) +} + +func waitForRedisStats(t *testing.T, bus *RedisStreamsEventBus, wantPublished, wantDropped int) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + published, dropped := bus.Stats() + if published == wantPublished && dropped == wantDropped { + return + } + time.Sleep(10 * time.Millisecond) + } + published, dropped := bus.Stats() + t.Fatalf("expected published=%d dropped=%d, got published=%d dropped=%d", wantPublished, wantDropped, published, dropped) +} + +func waitForRedisStreamLength(t *testing.T, client *mockRedisStreamsClient, want int) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + client.mu.Lock() + length := len(client.stream) + client.mu.Unlock() + if length == want { + return + } + time.Sleep(10 * time.Millisecond) + } + client.mu.Lock() + length := len(client.stream) + client.mu.Unlock() + t.Fatalf("expected stream length %d, got %d", want, length) +} + +func waitForBusStreamLength(t *testing.T, bus *RedisStreamsEventBus, want int64) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + length, err := bus.StreamLength(context.Background()) + if err == nil && length == want { + return + } + time.Sleep(10 * time.Millisecond) + } + length, err := bus.StreamLength(context.Background()) + if err != nil { + t.Fatalf("expected stream length %d, got error %v", want, err) } + t.Fatalf("expected stream length %d, got %d", want, length) } // Test XClaim with no eligible messages diff --git a/internal/eventtransformer/decode_test.go b/internal/eventtransformer/decode_test.go index 8e98b3e8..fc766250 100644 --- a/internal/eventtransformer/decode_test.go +++ b/internal/eventtransformer/decode_test.go @@ -123,6 +123,18 @@ func TestExtractAssistantReplyContent(t *testing.T) { want: "", wantErr: false, }, + { + name: "responses api output text", + resp: `{"output":[{"type":"message","content":[{"type":"output_text","text":"Hello from responses"}]}]}`, + want: "Hello from responses", + wantErr: false, + }, + { + name: "responses api nested text value", + resp: `{"output":[{"type":"message","content":[{"type":"output_text","text":{"value":"Hello nested"}}]}]}`, + want: "Hello nested", + wantErr: false, + }, } for _, tt := range tests { diff --git a/internal/eventtransformer/openai.go b/internal/eventtransformer/openai.go index 7d28d0c3..d3f92824 100644 --- a/internal/eventtransformer/openai.go +++ b/internal/eventtransformer/openai.go @@ -30,23 +30,57 @@ func MergeOpenAIStreamingChunks(body string) (map[string]any, error) { TotalTokens int `json:"total_tokens"` } var ( - id string - object string - created int - model string - content strings.Builder - finish string - usage Usage + id string + object string + created int + model string + content strings.Builder + finish string + usage Usage + responseState map[string]any + completedResponse map[string]any ) lines := strings.Split(body, "\n") + currentEvent := "" for _, line := range lines { line = strings.TrimSpace(line) + if strings.HasPrefix(line, "event: ") { + currentEvent = strings.TrimSpace(strings.TrimPrefix(line, "event: ")) + continue + } if !strings.HasPrefix(line, "data: ") || line == "data: [DONE]" { continue } data := strings.TrimPrefix(line, "data: ") var chunk map[string]any if err := json.Unmarshal([]byte(data), &chunk); err != nil { + currentEvent = "" + continue + } + eventType := currentEvent + if eventType == "" { + if v, ok := chunk["type"].(string); ok { + eventType = v + } + } + currentEvent = "" + + if strings.HasPrefix(eventType, "response.") { + if response, ok := chunk["response"].(map[string]any); ok { + if responseState == nil { + responseState = cloneAnyMap(response) + } else { + mergeMissingAnyMap(responseState, response) + } + if eventType == "response.completed" { + completedResponse = cloneAnyMap(response) + } + } + if eventType == "response.output_text.delta" { + if delta, ok := chunk["delta"].(string); ok && delta != "" { + content.WriteString(delta) + } + } continue } if v, ok := chunk["id"].(string); ok && id == "" { @@ -75,16 +109,37 @@ func MergeOpenAIStreamingChunks(body string) (map[string]any, error) { } // Merge usage if present (usually only in last chunk) if u, ok := chunk["usage"].(map[string]any); ok { - if v, ok := u["prompt_tokens"].(float64); ok { - usage.PromptTokens = int(v) - } - if v, ok := u["completion_tokens"].(float64); ok { - usage.CompletionTokens = int(v) + if normalized, ok := normalizeOpenAIUsageMap(u); ok { + usage.PromptTokens, _ = usageNumber(normalized, "input_tokens") + usage.CompletionTokens, _ = usageNumber(normalized, "output_tokens") + usage.TotalTokens, _ = usageNumber(normalized, "total_tokens") } - if v, ok := u["total_tokens"].(float64); ok { - usage.TotalTokens = int(v) + } + } + if responseState != nil || completedResponse != nil { + merged := cloneAnyMap(responseState) + if completedResponse != nil { + merged = cloneAnyMap(completedResponse) + mergeMissingAnyMap(merged, responseState) + } + if merged == nil { + merged = map[string]any{} + } + if _, ok := merged["output"]; !ok && content.Len() > 0 { + merged["output"] = []any{map[string]any{ + "type": "message", + "content": []any{map[string]any{ + "type": "output_text", + "text": content.String(), + }}, + }} + } + if responseUsage, ok := merged["usage"].(map[string]any); ok { + if normalized, ok := normalizeOpenAIUsageMap(responseUsage); ok { + merged["usage"] = normalized } } + return merged, nil } merged := map[string]any{ "id": id, @@ -102,15 +157,38 @@ func MergeOpenAIStreamingChunks(body string) (map[string]any, error) { } if usage.PromptTokens > 0 || usage.CompletionTokens > 0 || usage.TotalTokens > 0 { merged["usage"] = map[string]any{ - "prompt_tokens": usage.PromptTokens, - "completion_tokens": usage.CompletionTokens, - "total_tokens": usage.TotalTokens, + "input_tokens": usage.PromptTokens, + "output_tokens": usage.CompletionTokens, + "total_tokens": usage.TotalTokens, } } // Merged OpenAI chunks successfully return merged, nil } +func cloneAnyMap(source map[string]any) map[string]any { + if source == nil { + return nil + } + cloned := make(map[string]any, len(source)) + for key, value := range source { + cloned[key] = value + } + return cloned +} + +func mergeMissingAnyMap(target map[string]any, source map[string]any) { + if target == nil || source == nil { + return + } + for key, value := range source { + if _, exists := target[key]; exists { + continue + } + target[key] = value + } +} + // MergeThreadStreamingChunks parses and merges assistant thread.run streaming events. func MergeThreadStreamingChunks(body string) (map[string]any, error) { type Usage struct { @@ -236,9 +314,9 @@ func MergeThreadStreamingChunks(body string) (map[string]any, error) { } if usage.PromptTokens+usage.CompletionTokens+usage.TotalTokens > 0 { merged["usage"] = map[string]any{ - "prompt_tokens": usage.PromptTokens, - "completion_tokens": usage.CompletionTokens, - "total_tokens": usage.TotalTokens, + "input_tokens": usage.PromptTokens, + "output_tokens": usage.CompletionTokens, + "total_tokens": usage.TotalTokens, } } // Merged thread.run chunks successfully @@ -361,7 +439,9 @@ func (t *OpenAITransformer) TransformEvent(evt map[string]any) (map[string]any, comp, _ := json.Marshal(merged) evt["response_body"] = string(comp) if usage, ok := merged["usage"].(map[string]any); ok { - evt["TokenUsage"] = usage + if normalized, ok := normalizeOpenAIUsageMap(usage); ok { + evt["TokenUsage"] = normalized + } } } else { // Merge stream error - using decoded response body @@ -380,7 +460,9 @@ func (t *OpenAITransformer) TransformEvent(evt map[string]any) (map[string]any, var respObj map[string]any if err := json.Unmarshal([]byte(resp), &respObj); err == nil { if usage, ok := respObj["usage"].(map[string]any); ok { - evt["TokenUsage"] = usage + if normalized, ok := normalizeOpenAIUsageMap(usage); ok { + evt["TokenUsage"] = normalized + } delete(respObj, "usage") b, _ := json.Marshal(respObj) evt["response_body"] = string(b) @@ -411,6 +493,9 @@ func (t *OpenAITransformer) TransformEvent(evt map[string]any) (map[string]any, promptTokenSource = instr } } + if inputSource := promptInputSource(reqObj["input"]); inputSource != "" { + promptTokenSource += inputSource + } } } if promptTokenSource != "" { @@ -438,7 +523,9 @@ func (t *OpenAITransformer) TransformEvent(evt map[string]any) (map[string]any, tk, _ := CountOpenAITokensForModel(cnt, modelName) ct = tk } - evt["TokenUsage"] = map[string]int{"prompt_tokens": pt, "completion_tokens": ct, "total_tokens": pt + ct} + if pt > 0 || ct > 0 { + evt["TokenUsage"] = map[string]any{"input_tokens": pt, "output_tokens": ct, "total_tokens": pt + ct} + } } } @@ -501,6 +588,88 @@ func normalizeToCompactJSON(input string) (string, string, bool) { func isValidUTF8(s string) bool { return utf8.ValidString(s) } +func normalizeOpenAIUsageMap(usage map[string]any) (map[string]any, bool) { + inputTokens, hasInputTokens := usageNumber(usage, "input_tokens", "prompt_tokens") + outputTokens, hasOutputTokens := usageNumber(usage, "output_tokens", "completion_tokens") + totalTokens, hasTotalTokens := usageNumber(usage, "total_tokens") + + if !hasTotalTokens && (hasInputTokens || hasOutputTokens) { + totalTokens = inputTokens + outputTokens + hasTotalTokens = true + } + + if !hasInputTokens && !hasOutputTokens && !hasTotalTokens { + return nil, false + } + + normalized := map[string]any{ + "input_tokens": inputTokens, + "output_tokens": outputTokens, + "total_tokens": totalTokens, + } + if details, ok := usageDetailsMap(usage, "input_tokens_details", "prompt_tokens_details"); ok { + normalized["input_tokens_details"] = details + } + if details, ok := usageDetailsMap(usage, "output_tokens_details", "completion_tokens_details"); ok { + normalized["output_tokens_details"] = details + } + + return normalized, true +} + +func usageDetailsMap(usage map[string]any, keys ...string) (map[string]any, bool) { + for _, key := range keys { + value, ok := usage[key] + if !ok { + continue + } + if typed, ok := value.(map[string]any); ok && len(typed) > 0 { + return cloneAnyMap(typed), true + } + } + + return nil, false +} + +func usageNumber(usage map[string]any, keys ...string) (int, bool) { + for _, key := range keys { + value, ok := usage[key] + if !ok { + continue + } + + switch typed := value.(type) { + case int: + return typed, true + case int32: + return int(typed), true + case int64: + return int(typed), true + case float32: + return int(typed), true + case float64: + return int(typed), true + } + } + + return 0, false +} + +func promptInputSource(input any) string { + switch typed := input.(type) { + case string: + return typed + case nil: + return "" + default: + encodedInput, err := json.Marshal(typed) + if err != nil { + return "" + } + return string(encodedInput) + } +} + func extractAssistantReplyContent(resp string) (string, error) { var obj map[string]any if err := json.Unmarshal([]byte(resp), &obj); err != nil { @@ -515,5 +684,40 @@ func extractAssistantReplyContent(resp string) (string, error) { } } } + if outputItems, ok := obj["output"].([]any); ok { + var content strings.Builder + for _, outputItem := range outputItems { + itemMap, ok := outputItem.(map[string]any) + if !ok { + continue + } + segments, ok := itemMap["content"].([]any) + if !ok { + continue + } + for _, segment := range segments { + segmentMap, ok := segment.(map[string]any) + if !ok { + continue + } + if text, ok := segmentMap["text"].(string); ok && text != "" { + content.WriteString(text) + continue + } + if textMap, ok := segmentMap["text"].(map[string]any); ok { + if value, ok := textMap["value"].(string); ok && value != "" { + content.WriteString(value) + continue + } + } + if text, ok := segmentMap["output_text"].(string); ok && text != "" { + content.WriteString(text) + } + } + } + if content.Len() > 0 { + return content.String(), nil + } + } return "", nil } diff --git a/internal/eventtransformer/openai_test.go b/internal/eventtransformer/openai_test.go index dffd3ade..c7000bd8 100644 --- a/internal/eventtransformer/openai_test.go +++ b/internal/eventtransformer/openai_test.go @@ -43,6 +43,47 @@ func TestMergeOpenAIStreamingChunks(t *testing.T) { } } +func TestMergeOpenAIStreamingChunks_ResponsesCompletedEvent(t *testing.T) { + input := "event: response.created\n" + + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":123,\"model\":\"gpt-4.1-mini\",\"status\":\"in_progress\",\"output\":[]}}\n\n" + + "event: response.output_text.delta\n" + + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"Hallo Welt\"}\n\n" + + "event: response.completed\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"object\":\"response\",\"created_at\":123,\"status\":\"completed\",\"model\":\"gpt-4.1-mini\",\"usage\":{\"input_tokens\":40,\"output_tokens\":12,\"total_tokens\":52},\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"Hallo Welt\"}]}]}}\n\n" + + "data: [DONE]\n" + + merged, err := MergeOpenAIStreamingChunks(input) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if merged["status"] != "completed" { + t.Fatalf("expected completed status, got %v", merged["status"]) + } + if merged["model"] != "gpt-4.1-mini" { + t.Fatalf("expected model gpt-4.1-mini, got %v", merged["model"]) + } + usage, ok := merged["usage"].(map[string]any) + if !ok { + t.Fatalf("expected usage map, got %T", merged["usage"]) + } + if usage["input_tokens"] != 40 { + t.Fatalf("expected input_tokens 40, got %v", usage["input_tokens"]) + } + if usage["output_tokens"] != 12 { + t.Fatalf("expected output_tokens 12, got %v", usage["output_tokens"]) + } + output, ok := merged["output"].([]any) + if !ok || len(output) == 0 { + t.Fatalf("expected response output, got %T", merged["output"]) + } + message := output[0].(map[string]any) + content := message["content"].([]any) + segment := content[0].(map[string]any) + if segment["text"] != "Hallo Welt" { + t.Fatalf("expected merged response text, got %v", segment["text"]) + } +} + func TestMergeThreadStreamingChunks(t *testing.T) { input := "event: thread.run.created\ndata: {\"id\":\"run1\",\"assistant_id\":\"asst1\",\"thread_id\":\"th1\",\"status\":\"queued\",\"created_at\":123}\nevent: thread.message.delta\ndata: {\"delta\":{\"content\":[{\"type\":\"text\",\"text\":{\"value\":\"Hello\"}}]}}\nevent: thread.message.completed\ndata: {\"content\":[{\"type\":\"text\",\"text\":{\"value\":\"Final\"}}]}\n" merged, err := MergeThreadStreamingChunks(input) @@ -316,6 +357,72 @@ func TestOpenAITransformer_TransformEvent_ErrorsAndEdgeCases(t *testing.T) { } }, }, + { + name: "responses api usage normalization", + input: map[string]interface{}{ + "Method": "POST", + "ResponseHeaders": map[string]interface{}{"Content-Type": "application/json"}, + "ResponseBody": base64.StdEncoding.EncodeToString([]byte(`{"id":"resp_123","model":"gpt-4.1-mini","usage":{"input_tokens":51,"output_tokens":19,"total_tokens":70},"output":[{"type":"message","content":[{"type":"output_text","text":"Hello from responses"}]}]}`)), + "RequestBody": base64.StdEncoding.EncodeToString([]byte(`{"model":"gpt-4.1-mini","input":"hello"}`)), + }, + check: func(t *testing.T, out map[string]interface{}, err error) { + if err != nil { + t.Errorf("unexpected error: %v", err) + } + switch usage := out["token_usage"].(type) { + case map[string]int: + if usage["input_tokens"] != 51 { + t.Fatalf("expected input_tokens 51, got %v", usage["input_tokens"]) + } + if usage["output_tokens"] != 19 { + t.Fatalf("expected output_tokens 19, got %v", usage["output_tokens"]) + } + if usage["total_tokens"] != 70 { + t.Fatalf("expected total_tokens 70, got %v", usage["total_tokens"]) + } + case map[string]any: + if usage["input_tokens"] != 51 { + t.Fatalf("expected input_tokens 51, got %v", usage["input_tokens"]) + } + if usage["output_tokens"] != 19 { + t.Fatalf("expected output_tokens 19, got %v", usage["output_tokens"]) + } + if usage["total_tokens"] != 70 { + t.Fatalf("expected total_tokens 70, got %v", usage["total_tokens"]) + } + default: + t.Fatalf("expected token_usage map, got %T", out["token_usage"]) + } + }, + }, + { + name: "responses api fallback counts input prompt", + input: map[string]interface{}{ + "Method": "POST", + "ResponseHeaders": map[string]interface{}{"Content-Type": "application/json"}, + "ResponseBody": base64.StdEncoding.EncodeToString([]byte(`{"id":"resp_456","model":"gpt-4.1-mini","status":"in_progress","output":[]}`)), + "RequestBody": base64.StdEncoding.EncodeToString([]byte(`{"model":"gpt-4.1-mini","input":"hello from responses"}`)), + }, + check: func(t *testing.T, out map[string]interface{}, err error) { + if err != nil { + t.Errorf("unexpected error: %v", err) + } + usage, ok := out["token_usage"].(map[string]any) + if !ok { + t.Fatalf("expected token_usage map[string]any, got %T", out["token_usage"]) + } + inputTokens, ok := usage["input_tokens"].(int) + if !ok { + t.Fatalf("expected input_tokens int, got %T", usage["input_tokens"]) + } + if inputTokens <= 0 { + t.Fatalf("expected input_tokens > 0, got %v", usage["input_tokens"]) + } + if usage["output_tokens"] != 0 { + t.Fatalf("expected output_tokens 0 for in_progress response, got %v", usage["output_tokens"]) + } + }, + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -374,6 +481,49 @@ func TestMergeOpenAIStreamingChunks_EdgeCases(t *testing.T) { } }, }, + { + name: "responses usage keys are normalized", + input: "data: {\"usage\":{\"input_tokens\":3,\"output_tokens\":4,\"total_tokens\":7},\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}", + check: func(t *testing.T, merged map[string]any, err error) { + if err != nil { + t.Errorf("unexpected error: %v", err) + } + usage, ok := merged["usage"].(map[string]any) + if !ok { + t.Fatalf("expected usage map, got %T", merged["usage"]) + } + if usage["input_tokens"] != 3.0 && usage["input_tokens"] != 3 { + t.Fatalf("expected input_tokens 3, got %v", usage["input_tokens"]) + } + if usage["output_tokens"] != 4.0 && usage["output_tokens"] != 4 { + t.Fatalf("expected output_tokens 4, got %v", usage["output_tokens"]) + } + }, + }, + { + name: "responses completed event is authoritative", + input: "event: response.created\n" + + "data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_1\",\"status\":\"in_progress\",\"model\":\"gpt-4.1-mini\"}}\n\n" + + "event: response.output_text.delta\n" + + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n\n" + + "event: response.completed\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"status\":\"completed\",\"model\":\"gpt-4.1-mini\",\"usage\":{\"input_tokens\":3,\"output_tokens\":4,\"total_tokens\":7},\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"hi\"}]}]}}", + check: func(t *testing.T, merged map[string]any, err error) { + if err != nil { + t.Errorf("unexpected error: %v", err) + } + if merged["status"] != "completed" { + t.Fatalf("expected completed status, got %v", merged["status"]) + } + usage, ok := merged["usage"].(map[string]any) + if !ok { + t.Fatalf("expected usage map, got %T", merged["usage"]) + } + if usage["output_tokens"] != 4 { + t.Fatalf("expected output_tokens 4, got %v", usage["output_tokens"]) + } + }, + }, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { @@ -463,8 +613,8 @@ func TestMergeThreadStreamingChunks_MoreBranches(t *testing.T) { input: "event: thread.run.step.completed\ndata: {\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}", check: func(t *testing.T, merged map[string]any, err error) { usage, ok := merged["usage"].(map[string]any) - if !ok || usage["prompt_tokens"] != 1 { - t.Errorf("expected usage to be set with prompt_tokens=1, got %v", usage["prompt_tokens"]) + if !ok || usage["input_tokens"] != 1 { + t.Errorf("expected usage to be set with input_tokens=1, got %v", usage["input_tokens"]) } }, }, diff --git a/internal/logging/logger.go b/internal/logging/logger.go index f3e8e509..723b3937 100644 --- a/internal/logging/logger.go +++ b/internal/logging/logger.go @@ -2,14 +2,76 @@ package logging import ( "context" + "errors" "os" "strings" + "syscall" "github.com/sofatutor/llm-proxy/internal/obfuscate" "go.uber.org/zap" "go.uber.org/zap/zapcore" ) +type safeWriteSyncer struct { + zapcore.WriteSyncer + ignoreSyncErrors bool +} + +func (s safeWriteSyncer) Sync() error { + err := s.WriteSyncer.Sync() + if err == nil { + return nil + } + if s.ignoreSyncErrors && isIgnorableSyncError(err) { + return nil + } + return err +} + +func isIgnorableSyncError(err error) bool { + if errors.Is(err, syscall.EINVAL) { + return true + } + + return strings.Contains(strings.ToLower(err.Error()), "inappropriate ioctl for device") +} + +func shouldIgnoreSyncErrors(file *os.File, filePath string) bool { + if file == nil { + return false + } + + if filePath == "" || filePath == "/dev/stdout" || filePath == "/dev/stderr" { + return true + } + + info, err := file.Stat() + if err != nil { + return false + } + + return info.Mode()&os.ModeType != 0 +} + +func newWriteSyncer(filePath string) (zapcore.WriteSyncer, error) { + if filePath == "" { + return safeWriteSyncer{ + WriteSyncer: zapcore.AddSync(os.Stdout), + ignoreSyncErrors: true, + }, nil + } + + f, err := os.OpenFile(filePath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644) + if err != nil { + return nil, err + } + + return safeWriteSyncer{ + WriteSyncer: f, + ignoreSyncErrors: shouldIgnoreSyncErrors(f, filePath), + }, nil +} + // NewLogger creates a zap.Logger with the specified level, format, and optional file output. // level can be debug, info, warn, or error. format can be json or console. // If filePath is empty, logs are written to stdout. @@ -47,13 +109,9 @@ func NewLogger(level, format, filePath string) (*zap.Logger, error) { encoder = zapcore.NewJSONEncoder(encCfg) } - var ws = zapcore.AddSync(os.Stdout) - if filePath != "" { - f, err := os.OpenFile(filePath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0644) - if err != nil { - return nil, err - } - ws = f + ws, err := newWriteSyncer(filePath) + if err != nil { + return nil, err } core := zapcore.NewCore(encoder, ws, lvl) diff --git a/internal/logging/logger_test.go b/internal/logging/logger_test.go index 5b8b80a9..5ffc4613 100644 --- a/internal/logging/logger_test.go +++ b/internal/logging/logger_test.go @@ -2,8 +2,10 @@ package logging import ( "context" + "errors" "os" "path/filepath" + "syscall" "testing" "github.com/stretchr/testify/assert" @@ -13,6 +15,18 @@ import ( "go.uber.org/zap/zaptest/observer" ) +type stubWriteSyncer struct { + syncErr error +} + +func (s stubWriteSyncer) Write(p []byte) (int, error) { + return len(p), nil +} + +func (s stubWriteSyncer) Sync() error { + return s.syncErr +} + func TestNewLogger_FileOutput(t *testing.T) { tmpDir := t.TempDir() logFile := filepath.Join(tmpDir, "test.log") @@ -33,6 +47,34 @@ func TestNewLogger_StdoutOutput(t *testing.T) { assert.NotNil(t, logger) } +func TestSafeWriteSyncer_IgnoresStdoutLikeEINVAL(t *testing.T) { + syncer := safeWriteSyncer{ + WriteSyncer: stubWriteSyncer{syncErr: &os.PathError{Op: "sync", Path: "/dev/stdout", Err: syscall.EINVAL}}, + ignoreSyncErrors: true, + } + + require.NoError(t, syncer.Sync()) +} + +func TestSafeWriteSyncer_IgnoresIoctlDeviceErrorsWhenConfigured(t *testing.T) { + syncer := safeWriteSyncer{ + WriteSyncer: stubWriteSyncer{syncErr: errors.New("sync stdout: inappropriate ioctl for device")}, + ignoreSyncErrors: true, + } + + require.NoError(t, syncer.Sync()) +} + +func TestSafeWriteSyncer_PreservesRealSyncErrors(t *testing.T) { + wantErr := errors.New("disk full") + syncer := safeWriteSyncer{ + WriteSyncer: stubWriteSyncer{syncErr: wantErr}, + ignoreSyncErrors: true, + } + + require.ErrorIs(t, syncer.Sync(), wantErr) +} + func TestNewLogger_AllLevels(t *testing.T) { tests := []struct { level string diff --git a/internal/middleware/instrumentation.go b/internal/middleware/instrumentation.go index 4fbd68a6..230305d6 100644 --- a/internal/middleware/instrumentation.go +++ b/internal/middleware/instrumentation.go @@ -23,6 +23,8 @@ type Middleware func(http.Handler) http.Handler type ObservabilityConfig struct { Enabled bool EventBus eventbus.EventBus + // EventEnricher adds request-scoped metadata to the event before publication. + EventEnricher func(*http.Request, *eventbus.Event) // MaxRequestBodyBytes limits request body capture for observability events. 0 means "use default". MaxRequestBodyBytes int64 // MaxResponseBodyBytes limits response body capture for observability events. 0 means "use default". @@ -119,6 +121,10 @@ func (m *ObservabilityMiddleware) Middleware() Middleware { RequestBody: reqBody, } + if m.cfg.EventEnricher != nil { + m.cfg.EventEnricher(r, &evt) + } + // Publish is non-blocking; avoid spawning a goroutine per request. // Any heavy transformations (e.g., OpenAI metadata extraction) should happen in downstream consumers. m.cfg.EventBus.Publish(context.Background(), evt) diff --git a/internal/middleware/instrumentation_test.go b/internal/middleware/instrumentation_test.go index 8fe8e91c..9c45aebb 100644 --- a/internal/middleware/instrumentation_test.go +++ b/internal/middleware/instrumentation_test.go @@ -13,6 +13,7 @@ import ( "github.com/sofatutor/llm-proxy/internal/eventbus" "github.com/sofatutor/llm-proxy/internal/logging" + "github.com/sofatutor/llm-proxy/internal/token" "github.com/stretchr/testify/require" ) @@ -368,6 +369,41 @@ func TestObservabilityMiddleware_RequestIDResolution(t *testing.T) { } } +func TestObservabilityMiddleware_EventEnricherAddsTokenMetadata(t *testing.T) { + bus := eventbus.NewInMemoryEventBus(10) + mw := NewObservabilityMiddleware(ObservabilityConfig{ + Enabled: true, + EventBus: bus, + EventEnricher: func(r *http.Request, evt *eventbus.Event) { + evt.ProjectID = "project-1" + evt.TokenID = "token-uuid-1" + evt.TokenMetadata = token.CloneMetadata(map[string]string{"feature": "sofabuddy", "user_id": "42"}) + }, + }, nil) + + handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + }) + + wrapped := mw.Middleware()(handler) + req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader([]byte(`{"input":"hello"}`))) + req.Header.Set("X-Request-ID", "req-meta") + rr := httptest.NewRecorder() + ch := bus.Subscribe() + + wrapped.ServeHTTP(rr, req) + + select { + case evt := <-ch: + require.Equal(t, "project-1", evt.ProjectID) + require.Equal(t, "token-uuid-1", evt.TokenID) + require.Equal(t, map[string]string{"feature": "sofabuddy", "user_id": "42"}, evt.TokenMetadata) + case <-time.After(time.Second): + t.Fatal("event not received") + } +} + func TestObservabilityMiddleware_RequestBodyReadError(t *testing.T) { bus := eventbus.NewInMemoryEventBus(10) mw := NewObservabilityMiddleware(ObservabilityConfig{Enabled: true, EventBus: bus}, nil) diff --git a/internal/proxy/cache_helpers.go b/internal/proxy/cache_helpers.go index 2bdb172d..a57f4e91 100644 --- a/internal/proxy/cache_helpers.go +++ b/internal/proxy/cache_helpers.go @@ -11,6 +11,10 @@ import ( "time" ) +const maxCacheNamespaceLength = 128 + +var cacheNamespaceHeader = textproto.CanonicalMIMEHeaderKey("X-LLM-Proxy-Cache-Namespace") + // generateCacheKey builds a cache key from method, path, query and an optional // ordered list of header names to include. When headersToInclude is nil or empty, // no header values are incorporated. For methods carrying a body (POST/PUT/PATCH) @@ -62,6 +66,11 @@ func generateCacheKey(r *http.Request, headersToInclude []string) string { final := strings.Builder{} final.WriteString(varyKey) + if namespace := cacheNamespaceFromRequest(r); namespace != "" { + final.WriteString("|ns=") + final.WriteString(namespace) + } + // For methods with body, include X-Body-Hash when present (computed in proxy) if r.Header.Get("X-Body-Hash") != "" { final.WriteString("|body=") @@ -89,6 +98,38 @@ func generateCacheKey(r *http.Request, headersToInclude []string) string { return final.String() } +func cacheNamespaceFromRequest(r *http.Request) string { + if r == nil { + return "" + } + + raw := strings.TrimSpace(r.Header.Get(cacheNamespaceHeader)) + if raw == "" { + return "" + } + + if len(raw) > maxCacheNamespaceLength { + raw = raw[:maxCacheNamespaceLength] + } + + var builder strings.Builder + builder.Grow(len(raw)) + for _, char := range raw { + switch { + case char >= 'a' && char <= 'z': + builder.WriteRune(char) + case char >= 'A' && char <= 'Z': + builder.WriteRune(char) + case char >= '0' && char <= '9': + builder.WriteRune(char) + case char == '-', char == '_', char == ':': + builder.WriteRune(char) + } + } + + return builder.String() +} + func CacheKeyFromRequest(r *http.Request) string { // Conservative subset until per-response Vary handling is applied by caller return generateCacheKey(r, []string{"Accept", "Accept-Encoding", "Accept-Language"}) diff --git a/internal/proxy/cache_helpers_test.go b/internal/proxy/cache_helpers_test.go index eec2a8da..ea66866e 100644 --- a/internal/proxy/cache_helpers_test.go +++ b/internal/proxy/cache_helpers_test.go @@ -87,6 +87,20 @@ func TestCacheKeyFromRequest_VaryAndBodyHash(t *testing.T) { if CacheKeyFromRequest(r5) == CacheKeyFromRequest(r6) { t.Fatalf("expected cache-forced TTL to influence key for POST") } + + // Request cache namespace should influence key when provided + r7 := &http.Request{Method: http.MethodPost, Header: http.Header{cacheNamespaceHeader: {"search_v2_bucket_9_13"}}, URL: &url.URL{Path: "/v1/x"}} + r8 := &http.Request{Method: http.MethodPost, Header: http.Header{cacheNamespaceHeader: {"search_v3_bucket_9_13"}}, URL: &url.URL{Path: "/v1/x"}} + if CacheKeyFromRequest(r7) == CacheKeyFromRequest(r8) { + t.Fatalf("expected cache namespace to influence key") + } + + // Invalid namespace characters should be stripped, producing the same key as sanitized form + r9 := &http.Request{Method: http.MethodPost, Header: http.Header{cacheNamespaceHeader: {"search!@#_v2"}}, URL: &url.URL{Path: "/v1/x"}} + r10 := &http.Request{Method: http.MethodPost, Header: http.Header{cacheNamespaceHeader: {"search_v2"}}, URL: &url.URL{Path: "/v1/x"}} + if CacheKeyFromRequest(r9) != CacheKeyFromRequest(r10) { + t.Fatalf("expected sanitized cache namespace to match") + } } func TestIsResponseCacheable_Policy(t *testing.T) { diff --git a/internal/proxy/cache_hit_fastpath_test.go b/internal/proxy/cache_hit_fastpath_test.go index a70fe1d3..909bd592 100644 --- a/internal/proxy/cache_hit_fastpath_test.go +++ b/internal/proxy/cache_hit_fastpath_test.go @@ -10,14 +10,15 @@ import ( "testing" "time" + "github.com/sofatutor/llm-proxy/internal/token" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/require" "go.uber.org/zap" ) func TestProxy_CacheHitAvoidsTrackingAndAPIKeyLookup(t *testing.T) { - validator := &MockTokenValidator{} - validator.On("ValidateToken", mock.Anything, "valid-token").Return("project-1", nil).Once() + validator := &MockTokenDataValidator{} + validator.On("ValidateTokenData", mock.Anything, "valid-token").Return(token.TokenData{ProjectID: "project-1", Token: "valid-token", IsActive: true}, nil).Once() // Intentionally do NOT set expectations for ValidateTokenWithTracking. // If it is called, the mock will fail. @@ -64,13 +65,14 @@ func TestProxy_CacheHitAvoidsTrackingAndAPIKeyLookup(t *testing.T) { require.Equal(t, "llm-proxy; hit", rr.Header().Get("Cache-Status")) require.Equal(t, "hit", rr.Header().Get("X-PROXY-CACHE")) require.NotEmpty(t, rr.Header().Get("X-PROXY-CACHE-KEY")) + validator.AssertExpectations(t) } func TestProxy_POSTCacheHitAvoidsTracking(t *testing.T) { - validator := &MockTokenValidator{} + validator := &MockTokenDataValidator{} // Only expect ValidateToken (no tracking) to be called during cache hit. // The cache is manually pre-populated, so no actual upstream request occurs. - validator.On("ValidateToken", mock.Anything, "valid-token").Return("project-1", nil).Once() + validator.On("ValidateTokenData", mock.Anything, "valid-token").Return(token.TokenData{ProjectID: "project-1", Token: "valid-token", IsActive: true}, nil).Once() // Intentionally do NOT set expectations for ValidateTokenWithTracking on cache hit. // If it is called during the cache hit, the mock will fail. diff --git a/internal/proxy/interfaces.go b/internal/proxy/interfaces.go index 52dea358..014bfe46 100644 --- a/internal/proxy/interfaces.go +++ b/internal/proxy/interfaces.go @@ -7,6 +7,7 @@ import ( "time" "github.com/sofatutor/llm-proxy/internal/audit" + tokenpkg "github.com/sofatutor/llm-proxy/internal/token" ) // TokenValidator defines the interface for token validation @@ -18,6 +19,12 @@ type TokenValidator interface { ValidateTokenWithTracking(ctx context.Context, token string) (string, error) } +// TokenDataValidator exposes token metadata from the validation step without a second store lookup. +type TokenDataValidator interface { + ValidateTokenData(ctx context.Context, token string) (tokenpkg.TokenData, error) + ValidateTokenDataWithTracking(ctx context.Context, token string) (tokenpkg.TokenData, error) +} + // ProjectStore defines the interface for retrieving and managing project information // (extended for management API) type ProjectStore interface { @@ -164,6 +171,12 @@ const ( ctxKeyProjectID contextKey = "project_id" // ctxKeyTokenID is the context key for the token ID (used for cache stats) ctxKeyTokenID contextKey = "token_id" + // ctxKeyTokenRecordID is the context key for the stored token UUID + ctxKeyTokenRecordID contextKey = "token_record_id" + // ctxKeyTokenMetadata carries non-secret token metadata for observability + ctxKeyTokenMetadata contextKey = "token_metadata" + // ctxKeyObservabilityEventData carries request-scoped metadata shared with observability middleware. + ctxKeyObservabilityEventData contextKey = "observability_event_data" // ctxKeyLogger is the context key for a request-scoped logger ctxKeyLogger contextKey = "logger" // ctxKeyOriginalPath stores the original request path before proxy rewriting diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index b4d7da65..fd7d7e7d 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -23,6 +23,7 @@ import ( "github.com/google/uuid" "github.com/redis/go-redis/v9" + "github.com/sofatutor/llm-proxy/internal/eventbus" "github.com/sofatutor/llm-proxy/internal/logging" "github.com/sofatutor/llm-proxy/internal/middleware" "github.com/sofatutor/llm-proxy/internal/token" @@ -288,6 +289,7 @@ func NewTransparentProxyWithLoggerAndObservability(config ProxyConfig, validator if err != nil { return nil, err } + obsCfg.EventEnricher = p.enrichObservabilityEvent(obsCfg.EventEnricher) p.obsMiddleware = middleware.NewObservabilityMiddleware(obsCfg, logger) return p, nil } @@ -361,6 +363,7 @@ func (p *TransparentProxy) processRequestHeaders(req *http.Request) { "CF-IPCountry", // Cloudflare headers "X-Client-IP", // Other proxies "X-Original-Forwarded-For", // Chain of proxies + cacheNamespaceHeader, // Internal cache-key namespace hint; never forward upstream } // Remove headers that shouldn't be passed to the upstream @@ -807,14 +810,24 @@ func (p *TransparentProxy) Handler() http.Handler { return } var ( + tokenData token.TokenData projectID string err error ) - if preCacheOK { - // Avoid per-request DB updates on true cache hits. - projectID, err = p.tokenValidator.ValidateToken(r.Context(), tokenStr) + if validator, ok := p.tokenValidator.(TokenDataValidator); ok { + if preCacheOK { + // Avoid per-request DB updates on true cache hits. + tokenData, err = validator.ValidateTokenData(r.Context(), tokenStr) + } else { + tokenData, err = validator.ValidateTokenDataWithTracking(r.Context(), tokenStr) + } + projectID = tokenData.ProjectID } else { - projectID, err = p.tokenValidator.ValidateTokenWithTracking(r.Context(), tokenStr) + if preCacheOK { + projectID, err = p.tokenValidator.ValidateToken(r.Context(), tokenStr) + } else { + projectID, err = p.tokenValidator.ValidateTokenWithTracking(r.Context(), tokenStr) + } } if err != nil { p.handleValidationError(w, r, err) @@ -822,6 +835,17 @@ func (p *TransparentProxy) Handler() http.Handler { } ctx = context.WithValue(r.Context(), ctxKeyProjectID, projectID) ctx = context.WithValue(ctx, ctxKeyTokenID, tokenStr) + if tokenData.ID != "" { + ctx = context.WithValue(ctx, ctxKeyTokenRecordID, tokenData.ID) + } + if len(tokenData.Metadata) > 0 { + ctx = context.WithValue(ctx, ctxKeyTokenMetadata, token.CloneMetadata(tokenData.Metadata)) + } + if eventData, ok := ctx.Value(ctxKeyObservabilityEventData).(*observabilityEventData); ok && eventData != nil { + eventData.ProjectID = projectID + eventData.TokenID = tokenData.ID + eventData.TokenMetadata = token.CloneMetadata(tokenData.Metadata) + } r = r.WithContext(ctx) // Defer upstream API key lookup until we actually need to proxy upstream. // This keeps cache-hit latency low under concurrency. @@ -995,11 +1019,11 @@ func (p *TransparentProxy) Handler() http.Handler { }) var handler http.Handler = baseHandler - handler = p.ValidateRequestMiddleware()(handler) - if p.obsMiddleware != nil { handler = p.obsMiddleware.Middleware()(handler) + handler = p.observabilityEventContextMiddleware()(handler) } + handler = p.ValidateRequestMiddleware()(handler) handler = CircuitBreakerMiddleware(5, 30*time.Second, func(status int) bool { return status == http.StatusBadGateway || status == http.StatusServiceUnavailable || status == http.StatusGatewayTimeout })(handler) @@ -1007,6 +1031,56 @@ func (p *TransparentProxy) Handler() http.Handler { return handler } +type observabilityEventData struct { + ProjectID string + TokenID string + TokenMetadata map[string]string +} + +func (p *TransparentProxy) observabilityEventContextMiddleware() Middleware { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if _, ok := r.Context().Value(ctxKeyObservabilityEventData).(*observabilityEventData); ok { + next.ServeHTTP(w, r) + return + } + + eventData := &observabilityEventData{} + next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), ctxKeyObservabilityEventData, eventData))) + }) + } +} + +func (p *TransparentProxy) enrichObservabilityEvent(existing func(*http.Request, *eventbus.Event)) func(*http.Request, *eventbus.Event) { + return func(r *http.Request, evt *eventbus.Event) { + if existing != nil { + existing(r, evt) + } + + if eventData, ok := r.Context().Value(ctxKeyObservabilityEventData).(*observabilityEventData); ok && eventData != nil { + if eventData.ProjectID != "" { + evt.ProjectID = eventData.ProjectID + } + if eventData.TokenID != "" { + evt.TokenID = eventData.TokenID + } + if len(eventData.TokenMetadata) > 0 { + evt.TokenMetadata = token.CloneMetadata(eventData.TokenMetadata) + } + } + + if projectID, ok := r.Context().Value(ctxKeyProjectID).(string); ok { + evt.ProjectID = projectID + } + if tokenRecordID, ok := r.Context().Value(ctxKeyTokenRecordID).(string); ok { + evt.TokenID = tokenRecordID + } + if tokenMetadata, ok := r.Context().Value(ctxKeyTokenMetadata).(map[string]string); ok { + evt.TokenMetadata = token.CloneMetadata(tokenMetadata) + } + } +} + type timingResponseWriter struct { http.ResponseWriter firstByteOnce sync.Once diff --git a/internal/proxy/proxy_middleware_test.go b/internal/proxy/proxy_middleware_test.go index b9b5e051..a6094552 100644 --- a/internal/proxy/proxy_middleware_test.go +++ b/internal/proxy/proxy_middleware_test.go @@ -71,6 +71,7 @@ func TestProcessRequestHeaders(t *testing.T) { req.Header.Set("X-Forwarded-For", "192.168.1.1") req.Header.Set("X-Real-IP", "192.168.1.2") req.Header.Set("CF-Connecting-IP", "192.168.1.3") + req.Header.Set(cacheNamespaceHeader, "search_v2_bucket_9_13") req.RemoteAddr = "192.168.1.4:12345" // Create a streaming request @@ -86,6 +87,7 @@ func TestProcessRequestHeaders(t *testing.T) { // Verify prohibited headers were removed assert.Empty(t, req.Header.Get("X-Real-IP")) assert.Empty(t, req.Header.Get("CF-Connecting-IP")) + assert.Empty(t, req.Header.Get(cacheNamespaceHeader)) // Verify X-Forwarded-For was set to client IP assert.Equal(t, "192.168.1.4", req.Header.Get("X-Forwarded-For")) diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index 55b0f45f..8fadacad 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -16,6 +16,7 @@ import ( "errors" "github.com/sofatutor/llm-proxy/internal/audit" + "github.com/sofatutor/llm-proxy/internal/eventbus" "github.com/sofatutor/llm-proxy/internal/middleware" "github.com/sofatutor/llm-proxy/internal/token" "github.com/stretchr/testify/assert" @@ -39,6 +40,30 @@ func (m *MockTokenValidator) ValidateTokenWithTracking(ctx context.Context, toke return args.String(0), args.Error(1) } +type MockTokenDataValidator struct { + mock.Mock +} + +func (m *MockTokenDataValidator) ValidateToken(ctx context.Context, tokenID string) (string, error) { + args := m.Called(ctx, tokenID) + return args.String(0), args.Error(1) +} + +func (m *MockTokenDataValidator) ValidateTokenWithTracking(ctx context.Context, tokenID string) (string, error) { + args := m.Called(ctx, tokenID) + return args.String(0), args.Error(1) +} + +func (m *MockTokenDataValidator) ValidateTokenData(ctx context.Context, tokenID string) (token.TokenData, error) { + args := m.Called(ctx, tokenID) + return args.Get(0).(token.TokenData), args.Error(1) +} + +func (m *MockTokenDataValidator) ValidateTokenDataWithTracking(ctx context.Context, tokenID string) (token.TokenData, error) { + args := m.Called(ctx, tokenID) + return args.Get(0).(token.TokenData), args.Error(1) +} + // MockProjectStore mocks the project storage interface type MockProjectStore struct { mock.Mock @@ -1770,6 +1795,58 @@ func TestNewTransparentProxyWithLoggerAndObservability(t *testing.T) { assert.NotNil(t, proxy.obsMiddleware) } +func TestTransparentProxy_ObservabilityEnrichesTokenMetadata(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"id":"resp_1","model":"gpt-4.1-mini","usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}`)) + })) + defer upstream.Close() + + validator := &MockTokenDataValidator{} + projectStore := &MockProjectStore{} + bus := eventbus.NewInMemoryEventBus(10) + obsCfg := middleware.ObservabilityConfig{Enabled: true, EventBus: bus} + + validator.On("ValidateTokenDataWithTracking", mock.Anything, "tok").Return(token.TokenData{ + ID: "token-uuid-1", + Token: "tok", + ProjectID: "project-1", + Metadata: map[string]string{"feature": "sofabuddy", "user_id": "42"}, + IsActive: true, + }, nil).Once() + projectStore.On("GetAPIKeyForProject", mock.Anything, "project-1").Return("upstream-key", nil).Once() + + p, err := NewTransparentProxyWithAudit(ProxyConfig{ + TargetBaseURL: upstream.URL, + AllowedEndpoints: []string{"/v1/responses"}, + AllowedMethods: []string{"POST"}, + }, validator, projectStore, zap.NewNop(), &TestAuditLogger{}, obsCfg) + require.NoError(t, err) + + req := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{"model":"gpt-4.1-mini","input":"hello"}`)) + req.Header.Set("Authorization", "Bearer tok") + req.Header.Set("Content-Type", "application/json") + req.Header.Set("X-Request-ID", "req-observe") + w := httptest.NewRecorder() + ch := bus.Subscribe() + + p.Handler().ServeHTTP(w, req) + require.Equal(t, http.StatusOK, w.Code) + + select { + case evt := <-ch: + require.Equal(t, "project-1", evt.ProjectID) + require.Equal(t, "token-uuid-1", evt.TokenID) + require.Equal(t, map[string]string{"feature": "sofabuddy", "user_id": "42"}, evt.TokenMetadata) + case <-time.After(time.Second): + t.Fatal("event not received") + } + + validator.AssertExpectations(t) + projectStore.AssertExpectations(t) +} + // --- Project Active Enforcement Coverage --- func TestProjectActiveEnforcement_DBErrorReturns503(t *testing.T) { // Upstream that would return 200 if reached diff --git a/internal/server/management_api_test.go b/internal/server/management_api_test.go index ac4ea740..55f5ac58 100644 --- a/internal/server/management_api_test.go +++ b/internal/server/management_api_test.go @@ -476,6 +476,9 @@ func TestHandleTokens(t *testing.T) { reqBody := map[string]interface{}{ "project_id": "project-1", "duration_minutes": 60 * 24, + "metadata": map[string]string{ + "user_id": "42", + }, } body, _ := json.Marshal(reqBody) req := httptest.NewRequest("POST", "/manage/tokens", bytes.NewReader(body)) @@ -494,6 +497,8 @@ func TestHandleTokens(t *testing.T) { assert.Contains(t, response, "id") assert.Contains(t, response, "token") assert.Contains(t, response, "expires_at") + require.NotEmpty(t, tokenStore.tokens) + assert.Equal(t, map[string]string{"user_id": "42"}, tokenStore.tokens[len(tokenStore.tokens)-1].Metadata) }) t.Run("POST_Token_InvalidRequest", func(t *testing.T) { diff --git a/internal/server/server.go b/internal/server/server.go index cf89830a..04e1d47e 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -151,13 +151,14 @@ func NewWithDatabase(cfg *config.Config, tokenStore token.TokenStore, projectSto } streamsConfig := eventbus.RedisStreamsConfig{ - StreamKey: cfg.RedisStreamKey, - ConsumerGroup: cfg.RedisConsumerGroup, - ConsumerName: consumerName, - MaxLen: cfg.RedisStreamMaxLen, - BlockTimeout: cfg.RedisStreamBlockTime, - ClaimMinIdleTime: cfg.RedisStreamClaimTime, - BatchSize: cfg.RedisStreamBatchSize, + StreamKey: cfg.RedisStreamKey, + ConsumerGroup: cfg.RedisConsumerGroup, + ConsumerName: consumerName, + PublishBufferSize: cfg.ObservabilityBufferSize, + MaxLen: cfg.RedisStreamMaxLen, + BlockTimeout: cfg.RedisStreamBlockTime, + ClaimMinIdleTime: cfg.RedisStreamClaimTime, + BatchSize: cfg.RedisStreamBatchSize, } adapter := &eventbus.RedisStreamsClientAdapter{Client: client} bus = eventbus.NewRedisStreamsEventBus(adapter, streamsConfig) @@ -1215,9 +1216,10 @@ func (s *Server) handleTokens(w http.ResponseWriter, r *http.Request) { switch r.Method { case http.MethodPost: var req struct { - ProjectID string `json:"project_id"` - DurationMinutes int `json:"duration_minutes"` - MaxRequests *int `json:"max_requests"` + ProjectID string `json:"project_id"` + DurationMinutes int `json:"duration_minutes"` + MaxRequests *int `json:"max_requests"` + Metadata map[string]string `json:"metadata"` } if err := json.NewDecoder(r.Body).Decode(&req); err != nil { s.logger.Error("invalid token create request body", zap.Error(err), zap.String("request_id", requestID)) @@ -1334,6 +1336,7 @@ func (s *Server) handleTokens(w http.ResponseWriter, r *http.Request) { ID: tokenID, Token: tokenStr, ProjectID: req.ProjectID, + Metadata: token.CloneMetadata(req.Metadata), ExpiresAt: expiresAt, IsActive: true, RequestCount: 0, diff --git a/internal/token/cache.go b/internal/token/cache.go index c10619b6..8b38228e 100644 --- a/internal/token/cache.go +++ b/internal/token/cache.go @@ -120,22 +120,42 @@ func NewCachedValidator(validator TokenValidator, options ...CacheOptions) *Cach // ValidateToken validates a token using the cache when possible func (cv *CachedValidator) ValidateToken(ctx context.Context, tokenID string) (string, error) { + tokenData, err := cv.ValidateTokenData(ctx, tokenID) + if err != nil { + return "", err + } + + return tokenData.ProjectID, nil +} + +// ValidateTokenData validates a token using the cache when possible and returns token data. +func (cv *CachedValidator) ValidateTokenData(ctx context.Context, tokenID string) (TokenData, error) { // Check cache first - projectID, found := cv.checkCache(tokenID) + tokenData, found := cv.checkCacheData(tokenID) if found { - return projectID, nil + return tokenData, nil + } + + if validator, ok := cv.validator.(interface { + ValidateTokenData(context.Context, string) (TokenData, error) + }); ok { + tokenData, err := validator.ValidateTokenData(ctx, tokenID) + if err != nil { + return TokenData{}, err + } + + cv.cacheTokenData(tokenID, tokenData) + return tokenData, nil } - // Cache miss, validate using the underlying validator projectID, err := cv.validator.ValidateToken(ctx, tokenID) if err != nil { - return "", err + return TokenData{}, err } - // Cache the successful validation cv.cacheToken(ctx, tokenID) - return projectID, nil + return TokenData{Token: tokenID, ProjectID: projectID}, nil } // ValidateTokenWithTracking validates a token and tracks usage. @@ -147,6 +167,16 @@ func (cv *CachedValidator) ValidateToken(ctx context.Context, tokenID string) (s // by using the cached token metadata (active/expiry/project) and performing a synchronous usage // increment (which is where max_requests is enforced). func (cv *CachedValidator) ValidateTokenWithTracking(ctx context.Context, tokenID string) (string, error) { + tokenData, err := cv.ValidateTokenDataWithTracking(ctx, tokenID) + if err != nil { + return "", err + } + + return tokenData.ProjectID, nil +} + +// ValidateTokenDataWithTracking validates a token, tracks usage, and returns token data. +func (cv *CachedValidator) ValidateTokenDataWithTracking(ctx context.Context, tokenID string) (TokenData, error) { cv.cacheMutex.RLock() entry, found := cv.cache[tokenID] cv.cacheMutex.RUnlock() @@ -174,7 +204,7 @@ func (cv *CachedValidator) ValidateTokenWithTracking(ctx context.Context, tokenI if getter, ok := cv.validator.(usageStatsAggregatorGetter); ok { if agg := getter.usageStatsAggregator(); agg != nil { agg.RecordTokenUsage(tokenID) - return entry.Data.ProjectID, nil + return cloneTokenData(entry.Data), nil } } } else if sv, ok := cv.validator.(*StandardValidator); ok && sv != nil && sv.store != nil { @@ -185,9 +215,9 @@ func (cv *CachedValidator) ValidateTokenWithTracking(ctx context.Context, tokenI if err == ErrTokenRateLimit || err == ErrTokenInactive || err == ErrTokenExpired { cv.invalidateCache(tokenID) } - return "", err + return TokenData{}, err } - return entry.Data.ProjectID, nil + return cloneTokenData(entry.Data), nil } } else { // Token became invalid (inactive/expired). Be defensive and drop the cache entry. @@ -200,20 +230,31 @@ func (cv *CachedValidator) ValidateTokenWithTracking(ctx context.Context, tokenI } } + if validator, ok := cv.validator.(interface { + ValidateTokenDataWithTracking(context.Context, string) (TokenData, error) + }); ok { + tokenData, err := validator.ValidateTokenDataWithTracking(ctx, tokenID) + if err != nil { + return TokenData{}, err + } + + cv.cacheTokenData(tokenID, tokenData) + return tokenData, nil + } + projectID, err := cv.validator.ValidateTokenWithTracking(ctx, tokenID) if err != nil { - return "", err + return TokenData{}, err } // Populate the cache after successful tracking so subsequent requests can avoid extra DB reads. // (Only works for StandardValidator; others are safely ignored.) cv.cacheToken(ctx, tokenID) - return projectID, nil + return TokenData{Token: tokenID, ProjectID: projectID}, nil } -// checkCache checks if a token is in the cache and still valid -func (cv *CachedValidator) checkCache(tokenID string) (string, bool) { +func (cv *CachedValidator) checkCacheData(tokenID string) (TokenData, bool) { cv.cacheMutex.RLock() entry, found := cv.cache[tokenID] cv.cacheMutex.RUnlock() @@ -223,7 +264,7 @@ func (cv *CachedValidator) checkCache(tokenID string) (string, bool) { cv.statsMutex.Lock() cv.misses++ cv.statsMutex.Unlock() - return "", false + return TokenData{}, false } // In cache but expired @@ -238,7 +279,7 @@ func (cv *CachedValidator) checkCache(tokenID string) (string, bool) { cv.evictions++ cv.statsMutex.Unlock() - return "", false + return TokenData{}, false } // In cache but token is no longer valid (inactive/expired) @@ -248,7 +289,7 @@ func (cv *CachedValidator) checkCache(tokenID string) (string, bool) { cv.misses++ cv.evictions++ cv.statsMutex.Unlock() - return "", false + return TokenData{}, false } // In cache and valid @@ -256,7 +297,13 @@ func (cv *CachedValidator) checkCache(tokenID string) (string, bool) { cv.hits++ cv.statsMutex.Unlock() - return entry.Data.ProjectID, true + return cloneTokenData(entry.Data), true +} + +func cloneTokenData(data TokenData) TokenData { + cloned := data + cloned.Metadata = CloneMetadata(data.Metadata) + return cloned } // isCacheableTokenValid determines whether a cached token can be used for authentication. @@ -287,6 +334,11 @@ func (cv *CachedValidator) cacheToken(ctx context.Context, tokenID string) { if err != nil { return } + + cv.cacheTokenData(tokenID, tokenData) +} + +func (cv *CachedValidator) cacheTokenData(tokenID string, tokenData TokenData) { if !isCacheableTokenValid(tokenData) { return } @@ -300,7 +352,7 @@ func (cv *CachedValidator) cacheToken(ctx context.Context, tokenID string) { cv.insertCounter++ cv.cache[tokenID] = CacheEntry{ - Data: tokenData, + Data: cloneTokenData(tokenData), ValidUntil: validUntil, } // Remove old heap entry if present diff --git a/internal/token/cache_test.go b/internal/token/cache_test.go index a347b55a..b8d67b68 100644 --- a/internal/token/cache_test.go +++ b/internal/token/cache_test.go @@ -332,10 +332,56 @@ func TestCachedValidator_ValidateTokenWithTracking_LimitedToken_UsesCacheAndAvoi if store.incCalls != 2 { t.Fatalf("IncrementTokenUsage calls = %d, want 2", store.incCalls) } - // First call: 2x reads (1x via validateTokenData + 1x via cacheToken population) and 1x write. + // First call: 1x read (validateTokenData) and 1x write. // Second call (cache hit): 0x reads and 1x write (only IncrementTokenUsage). - if store.getByTokenCalls != 2 { - t.Fatalf("GetTokenByToken calls = %d, want 2", store.getByTokenCalls) + if store.getByTokenCalls != 1 { + t.Fatalf("GetTokenByToken calls = %d, want 1", store.getByTokenCalls) + } +} + +func TestCachedValidator_ValidateTokenData_CachesReturnedTokenData(t *testing.T) { + ctx := context.Background() + store := newCountingStore() + validator := &StandardValidator{store: store} + cv := NewCachedValidator(validator, CacheOptions{TTL: time.Minute, MaxSize: 10, EnableCleanup: false}) + + now := time.Now() + future := now.Add(1 * time.Hour) + tok, _ := GenerateToken() + + store.tokens[tok] = TokenData{ + Token: tok, + ProjectID: "p1", + Metadata: map[string]string{"feature": "sofabuddy"}, + ExpiresAt: &future, + IsActive: true, + CreatedAt: now, + } + + first, err := cv.ValidateTokenData(ctx, tok) + if err != nil { + t.Fatalf("ValidateTokenData() first error = %v", err) + } + second, err := cv.ValidateTokenData(ctx, tok) + if err != nil { + t.Fatalf("ValidateTokenData() second error = %v", err) + } + + if first.ProjectID != "p1" || second.ProjectID != "p1" { + t.Fatalf("unexpected project IDs: first=%q second=%q", first.ProjectID, second.ProjectID) + } + if second.Metadata["feature"] != "sofabuddy" { + t.Fatalf("expected cached metadata to be preserved, got %#v", second.Metadata) + } + if store.getByTokenCalls != 1 { + t.Fatalf("GetTokenByToken calls = %d, want 1", store.getByTokenCalls) + } + if store.incCalls != 0 { + t.Fatalf("IncrementTokenUsage calls = %d, want 0", store.incCalls) + } + first.Metadata["feature"] = "changed" + if second.Metadata["feature"] != "sofabuddy" { + t.Fatalf("expected cached token metadata clone to be isolated, got %#v", second.Metadata) } } diff --git a/internal/token/manager.go b/internal/token/manager.go index ee7752b5..6d729d33 100644 --- a/internal/token/manager.go +++ b/internal/token/manager.go @@ -67,6 +67,7 @@ func (m *Manager) CreateToken(ctx context.Context, projectID string, options Tok token := TokenData{ Token: tokenStr, ProjectID: projectID, + Metadata: CloneMetadata(options.Metadata), ExpiresAt: expiresAt, IsActive: true, RequestCount: 0, @@ -231,6 +232,20 @@ type TokenOptions struct { Metadata map[string]string } +// CloneMetadata returns a shallow copy of metadata and nil for empty input. +func CloneMetadata(metadata map[string]string) map[string]string { + if len(metadata) == 0 { + return nil + } + + cloned := make(map[string]string, len(metadata)) + for key, value := range metadata { + cloned[key] = value + } + + return cloned +} + // TokenStats contains statistics about token usage type TokenStats struct { Token string diff --git a/internal/token/validate.go b/internal/token/validate.go index 60383bd6..b1cdd174 100644 --- a/internal/token/validate.go +++ b/internal/token/validate.go @@ -51,9 +51,10 @@ type TokenStore interface { // TokenData represents the data associated with a token type TokenData struct { - ID string // The token ID (UUID) - used for management operations - Token string // The token string (sk-...) - used for authentication - ProjectID string // The associated project ID + ID string // The token ID (UUID) - used for management operations + Token string // The token string (sk-...) - used for authentication + ProjectID string // The associated project ID + Metadata map[string]string ExpiresAt *time.Time // When the token expires (nil for no expiration) IsActive bool // Whether the token is active DeactivatedAt *time.Time // When the token was deactivated (nil if not deactivated) @@ -144,50 +145,69 @@ func (v *StandardValidator) validateTokenData(ctx context.Context, tokenString s // ValidateToken validates a token without incrementing usage func (v *StandardValidator) ValidateToken(ctx context.Context, tokenString string) (string, error) { - tokenData, err := v.validateTokenData(ctx, tokenString) + tokenData, err := v.ValidateTokenData(ctx, tokenString) if err != nil { return "", err } + return tokenData.ProjectID, nil +} + +// ValidateTokenData validates a token without incrementing usage and returns the token data. +func (v *StandardValidator) ValidateTokenData(ctx context.Context, tokenString string) (TokenData, error) { + tokenData, err := v.validateTokenData(ctx, tokenString) + if err != nil { + return TokenData{}, err + } + // Check if the token has reached its rate limit if tokenData.IsRateLimited() { - return "", ErrTokenRateLimit + return TokenData{}, ErrTokenRateLimit } - // Token is valid, return the project ID - return tokenData.ProjectID, nil + return tokenData, nil } // ValidateTokenWithTracking validates a token and increments its usage count func (v *StandardValidator) ValidateTokenWithTracking(ctx context.Context, tokenString string) (string, error) { + tokenData, err := v.ValidateTokenDataWithTracking(ctx, tokenString) + if err != nil { + return "", err + } + + return tokenData.ProjectID, nil +} + +// ValidateTokenDataWithTracking validates a token, increments usage, and returns the token data. +func (v *StandardValidator) ValidateTokenDataWithTracking(ctx context.Context, tokenString string) (TokenData, error) { // Validate token and load token data once. tokenData, err := v.validateTokenData(ctx, tokenString) if err != nil { - return "", err + return TokenData{}, err } // Enforce rate limit before tracking. if tokenData.IsRateLimited() { - return "", ErrTokenRateLimit + return TokenData{}, ErrTokenRateLimit } // For unlimited tokens, move request_count/last_used_at updates off the hot path. if tokenData.MaxRequests == nil || *tokenData.MaxRequests <= 0 { if agg := v.usageStatsAgg.Load(); agg != nil { agg.RecordTokenUsage(tokenString) - return tokenData.ProjectID, nil + return tokenData, nil } } // Limited tokens (or no async aggregator configured): do synchronous tracking. if err := v.store.IncrementTokenUsage(ctx, tokenString); err != nil { if errors.Is(err, ErrTokenRateLimit) { - return "", ErrTokenRateLimit + return TokenData{}, ErrTokenRateLimit } - return "", fmt.Errorf("failed to track token usage: %w", err) + return TokenData{}, fmt.Errorf("failed to track token usage: %w", err) } - return tokenData.ProjectID, nil + return tokenData, nil } // ValidateTokenFormat checks if a token string has the correct format diff --git a/scripts/schema.sql b/scripts/schema.sql index eaa9a5c7..6aaefea7 100644 --- a/scripts/schema.sql +++ b/scripts/schema.sql @@ -22,6 +22,7 @@ CREATE TABLE IF NOT EXISTS tokens ( id TEXT PRIMARY KEY, token TEXT NOT NULL UNIQUE, project_id TEXT NOT NULL, + metadata TEXT, expires_at DATETIME, is_active BOOLEAN NOT NULL DEFAULT 1, deactivated_at DATETIME,