fix: introduce fallback stub for ToolMetricsClient in dev and prod (#580)

- Refactor `ShouldFlag` signatures to accept an organisation ID
- Introduce a fallback stub for `ToolMetricsClient` in dev and prod.

Author: jirevwe

server/cmd/gram/deps.go

@@ -97,22 +97,23 @@ func newToolMetricsClient(ctx context.Context, logger *slog.Logger, c *cli.Conte
 		},
 	})
 	if err != nil {
-		return nil, nilFunc, fmt.Errorf("connect to clickhouse: %w", err)
+		logger.WarnContext(ctx, "error connecting to clickhouse; falling back to stub tool call metrics client")
+		return &tm.StubToolMetricsClient{}, func(context.Context) error { return nil }, nil
 	}
 
 	if err = conn.Ping(ctx); err != nil {
-		return nil, nilFunc, fmt.Errorf("ping clickhouse: %w", err)
+		logger.WarnContext(ctx, "failed to ping clickhouse; falling back to stub tool call metrics client", attr.SlogError(err))
+		return &tm.StubToolMetricsClient{}, func(context.Context) error { return nil }, nil
 	}
 
-	cc := tm.New(logger, tracerProvider, conn, func(ctx context.Context, log tm.ToolHTTPRequest) (bool, error) {
+	cc := tm.New(logger, tracerProvider, conn, func(ctx context.Context, orgId string) (bool, error) {
 		f := conv.Default[feature.Provider](features, &feature.InMemory{})
-		isEnabled, err := f.IsFlagEnabled(ctx, feature.FlagClickhouseToolMetrics, log.OrganizationID)
+		isEnabled, err := f.IsFlagEnabled(ctx, feature.FlagClickhouseToolMetrics, orgId)
 		if err != nil {
 			logger.ErrorContext(
 				ctx, "error checking clickhouse feature flag",
 				attr.SlogError(err),
-				attr.SlogOrganizationSlug(log.OrganizationID),
-				attr.SlogProjectID(log.ProjectID),
+				attr.SlogOrganizationSlug(orgId),
 			)
 			return false, fmt.Errorf("check clickhouse feature flag: %w", err)
 		}

server/internal/gateway/proxy.go

@@ -625,13 +625,31 @@ func reverseProxyRequest(
 		MaxIdleConnsPerHost:   runtime.GOMAXPROCS(0) + 1,
 	}
 
-	// Wrap with HTTP logging round tripper
-	loggingTransport := tm.NewHTTPLoggingRoundTripper(transport, tcm, logger, tracer)
-
-	otelTransport := otelhttp.NewTransport(
-		loggingTransport,
-		otelhttp.WithPropagators(propagation.TraceContext{}),
-	)
+	isAllowed, err := tcm.ShouldLog(ctx, tool.OrganizationID)
+	if err != nil {
+		// If we can't determine if the tool is allowed to log, we won't log the request.
+		isAllowed = false
+		logger.ErrorContext(ctx,
+			"failed to determine if tool is allowed to log",
+			attr.SlogOrganizationID(tool.OrganizationID),
+			attr.SlogToolName(tool.Name),
+			attr.SlogError(err))
+	}
+
+	var otelTransport *otelhttp.Transport
+	if isAllowed {
+		// Wrap with HTTP logging round tripper
+		loggingTransport := tm.NewHTTPLoggingRoundTripper(transport, tcm, logger, tracer)
+		otelTransport = otelhttp.NewTransport(
+			loggingTransport,
+			otelhttp.WithPropagators(propagation.TraceContext{}),
+		)
+	} else {
+		otelTransport = otelhttp.NewTransport(
+			transport,
+			otelhttp.WithPropagators(propagation.TraceContext{}),
+		)
+	}
 
 	client := &http.Client{
 		Timeout:   60 * time.Second,

server/internal/gateway/proxy_test.go

@@ -70,7 +70,7 @@ func newClickhouseClient(t *testing.T, orgId string) *toolmetrics.Queries {
 
 	tracerProvider := testenv.NewTracerProvider(t)
 
-	ch := toolmetrics.New(testenv.NewLogger(t), tracerProvider, chConn, func(ctx context.Context, log toolmetrics.ToolHTTPRequest) (bool, error) {
+	ch := toolmetrics.New(testenv.NewLogger(t), tracerProvider, chConn, func(context.Context, string) (bool, error) {
 		return true, nil
 	})
 

server/internal/logs/setup_test.go

@@ -74,7 +74,7 @@ func newTestLogsService(t *testing.T) (context.Context, *testInstance) {
 
 	tracerProvider := testenv.NewTracerProvider(t)
 
-	chClient := toolmetrics.New(logger, tracerProvider, chConn, func(ctx context.Context, log toolmetrics.ToolHTTPRequest) (bool, error) {
+	chClient := toolmetrics.New(logger, tracerProvider, chConn, func(context.Context, string) (bool, error) {
 		return true, nil
 	})
 

server/internal/mcp/setup_test.go

@@ -94,7 +94,7 @@ func newTestMCPService(t *testing.T) (context.Context, *testInstance) {
 	chConn, err := infra.NewClickhouseClient(t)
 	require.NoError(t, err)
 
-	toolMetrics := toolmetrics.New(logger, tracerProvider, chConn, func(ctx context.Context, log toolmetrics.ToolHTTPRequest) (bool, error) {
+	toolMetrics := toolmetrics.New(logger, tracerProvider, chConn, func(context.Context, string) (bool, error) {
 		return true, nil
 	})
 

server/internal/thirdparty/toolmetrics/db.go

@@ -20,7 +20,7 @@ type Queries struct {
 	conn       CHTX
 	logger     *slog.Logger
 	tracer     trace.Tracer
-	ShouldFlag func(ctx context.Context, log ToolHTTPRequest) (bool, error)
+	ShouldFlag func(ctx context.Context, orgId string) (bool, error)
 }
 
 // WithConn returns a new Queries instance using the provided connection.
@@ -34,9 +34,9 @@ func (q *Queries) WithConn(conn CHTX) *Queries {
 }
 
 // New creates a new Queries instance with logger and tracer.
-func New(logger *slog.Logger, traceProvider trace.TracerProvider, conn CHTX, shouldFlag func(ctx context.Context, log ToolHTTPRequest) (bool, error)) *Queries {
+func New(logger *slog.Logger, traceProvider trace.TracerProvider, conn CHTX, shouldFlag func(ctx context.Context, orgId string) (bool, error)) *Queries {
 	if shouldFlag == nil {
-		shouldFlag = func(ctx context.Context, log ToolHTTPRequest) (bool, error) {
+		shouldFlag = func(ctx context.Context, orgId string) (bool, error) {
 			return true, nil
 		}
 	}

server/internal/thirdparty/toolmetrics/models.go

@@ -96,4 +96,6 @@ type ToolMetricsProvider interface {
 	List(ctx context.Context, opts ListToolLogsOptions) (*ListResult, error)
 	// Log tool call request/response
 	Log(context.Context, ToolHTTPRequest) error
+	// ShouldLog returns true if the tool call should be logged
+	ShouldLog(context.Context, string) (bool, error)
 }

server/internal/thirdparty/toolmetrics/queries.sql.go

@@ -37,6 +37,10 @@ order by ts
 limit $5
 `
 
+func (q *Queries) ShouldLog(ctx context.Context, orgId string) (bool, error) {
+	return q.ShouldFlag(ctx, orgId)
+}
+
 // List retrieves tool logs based on the provided options.
 func (q *Queries) List(ctx context.Context, opts ListToolLogsOptions) (res *ListResult, err error) {
 	projectID := opts.ProjectID
@@ -133,7 +137,7 @@ func (q *Queries) List(ctx context.Context, opts ListToolLogsOptions) (res *List
 
 // Log inserts a tool HTTP request log entry.
 func (q *Queries) Log(ctx context.Context, log ToolHTTPRequest) (err error) {
-	allow, err := q.ShouldFlag(ctx, log)
+	allow, err := q.ShouldFlag(ctx, log.OrganizationID)
 	if err != nil {
 		q.logger.ErrorContext(ctx, "failed to fetch feature flag", attr.SlogError(err))
 		return nil

server/internal/thirdparty/toolmetrics/stub.go

@@ -0,0 +1,19 @@
+package toolmetrics
+
+import (
+	"context"
+)
+
+type StubToolMetricsClient struct{}
+
+func (n *StubToolMetricsClient) List(_ context.Context, _ ListToolLogsOptions) (*ListResult, error) {
+	return nil, nil
+}
+
+func (n *StubToolMetricsClient) Log(_ context.Context, _ ToolHTTPRequest) error {
+	return nil
+}
+
+func (n *StubToolMetricsClient) ShouldLog(_ context.Context, _ string) (bool, error) {
+	return true, nil
+}