[jaegerai] Refactor route registration and request body size handling

cmd/jaeger/internal/extension/jaegermcp/server_test.go

@@ -45,7 +45,10 @@ func newMockQueryExtension(svc *querysvc.QueryService) *mockQueryExtension {
 	if svc == nil {
 		svc = querysvc.NewQueryService(&tracestoremocks.Reader{}, &depstoremocks.Reader{}, querysvc.QueryServiceOptions{})
 	}
-	return &mockQueryExtension{svc: svc, tm: tenancy.NewManager(&tenancy.Options{})}
+	return &mockQueryExtension{
+		svc: svc,
+		tm:  tenancy.NewManager(&tenancy.Options{}),
+	}
 }
 
 func (m *mockQueryExtension) QueryService() *querysvc.QueryService {

cmd/jaeger/internal/extension/jaegerquery/internal/flags.go

@@ -26,18 +26,28 @@ type UIConfig struct {
 	LogAccess bool `mapstructure:"log_access" valid:"optional"`
 }
 
+// DefaultMaxRequestBodySize is the fallback limit applied when
+// AIConfig.MaxRequestBodySize is left unset (zero).
+const DefaultMaxRequestBodySize int64 = 1 << 20 // 1 MiB
+
 type AIConfig struct {
 	// AgentURL is the WebSocket endpoint of an ACP-compatible agent sidecar.
 	// For example, ws://localhost:16688
 	// See https://agentclientprotocol.com/
 	AgentURL string `mapstructure:"agent_url" valid:"required"`
-	// MaxRequestBodySize is the maximum allowed size in bytes for the chat request body.
+	// A value of 0 selects DefaultMaxRequestBodySize; negative values are rejected.
 	MaxRequestBodySize int64 `mapstructure:"max_request_body_size" valid:"optional"`
 }
 
-func (c AIConfig) Validate() error {
+// Validate checks the AI config and applies DefaultMaxRequestBodySize in place
+// when MaxRequestBodySize is zero; the pointer receiver is required so the
+// default persists back to the caller's config.
+func (c *AIConfig) Validate() error {
 	if c.MaxRequestBodySize < 0 {
-		return errors.New("ai.max_request_body_size must not be negative")
+		return errors.New("ai.max_request_body_size must be a non-negative integer")
+	}
+	if c.MaxRequestBodySize == 0 {
+		c.MaxRequestBodySize = DefaultMaxRequestBodySize
 	}
 	return nil
 }
@@ -72,7 +82,7 @@ func DefaultQueryOptions() QueryOptions {
 		MaxClockSkewAdjust: 0, // disabled by default
 		AI: configoptional.Default(AIConfig{
 			AgentURL:           "ws://localhost:16688",
-			MaxRequestBodySize: 1 << 20, // 1 MiB
+			MaxRequestBodySize: DefaultMaxRequestBodySize,
 		}),
 		HTTP: confighttp.ServerConfig{
 			NetAddr: confignet.AddrConfig{

cmd/jaeger/internal/extension/jaegerquery/internal/flags_test.go

@@ -28,7 +28,14 @@ func TestAIConfigValidateRejectsNegativeBodySize(t *testing.T) {
 	require.Error(t, cfg.Validate())
 }
 
-func TestAIConfigValidateAcceptsZeroBodySize(t *testing.T) {
+func TestAIConfigValidateDefaultsZeroBodySize(t *testing.T) {
 	cfg := AIConfig{MaxRequestBodySize: 0}
 	require.NoError(t, cfg.Validate())
+	require.Equal(t, DefaultMaxRequestBodySize, cfg.MaxRequestBodySize)
+}
+
+func TestAIConfigValidateAcceptsPositiveBodySize(t *testing.T) {
+	cfg := AIConfig{MaxRequestBodySize: 1}
+	require.NoError(t, cfg.Validate())
+	require.Equal(t, int64(1), cfg.MaxRequestBodySize)
 }

cmd/jaeger/internal/extension/jaegerquery/internal/jaegerai/routes.go

@@ -0,0 +1,56 @@
+// Copyright (c) 2026 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package jaegerai
+
+import (
+	"net/http"
+	"strings"
+
+	"go.uber.org/zap"
+)
+
+const routeChat = "/api/ai/chat"
+
+// normalizeBasePath canonicalises the operator-supplied jaeger-query base
+// path so route registration agrees on a single prefix. Empty and "/" both
+// mean "no prefix" and are returned as "". Otherwise any trailing slash is
+// trimmed so concatenating "/api/..." can never produce a double slash.
+func normalizeBasePath(basePath string) string {
+	if basePath == "" || basePath == "/" {
+		return ""
+	}
+	return strings.TrimSuffix(basePath, "/")
+}
+
+// Handler is the entry point for the jaeger-query AI gateway. Callers
+// construct a Handler once (in jaegerquery's Start path), then call
+// RegisterRoutes when wiring the HTTP mux. This mirrors the APIHandler /
+// HTTPGateway pattern used by sibling jaeger-query subsystems and keeps all
+// AI routing inside the jaegerai package.
+type Handler struct {
+	logger             *zap.Logger
+	agentURL           string
+	basePath           string
+	maxRequestBodySize int64
+}
+
+// NewHandler constructs a jaegerai.Handler. agentURL is the WebSocket
+// endpoint of the ACP sidecar; basePath is the jaeger-query base path used
+// to prefix the AI routes. maxRequestBodySize bounds the chat request body
+// to prevent abuse. basePath is normalized once so the registered mux
+// pattern uses a single canonical prefix.
+func NewHandler(logger *zap.Logger, agentURL, basePath string, maxRequestBodySize int64) *Handler {
+	return &Handler{
+		logger:             logger,
+		agentURL:           agentURL,
+		basePath:           normalizeBasePath(basePath),
+		maxRequestBodySize: maxRequestBodySize,
+	}
+}
+
+// RegisterRoutes mounts the AI gateway endpoints on the provided mux. The
+// chat endpoint streams ACP turns to/from the sidecar.
+func (h *Handler) RegisterRoutes(router *http.ServeMux) {
+	router.HandleFunc(h.basePath+routeChat, NewChatHandler(h.logger, h.agentURL, h.maxRequestBodySize).ServeHTTP)
+}

cmd/jaeger/internal/extension/jaegerquery/internal/jaegerai/routes_test.go

@@ -0,0 +1,73 @@
+// Copyright (c) 2026 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package jaegerai
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
+)
+
+func TestNewHandlerStoresConfig(t *testing.T) {
+	h := NewHandler(zap.NewNop(), "ws://example", "/jaeger", 1<<20)
+	assert.Equal(t, "ws://example", h.agentURL)
+	assert.Equal(t, "/jaeger", h.basePath)
+	assert.Equal(t, int64(1<<20), h.maxRequestBodySize)
+}
+
+func TestRegisterRoutesMountsChatEndpoint(t *testing.T) {
+	tests := []struct {
+		name     string
+		basePath string
+		wantChat string
+	}{
+		{
+			name:     "no base path",
+			basePath: "",
+			wantChat: "/api/ai/chat",
+		},
+		{
+			name:     "single-slash base path is treated as no prefix",
+			basePath: "/",
+			wantChat: "/api/ai/chat",
+		},
+		{
+			name:     "with base path",
+			basePath: "/jaeger",
+			wantChat: "/jaeger/api/ai/chat",
+		},
+		{
+			// Operator-supplied trailing slash must be normalized away so we
+			// don't register a "/jaeger//api/..." pattern.
+			name:     "trailing slash in base path is normalized",
+			basePath: "/jaeger/",
+			wantChat: "/jaeger/api/ai/chat",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			h := NewHandler(zap.NewNop(), "ws://127.0.0.1:1", tc.basePath, 1<<20)
+			mux := http.NewServeMux()
+			h.RegisterRoutes(mux)
+
+			// Chat endpoint: GET (wrong method) is enough to confirm the
+			// route is mounted — the handler returns 405 instead of the
+			// mux returning 404.
+			req := httptest.NewRequest(http.MethodGet, tc.wantChat, http.NoBody)
+			rr := httptest.NewRecorder()
+			mux.ServeHTTP(rr, req)
+			assert.Equal(t, http.StatusMethodNotAllowed, rr.Code,
+				"chat endpoint should be mounted at %s", tc.wantChat)
+		})
+	}
+}
+
+func TestNewHandlerNormalizesTrailingSlash(t *testing.T) {
+	h := NewHandler(zap.NewNop(), "ws://127.0.0.1:1", "/jaeger/", 1<<20)
+	assert.Equal(t, "/jaeger", h.basePath, "NewHandler must trim the trailing slash")
+}

cmd/jaeger/internal/extension/jaegerquery/internal/server.go

@@ -193,15 +193,11 @@ func initRouter(
 
 	// AI Gateway Endpoints
 	if queryOpts.AI.HasValue() {
-		aiHandlerPath := "/api/ai/chat"
-		if queryOpts.BasePath != "" && queryOpts.BasePath != "/" {
-			aiHandlerPath = queryOpts.BasePath + aiHandlerPath
-		}
 		if aiCfg := queryOpts.AI.Get(); aiCfg != nil && aiCfg.AgentURL != "" {
 			if err := aiCfg.Validate(); err != nil {
 				telset.Logger.Error("Invalid AI config, AI handler disabled", zap.Error(err))
 			} else {
-				r.HandleFunc(aiHandlerPath, jaegerai.NewChatHandler(telset.Logger, aiCfg.AgentURL, aiCfg.MaxRequestBodySize).ServeHTTP)
+				jaegerai.NewHandler(telset.Logger, aiCfg.AgentURL, queryOpts.BasePath, aiCfg.MaxRequestBodySize).RegisterRoutes(r)
 			}
 		}
 	}