diff --git a/internal/core/chat_json.go b/internal/core/chat_json.go index ad2aa3a6d..bccd4cfd5 100644 --- a/internal/core/chat_json.go +++ b/internal/core/chat_json.go @@ -4,7 +4,7 @@ import "github.com/goccy/go-json" // chatRequestFields is derived from the struct's json tags at package init so // the known-field list cannot drift from the type definition. -var chatRequestFields = jsonFieldNames(ChatRequest{}) +var chatRequestFields = jsonFieldSetOf(ChatRequest{}) // UnmarshalJSON decodes the typed fields via an alias (so new fields are // picked up automatically) and captures every other member in ExtraFields. @@ -15,7 +15,7 @@ func (r *ChatRequest) UnmarshalJSON(data []byte) error { return err } - extraFields, err := extractUnknownJSONFields(data, chatRequestFields...) + extraFields, err := extractUnknownJSONFieldsSet(data, chatRequestFields) if err != nil { return err } diff --git a/internal/core/json_fields.go b/internal/core/json_fields.go index e35884e49..d5ba3b55f 100644 --- a/internal/core/json_fields.go +++ b/internal/core/json_fields.go @@ -292,6 +292,24 @@ func (fields UnknownJSONFields) Without(keys ...string) UnknownJSONFields { return UnknownJSONFields{raw: buf.Bytes()} } +// jsonFieldSet answers "is this key a known typed field?" in O(1). Use it via +// jsonFieldSetOf for the derived struct lists (dozens of fields, checked once +// per JSON key of every decoded request); the variadic +// extractUnknownJSONFields stays optimal for the hand-listed callers with a +// handful of fields, where a linear scan beats a map lookup. +type jsonFieldSet map[string]struct{} + +// jsonFieldSetOf derives the known-field set from v's struct definition, like +// jsonFieldNames but as a lookup set. +func jsonFieldSetOf(v any) jsonFieldSet { + names := jsonFieldNames(v) + set := make(jsonFieldSet, len(names)) + for _, name := range names { + set[name] = struct{}{} + } + return set +} + // extractUnknownJSONFields captures the object's keys that are not in // knownFields, preserving their raw bytes for passthrough (Postel's Law). // @@ -301,6 +319,21 @@ func (fields UnknownJSONFields) Without(keys ...string) UnknownJSONFields { // benefit. The cheap first-byte and IsObject checks remain to reject non-object // JSON explicitly. func extractUnknownJSONFields(data []byte, knownFields ...string) (UnknownJSONFields, error) { + return extractUnknownJSONFieldsWith(data, func(key string) bool { + return slices.Contains(knownFields, key) + }) +} + +// extractUnknownJSONFieldsSet is extractUnknownJSONFields with a precomputed +// known-field set, for the struct-derived lists too large to scan per key. +func extractUnknownJSONFieldsSet(data []byte, known jsonFieldSet) (UnknownJSONFields, error) { + return extractUnknownJSONFieldsWith(data, func(key string) bool { + _, ok := known[key] + return ok + }) +} + +func extractUnknownJSONFieldsWith(data []byte, isKnown func(string) bool) (UnknownJSONFields, error) { data = bytes.TrimSpace(data) if len(data) == 0 || data[0] != '{' { return UnknownJSONFields{}, fmt.Errorf("expected JSON object") @@ -320,7 +353,7 @@ func extractUnknownJSONFields(data []byte, knownFields ...string) (UnknownJSONFi buf.WriteByte('{') wrote := false root.ForEach(func(key, value gjson.Result) bool { - if slices.Contains(knownFields, key.String()) { + if isKnown(key.String()) { return true } if wrote { diff --git a/internal/core/responses_json.go b/internal/core/responses_json.go index 37e3503a2..5df2bbef2 100644 --- a/internal/core/responses_json.go +++ b/internal/core/responses_json.go @@ -12,15 +12,15 @@ import ( // captured as an unknown extra field. ContentSchema swagger phantoms share // tags with real fields, which is harmless here (duplicates in the list). var ( - responsesRequestFields = jsonFieldNames(ResponsesRequest{}) - responsesUtilityRequestFields = jsonFieldNames(ResponseInputTokensRequest{}) - responsesOutputItemFields = jsonFieldNames(ResponsesOutputItem{}) + responsesRequestFields = jsonFieldSetOf(ResponsesRequest{}) + responsesUtilityRequestFields = jsonFieldSetOf(ResponseInputTokensRequest{}) + responsesOutputItemFields = jsonFieldSetOf(ResponsesOutputItem{}) ) // responsesExtrasAndInput finishes a responses-shaped decode: it captures // unknown members and decodes the raw input union. -func responsesExtrasAndInput(data []byte, rawInput json.RawMessage, knownFields []string) (any, UnknownJSONFields, error) { - extraFields, err := extractUnknownJSONFields(data, knownFields...) +func responsesExtrasAndInput(data []byte, rawInput json.RawMessage, knownFields jsonFieldSet) (any, UnknownJSONFields, error) { + extraFields, err := extractUnknownJSONFieldsSet(data, knownFields) if err != nil { return nil, UnknownJSONFields{}, err } @@ -328,7 +328,7 @@ func (i *ResponsesOutputItem) UnmarshalJSON(data []byte) error { if err := json.Unmarshal(data, &decoded); err != nil { return err } - extraFields, err := extractUnknownJSONFields(data, responsesOutputItemFields...) + extraFields, err := extractUnknownJSONFieldsSet(data, responsesOutputItemFields) if err != nil { return err } diff --git a/internal/server/http.go b/internal/server/http.go index 99fd66e64..2af9fa325 100644 --- a/internal/server/http.go +++ b/internal/server/http.go @@ -138,6 +138,7 @@ func New(provider core.RoutableProvider, cfg *Config) *Server { AllowOverwritingRoute: true, NotFoundHandler: handleRouteNotFound, }), + JSONSerializer: goJSONSerializer{}, }) e.Logger = slog.Default() basePath := configuredBasePath(cfg) @@ -311,19 +312,13 @@ func New(provider core.RoutableProvider, cfg *Config) *Server { } e.Use(middleware.BodyLimit(parseBodySizeLimitBytes(bodySizeLimit))) - // Request ID middleware (always active — ensures every request has a unique ID - // for usage tracking, audit logging, and response correlation) - e.Use(func(next echo.HandlerFunc) echo.HandlerFunc { - return func(c *echo.Context) error { - req, id := ensureRequestID(c.Request()) - c.SetRequest(req) - c.Response().Header().Set("X-Request-ID", id) - return next(c) - } - }) e.Use(modelInteractionWriteDeadlineMiddleware()) - // Ingress capture (before auth/audit/model validation so they can consume shared raw request state) + // Ingress capture (before auth/audit/model validation so they can consume + // shared raw request state). Also assigns the per-request ID: the snapshot + // middleware runs unconditionally and calls ensureRequestID first thing, so + // a separate request-ID middleware would just repeat that work (a second + // context wrap + request copy) on every request. userPathHeaderName := configuredUserPathHeader(cfg) handler.userPathHeaderName = userPathHeaderName e.Use(RequestSnapshotCapture(userPathHeaderName)) diff --git a/internal/server/json_serializer.go b/internal/server/json_serializer.go new file mode 100644 index 000000000..d14bdb975 --- /dev/null +++ b/internal/server/json_serializer.go @@ -0,0 +1,28 @@ +package server + +import ( + "github.com/goccy/go-json" + "github.com/labstack/echo/v5" +) + +// goJSONSerializer serializes responses with goccy/go-json instead of echo's +// reflection-based encoding/json default. Response types already marshal with +// goccy elsewhere (internal/core), so this removes the second, slower encoder +// from the hot path. Deserialization keeps echo's default: request bodies are +// decoded through the core types, not c.Bind, so there is nothing to win there +// and the default's error mapping stays intact for the few admin binds. +type goJSONSerializer struct { + fallback echo.DefaultJSONSerializer +} + +func (s goJSONSerializer) Serialize(c *echo.Context, target any, indent string) error { + enc := json.NewEncoder(c.Response()) + if indent != "" { + enc.SetIndent("", indent) + } + return enc.Encode(target) +} + +func (s goJSONSerializer) Deserialize(c *echo.Context, target any) error { + return s.fallback.Deserialize(c, target) +} diff --git a/internal/session/canonical_test.go b/internal/session/canonical_test.go new file mode 100644 index 000000000..228f8686e --- /dev/null +++ b/internal/session/canonical_test.go @@ -0,0 +1,60 @@ +package session + +import ( + "bytes" + encjson "encoding/json" + "testing" + + "github.com/tidwall/gjson" +) + +// TestCanonicalSegmentMatchesStdlib pins that goccy-based canonicalization +// produces byte-identical output to the original encoding/json implementation. +// Auto-detected session ids hash these bytes, so any divergence would silently +// re-anchor every in-flight conversation on upgrade (breaking virtual-model +// affinity pins and Pro compression epochs mid-session). +func TestCanonicalSegmentMatchesStdlib(t *testing.T) { + segments := []string{ + `"gpt-4o-mini"`, + `"with \"escapes\" and é unicode 😀"`, + `"html & specials"`, + `123`, + `1e2`, + `1.50`, + `-0.0031415926535897932384626433e4`, + `9007199254740993`, + `true`, + `null`, + `[]`, + `{}`, + `[1, "two", {"three": 3}, [4]]`, + `{"b":2,"a":1,"nested":{"z":null,"y":[1.0,"x"]}}`, + `{"role":"user","content":[{"type":"text","text":"line one\nline two\ttabbed"}]}`, + `{ "spaced" : { "keys" : [ 1 , 2 ] } }`, + } + + for _, segment := range segments { + parsed := gjson.Parse(segment) + got := canonicalSegment(parsed) + want := stdlibCanonical(t, parsed.Raw) + if !bytes.Equal(got, want) { + t.Errorf("canonicalSegment(%s) = %s, stdlib canonical = %s", segment, got, want) + } + } +} + +// stdlibCanonical is the pre-switch implementation, kept verbatim as the oracle. +func stdlibCanonical(t *testing.T, raw string) []byte { + t.Helper() + decoder := encjson.NewDecoder(bytes.NewReader([]byte(raw))) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return []byte(raw) + } + canonical, err := encjson.Marshal(value) + if err != nil { + return []byte(raw) + } + return canonical +} diff --git a/internal/session/detect.go b/internal/session/detect.go index 3ed71f415..136b9a708 100644 --- a/internal/session/detect.go +++ b/internal/session/detect.go @@ -4,11 +4,11 @@ import ( "bytes" "crypto/sha256" "encoding/hex" - "encoding/json" "errors" "io" "strings" + "github.com/goccy/go-json" "github.com/tidwall/gjson" "github.com/enterpilot/gomodel/internal/core" @@ -196,6 +196,14 @@ func rawSegment(result gjson.Result) json.RawMessage { // equivalent string escapes must not split one conversation into multiple // auto-detected sessions. UseNumber preserves number spelling/precision while // arrays retain their original order. +// +// The goccy canonical bytes are byte-identical to encoding/json's (pinned by +// TestCanonicalSegmentMatchesStdlib), so auto-detected ids are stable across +// the library switch. The trailing-data guard must decode to io.EOF rather +// than check Decoder.More: More treats a stray closing bracket ("1]", "1}") +// as end of input, which would canonicalize malformed raw segments instead of +// falling back to their exact bytes. For valid input the extra decode reads +// only the empty remainder, so it costs nothing on the hot path. func canonicalSegment(result gjson.Result) json.RawMessage { raw := rawSegment(result) if len(raw) == 0 { @@ -207,7 +215,8 @@ func canonicalSegment(result gjson.Result) json.RawMessage { if err := decoder.Decode(&value); err != nil { return raw } - if err := ensureJSONEOF(decoder); err != nil { + var trailing any + if err := decoder.Decode(&trailing); !errors.Is(err, io.EOF) { return raw } canonical, err := json.Marshal(value) @@ -216,15 +225,3 @@ func canonicalSegment(result gjson.Result) json.RawMessage { } return canonical } - -func ensureJSONEOF(decoder *json.Decoder) error { - var trailing any - err := decoder.Decode(&trailing) - if errors.Is(err, io.EOF) { - return nil - } - if err == nil { - return errors.New("multiple JSON values") - } - return err -} diff --git a/internal/session/detect_test.go b/internal/session/detect_test.go index 365d53e46..d7ec50a6e 100644 --- a/internal/session/detect_test.go +++ b/internal/session/detect_test.go @@ -218,6 +218,10 @@ func TestCanonicalSegmentFallsBackToExactRawJSON(t *testing.T) { for _, raw := range []string{ `{"unterminated":`, `{"first":1}{"second":2}`, + // Stray closing brackets: Decoder.More treats these as end of input, + // so the trailing-data guard must decode to io.EOF to catch them. + `1]`, + `1}`, } { result := gjson.Result{Type: gjson.JSON, Raw: raw} if got := string(canonicalSegment(result)); got != raw { diff --git a/tests/perf/README.md b/tests/perf/README.md index 321053a42..5ff0ae685 100644 --- a/tests/perf/README.md +++ b/tests/perf/README.md @@ -30,3 +30,21 @@ bare one and is independent of catalog size. `BenchmarkSharedStreamingObserversDefaultConfig` covers streaming observation with audit body capture disabled (the default), where the observed stream skips JSON decoding for chunks no observer wants. + +## Production shape and ablation + +`BenchmarkGatewayHotPathProductionShape` runs the routed path with the +default-deployment middleware chain fully wired: master-key auth, audit +logging (bodies + headers), usage tracking, session keeping, and a configured +rate limit. The guard enforces allocation ceilings on this case too, so +regressions in any of those subsystems fail CI even though the bare cases +cannot see them. + +The `BenchmarkAblation*` family turns exactly one subsystem off relative to +that full shape; the delta against `BenchmarkAblationFull` attributes +per-request cost to that subsystem. These are diagnostic benchmarks (run via +`make perf-bench`), not guarded. + +`TestSessionIDVisibilityByBodySize` pins that content-based session +auto-detection is independent of request body size — there is no size above +which a request quietly stops carrying a session id to downstream consumers. diff --git a/tests/perf/ablation_test.go b/tests/perf/ablation_test.go new file mode 100644 index 000000000..8b2504deb --- /dev/null +++ b/tests/perf/ablation_test.go @@ -0,0 +1,85 @@ +package perf + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/enterpilot/gomodel/internal/auditlog" + "github.com/enterpilot/gomodel/internal/server" + "github.com/enterpilot/gomodel/internal/session" + "github.com/enterpilot/gomodel/internal/usage" +) + +// ablation isolates the per-subsystem cost of the default-on middleware stack. +// Each variant turns exactly one subsystem off relative to the full +// production shape, so the delta attributes cost to that subsystem. +func benchAblation(b *testing.B, mutate func(*server.Config)) { + cfg := &server.Config{ + LogOnlyModelInteractions: true, + MasterKey: "bench-master-key", + AuditLogger: benchAuditLogger{cfg: auditlog.Config{Enabled: true, LogBodies: true, LogHeaders: true}}, + UsageLogger: benchUsageLogger{cfg: usage.Config{Enabled: true}}, + SessionDetector: session.NewDetector(session.BuiltinRules(), true), + RateLimiter: newBenchRateLimiter(b), + } + mutate(cfg) + + srv := server.New(newBenchRouter(b, routedCatalogSize), cfg) + body := []byte(sampleChatRequest) + + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer bench-master-key") + + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + b.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusOK, rec.Body.String()) + } + } +} + +func BenchmarkAblationFull(b *testing.B) { + benchAblation(b, func(*server.Config) {}) +} + +// No content auto-detection: header/body rules still run, the sha256 conversation +// anchor does not. +func BenchmarkAblationNoAutoDetect(b *testing.B) { + benchAblation(b, func(c *server.Config) { + c.SessionDetector = session.NewDetector(session.BuiltinRules(), false) + }) +} + +func BenchmarkAblationNoSessionKeeping(b *testing.B) { + benchAblation(b, func(c *server.Config) { c.SessionDetector = nil }) +} + +func BenchmarkAblationNoAuditBodies(b *testing.B) { + benchAblation(b, func(c *server.Config) { + c.AuditLogger = benchAuditLogger{cfg: auditlog.Config{Enabled: true, LogHeaders: true}} + }) +} + +func BenchmarkAblationNoAudit(b *testing.B) { + benchAblation(b, func(c *server.Config) { c.AuditLogger = nil }) +} + +func BenchmarkAblationNoAuth(b *testing.B) { + benchAblation(b, func(c *server.Config) { c.MasterKey = "" }) +} + +func BenchmarkAblationNoUsage(b *testing.B) { + benchAblation(b, func(c *server.Config) { c.UsageLogger = nil }) +} + +func BenchmarkAblationNoRateLimit(b *testing.B) { + benchAblation(b, func(c *server.Config) { c.RateLimiter = nil }) +} diff --git a/tests/perf/hotpath_test.go b/tests/perf/hotpath_test.go index 30ced262a..c6f36b78b 100644 --- a/tests/perf/hotpath_test.go +++ b/tests/perf/hotpath_test.go @@ -389,7 +389,7 @@ func TestHotPathPerfGuard(t *testing.T) { { name: "gateway_chat_completion_hot_path", bench: BenchmarkGatewayHotPathChatCompletion, - maxAllocs: 112, // baseline 110 (incl. +1 strings.Clone that unpins the body from RouteHints) + maxAllocs: 108, // baseline 106 (goccy response serializer + single ensureRequestID) maxBytes: 14080, // baseline ~13.5 KB (incl. per-attempt response body/header capture fields) }, { @@ -400,9 +400,20 @@ func TestHotPathPerfGuard(t *testing.T) { // full catalog several times per request) would blow these limits. name: "gateway_chat_completion_hot_path_routed", bench: BenchmarkGatewayHotPathChatCompletionRouted, - maxAllocs: 130, // baseline 128 (incl. +1 strings.Clone that unpins the body from RouteHints) + maxAllocs: 128, // baseline 126 (goccy response serializer + single ensureRequestID) maxBytes: 14656, // baseline ~14.0 KB }, + { + // Default-deployment shape: auth + audit (bodies/headers) + usage + + // session keeping + a configured rate limit, through the routed + // catalog. The bare cases above cannot see regressions in any of + // those subsystems; this one holds the line for the configuration + // deployments actually run. + name: "gateway_chat_completion_production_shape", + bench: BenchmarkGatewayHotPathProductionShape, + maxAllocs: 306, // baseline 300 + maxBytes: 26496, // baseline ~25.1 KB + }, { // Typed chunk decoding + reused read buffer keep this converter at a // fraction of its former map[string]any-per-chunk cost (was 202/19.6KB). diff --git a/tests/perf/production_shape_test.go b/tests/perf/production_shape_test.go new file mode 100644 index 000000000..69fdc9791 --- /dev/null +++ b/tests/perf/production_shape_test.go @@ -0,0 +1,138 @@ +package perf + +import ( + "bytes" + "context" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/enterpilot/gomodel/internal/auditlog" + "github.com/enterpilot/gomodel/internal/core" + "github.com/enterpilot/gomodel/internal/providers" + "github.com/enterpilot/gomodel/internal/ratelimit" + "github.com/enterpilot/gomodel/internal/server" + "github.com/enterpilot/gomodel/internal/session" + "github.com/enterpilot/gomodel/internal/usage" +) + +// benchRateLimitStore is a minimal in-memory ratelimit.Store carrying one +// non-blocking rule, so the benchmark exercises the real per-request window +// accounting a deployment with any configured rate limit pays. +type benchRateLimitStore struct { + rules []ratelimit.Rule +} + +func (s *benchRateLimitStore) ListRules(context.Context) ([]ratelimit.Rule, error) { + return append([]ratelimit.Rule(nil), s.rules...), nil +} +func (s *benchRateLimitStore) UpsertRules(_ context.Context, rules []ratelimit.Rule) error { + s.rules = append(s.rules, rules...) + return nil +} +func (s *benchRateLimitStore) DeleteRule(context.Context, ratelimit.RuleScope, string, int64) error { + return nil +} +func (s *benchRateLimitStore) ReplaceConfigRules(context.Context, []ratelimit.Rule) error { + return nil +} +func (s *benchRateLimitStore) LoadCounters(context.Context) ([]ratelimit.WindowSnapshot, error) { + return nil, nil +} +func (s *benchRateLimitStore) SaveCounters(context.Context, []ratelimit.WindowSnapshot) error { + return nil +} +func (s *benchRateLimitStore) DeleteCounter(context.Context, ratelimit.RuleScope, string, int64) error { + return nil +} +func (s *benchRateLimitStore) DeleteAllCounters(context.Context) error { return nil } +func (s *benchRateLimitStore) Close() error { return nil } + +// newBenchRateLimiter builds a real ratelimit.Service with one high user-path +// request limit that matches every request but never rejects. +func newBenchRateLimiter(tb testing.TB) *ratelimit.Service { + tb.Helper() + + maxRequests := int64(1 << 40) + store := &benchRateLimitStore{rules: []ratelimit.Rule{{ + Scope: ratelimit.ScopeUserPath, + Subject: "/", + PeriodSeconds: 60, + MaxRequests: &maxRequests, + }}} + service, err := ratelimit.NewService(context.Background(), store) + if err != nil { + tb.Fatalf("new rate limit service: %v", err) + } + tb.Cleanup(service.Close) + return service +} + +// BenchmarkGatewayHotPathProductionShape wires the middleware chain the way a +// default deployment actually runs it: master-key auth, audit logging, usage +// tracking, session keeping, and a configured rate limit all enabled. The +// original guard benchmarks leave every one of those nil, so they measure a +// configuration nobody deploys — and therefore cannot see regressions in any +// feature added since the guard was written. TestHotPathPerfGuard enforces +// allocation ceilings on this benchmark too. +func BenchmarkGatewayHotPathProductionShape(b *testing.B) { + srv := newProductionBenchServer(b, routedCatalogSize) + body := []byte(sampleChatRequest) + + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer bench-master-key") + + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + b.Fatalf("status = %d, want %d: %s", rec.Code, http.StatusOK, rec.Body.String()) + } + } +} + +func newBenchRouter(tb testing.TB, modelCount int) *providers.Router { + tb.Helper() + + models := make([]core.Model, 0, modelCount) + models = append(models, core.Model{ID: "gpt-4o-mini", Object: "model", OwnedBy: "mock", Created: 1700000000}) + for i := 1; i < modelCount; i++ { + models = append(models, core.Model{ + ID: fmt.Sprintf("filler-model-%04d", i), + Object: "model", + OwnedBy: "mock", + Created: 1700000000, + }) + } + + registry := providers.NewModelRegistry() + registry.RegisterProviderWithNameAndType(&benchProvider{models: models}, "mock", "mock") + if err := registry.Initialize(context.Background()); err != nil { + tb.Fatalf("registry initialize: %v", err) + } + + router, err := providers.NewRouter(registry) + if err != nil { + tb.Fatalf("new router: %v", err) + } + return router +} + +func newProductionBenchServer(tb testing.TB, modelCount int) *server.Server { + tb.Helper() + + return server.New(newBenchRouter(tb, modelCount), &server.Config{ + LogOnlyModelInteractions: true, + MasterKey: "bench-master-key", + AuditLogger: benchAuditLogger{cfg: auditlog.Config{Enabled: true, LogBodies: true, LogHeaders: true}}, + UsageLogger: benchUsageLogger{cfg: usage.Config{Enabled: true}}, + SessionDetector: session.NewDetector(session.BuiltinRules(), true), + RateLimiter: newBenchRateLimiter(tb), + }) +} diff --git a/tests/perf/session_body_size_test.go b/tests/perf/session_body_size_test.go new file mode 100644 index 000000000..3ae22be6f --- /dev/null +++ b/tests/perf/session_body_size_test.go @@ -0,0 +1,107 @@ +package perf + +import ( + "bytes" + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/enterpilot/gomodel/ext" + "github.com/enterpilot/gomodel/internal/server" + "github.com/enterpilot/gomodel/internal/session" +) + +// recordingRewriter captures the ext.Input the stack hands to a request +// rewriter, so a test can assert what an extension actually sees at ingress. +type recordingRewriter struct { + mu sync.Mutex + last ext.Input + seen bool +} + +func (r *recordingRewriter) Name() string { return "recording" } + +func (r *recordingRewriter) Rewrite(_ context.Context, in ext.Input) (*ext.Result, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.last = in + r.seen = true + return nil, nil +} + +func (r *recordingRewriter) snapshot() (ext.Input, bool) { + r.mu.Lock() + defer r.mu.Unlock() + return r.last, r.seen +} + +// chatBodyOfSize builds a valid chat request whose encoded size is at least +// target bytes, by padding the first user message. The conversation prefix +// (model + leading messages) is what content auto-detection anchors on. +func chatBodyOfSize(target int) []byte { + var buf bytes.Buffer + buf.WriteString(`{"model":"gpt-4o-mini","messages":[{"role":"system","content":"You are a helpful assistant."},{"role":"user","content":"`) + for buf.Len() < target { + buf.WriteString("analyze this repeated log line and summarize it. ") + } + buf.WriteString(`"}]}`) + return buf.Bytes() +} + +// TestSessionIDVisibilityByBodySize pins that content-based session +// auto-detection is independent of request body size. +// +// RequestSnapshotCapture only inlines bodies with ContentLength <= 64 KiB +// (requestSnapshotInlineBodyLimit), so it would be reasonable to expect +// larger conversations to lose their auto-detected id. They do not: an id +// reaches ingress rewriters at every size, including well past the limit. +// +// This matters for cost, not just correctness. Every request carrying a +// session id engages per-session serialization in downstream consumers +// (sticky virtual-model routing, and GoModel Pro's compression epoch locks), +// so there is no size above which a request quietly opts out. +func TestSessionIDVisibilityByBodySize(t *testing.T) { + sizes := []int{ + 1 << 10, // 1 KiB + 32 << 10, // 32 KiB + 63 << 10, // just under the 64 KiB inline capture limit + 65 << 10, // just over it + 256 << 10, + } + + for _, size := range sizes { + t.Run(fmt.Sprintf("%dKiB", size/1024), func(t *testing.T) { + body := chatBodyOfSize(size) + + rewriter := &recordingRewriter{} + srv := server.New(benchProvider{}, &server.Config{ + LogOnlyModelInteractions: true, + SessionDetector: session.NewDetector(session.BuiltinRules(), true), + RequestRewriters: []ext.RequestRewriter{rewriter}, + }) + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.ContentLength = int64(len(body)) + + rec := httptest.NewRecorder() + srv.ServeHTTP(rec, req) + + in, seen := rewriter.snapshot() + if !seen { + t.Fatalf("rewriter never ran (status %d): %s", rec.Code, rec.Body.String()) + } + + if strings.TrimSpace(in.SessionID) == "" { + t.Fatalf("no session id detected for %d KiB body (status %d)", + len(body)/1024, rec.Code) + } + t.Logf("body=%d KiB status=%d session_id=%q", + len(body)/1024, rec.Code, in.SessionID) + }) + } +}