fix(mcp): make signals shortcut tools return data via templated selections

go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/DIMO-Network/credit-tracker v0.0.6
 	github.com/DIMO-Network/fetch-api v0.0.16
 	github.com/DIMO-Network/model-garage v1.0.13
-	github.com/DIMO-Network/server-garage v0.1.1
+	github.com/DIMO-Network/server-garage v0.4.0
 	github.com/DIMO-Network/shared v1.1.5
 	github.com/DIMO-Network/token-exchange-api v0.4.0
 	github.com/aarondl/sqlboiler/v4 v4.19.5

go.sum

@@ -24,8 +24,8 @@ github.com/DIMO-Network/fetch-api v0.0.16 h1:jOIM9AHOCTN9Omh5QoRrdVsy0UCPHwB+08/
 github.com/DIMO-Network/fetch-api v0.0.16/go.mod h1:pD9+5wPYjz3ORH9FBsXLlomsZJRPTyOzKy19urT6xew=
 github.com/DIMO-Network/model-garage v1.0.13 h1:KwFJ/v6XETGNtiqmP7/FbW9eQTZFVbIpZ6gdHzGL+Rc=
 github.com/DIMO-Network/model-garage v1.0.13/go.mod h1:oi7EGKQVxFVpXRsu2H+YbizbKcx06aQg2N1Yu4GqOp8=
-github.com/DIMO-Network/server-garage v0.1.1 h1:EYmyy+Fgi2BNW0Bufn04BViDtb8BCWaN7C7BbEuoI5s=
-github.com/DIMO-Network/server-garage v0.1.1/go.mod h1:Z3A1KDUsXey+XhrPhmw/wyCidfrQvmEdWp7nShno7ZM=
+github.com/DIMO-Network/server-garage v0.4.0 h1:3ukXvtldIhLldn9AtbtQBooBjT2gicvB3WphrsjNQho=
+github.com/DIMO-Network/server-garage v0.4.0/go.mod h1:Z3A1KDUsXey+XhrPhmw/wyCidfrQvmEdWp7nShno7ZM=
 github.com/DIMO-Network/shared v1.1.5 h1:cRI3BbeYOgolMkeBOSqbDVGxymnDfeP/q7xgIXJ5MkY=
 github.com/DIMO-Network/shared v1.1.5/go.mod h1:lDHUKwwT2LW6Zvd42Nb33dXklRNTmfyOlbUNx2dQfGY=
 github.com/DIMO-Network/token-exchange-api v0.4.0 h1:EayDrw9VdyAfc6rbpdnDxFhlN3lMhbonUJoouKZu35g=

internal/app/app.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/99designs/gqlgen/graphql"
+	"github.com/99designs/gqlgen/graphql/executor"
 	"github.com/99designs/gqlgen/graphql/handler"
 	"github.com/99designs/gqlgen/graphql/handler/extension"
 	"github.com/99designs/gqlgen/graphql/handler/transport"
@@ -81,11 +82,10 @@ func New(settings config.Settings) (*App, error) {
 	var costCalculator pricing.CostCalculator
 
 	es := graph.NewExecutableSchema(cfg)
-	server := newServer(es)
-	server.Use(dtcmiddleware.NewDCT(ctClient, &costCalculator))
+	dct := dtcmiddleware.NewDCT(ctClient, &costCalculator)
 
-	// Add query recording middleware
-	server.Use(queryRecorder.QueryRecordingExtension{Recorder: queryRec})
+	server := newServer(es)
+	configureGQLExtensions(server, dct, queryRec)
 
 	authMiddleware, err := auth.NewJWTMiddleware(settings.TokenExchangeIssuer, settings.TokenExchangeJWTKeySetURL)
 	if err != nil {
@@ -111,8 +111,19 @@ func New(settings config.Settings) (*App, error) {
 		),
 	)
 
-	mcpHandler, err := mcpserver.New(mcpserver.NewGQLGenExecutor(es), "DIMO Telemetry", "0.1.0", "telemetry",
-		mcpserver.WithTools(graph.MCPTools),
+	mcpTools, err := graph.OverrideMCPTools(graph.MCPTools)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't override MCP tools: %w", err)
+	}
+
+	// The MCP executor gets the same extension set as the HTTP GraphQL server
+	// so tool calls are billed, recorded, and complexity-limited identically.
+	mcpExec := mcpserver.NewGQLGenExecutor(es, func(e *executor.Executor) {
+		configureGQLExtensions(e, dct, queryRec)
+	})
+
+	mcpHandler, err := mcpserver.New(mcpExec, "DIMO Telemetry", "0.1.0", "telemetry",
+		mcpserver.WithTools(mcpTools),
 		mcpserver.WithCondensedSchema(graph.CondensedSchema),
 	)
 	if err != nil {
@@ -175,11 +186,29 @@ func newServer(es graphql.ExecutableSchema) *handler.Server {
 	srv.AddTransport(transport.GET{})
 	srv.AddTransport(transport.POST{})
 	srv.AddTransport(transport.MultipartForm{})
+	// srv.SetQueryCache(graphql.NoCache[*ast.QueryDocument]{})
+
+	return srv
+}
+
+// gqlExtensionTarget is satisfied by both *handler.Server and
+// *executor.Executor, letting the HTTP GraphQL server and the MCP executor
+// share one extension configuration.
+type gqlExtensionTarget interface {
+	Use(graphql.HandlerExtension)
+	SetErrorPresenter(graphql.ErrorPresenterFunc)
+}
+
+// configureGQLExtensions applies the extension set shared by the HTTP GraphQL
+// server and the MCP executor: complexity limit, introspection, metrics,
+// error shaping, DCX credit tracking, and query recording. Keeping both paths
+// on one list means MCP tool calls are billed and limited exactly like
+// regular GraphQL requests.
+func configureGQLExtensions(srv gqlExtensionTarget, dct *dtcmiddleware.DCT, queryRec *queryRecorder.QueryRecorder) {
 	srv.Use(extension.FixedComplexityLimit(100))
 	srv.Use(extension.Introspection{})
 	srv.Use(metrics.Tracer{})
-	// srv.SetQueryCache(graphql.NoCache[*ast.QueryDocument]{})
 	srv.SetErrorPresenter(errorPresenter)
-
-	return srv
+	srv.Use(dct)
+	srv.Use(queryRecorder.QueryRecordingExtension{Recorder: queryRec})
 }

internal/graph/mcp_overrides.go

@@ -0,0 +1,80 @@
+package graph
+
+import (
+	"fmt"
+
+	"github.com/DIMO-Network/server-garage/pkg/mcpserver"
+)
+
+// OverrideMCPTools patches the mcpgen-generated MCPTools slice so the two
+// shortcut tools that wrap `signals` and `signalsLatest` actually work.
+//
+// Neither GraphQL field declares per-signal arguments — callers name signals
+// by field-selection (e.g. `speed(agg: AVG)`). mcpgen, which builds a single
+// static query string per tool, can't express that shape. The two tools
+// therefore shipped with a selection of only `timestamp` / `lastSeen` and
+// returned empty data.
+//
+// Rather than change the GraphQL schema, this override uses the templated
+// selection + tool-only-argument features from server-garage (v0.3.0) to
+// reshape the two ToolDefinitions at startup:
+//
+//   - SelectionTemplate renders the concrete selection set at call time
+//     from the supplied signal list.
+//   - A ToolOnly ArgDefinition lets the model pass that list in, without
+//     the argument flowing through to the GraphQL executor.
+//
+// The GraphQL schema is untouched.
+//
+// Returns an error if either tool is missing from the generated slice, so a
+// rename in mcpgen output fails startup instead of silently shipping the
+// original empty-selection tools.
+func OverrideMCPTools(tools []mcpserver.ToolDefinition) ([]mcpserver.ToolDefinition, error) {
+	out := make([]mcpserver.ToolDefinition, len(tools))
+	copy(out, tools)
+	overridden := make(map[string]bool, 2)
+	for i := range out {
+		switch out[i].Name {
+		case "telemetry_get_signals_time_series":
+			overrideSignalsTimeSeries(&out[i])
+			overridden[out[i].Name] = true
+		case "telemetry_get_latest_signals":
+			overrideLatestSignals(&out[i])
+			overridden[out[i].Name] = true
+		}
+	}
+	for _, name := range []string{"telemetry_get_signals_time_series", "telemetry_get_latest_signals"} {
+		if !overridden[name] {
+			return nil, fmt.Errorf("graph: OverrideMCPTools did not find tool %q in generated MCPTools; mcpgen output may have renamed it", name)
+		}
+	}
+	return out, nil
+}
+
+func overrideSignalsTimeSeries(t *mcpserver.ToolDefinition) {
+	t.Description = "Get aggregated time series for a named list of float signals. Pass signalRequests as [{name, agg}] (e.g. [{name:\"speed\",agg:\"AVG\"},{name:\"powertrainTractionBatteryStateOfChargeCurrent\",agg:\"LAST\"}]). Returns buckets of {timestamp, <signal>: <value>, ...}. Signal names come from get_available_signals or get_data_summary. Aggregations: AVG, MED, MAX, MIN, RAND, FIRST, LAST."
+	t.Query = `query($tokenId: Int!, $interval: String!, $from: Time!, $to: Time!, $filter: SignalFilter) { signals(tokenId: $tokenId, interval: $interval, from: $from, to: $to, filter: $filter) { __MCPGEN_SELECTION__ } }`
+	t.SelectionTemplate = "timestamp{{range .signalRequests}} {{.name}}(agg: {{.agg}}){{end}}"
+	t.Args = append(t.Args, mcpserver.ArgDefinition{
+		Name:        "signalRequests",
+		Type:        "array",
+		ItemsType:   "object",
+		Required:    true,
+		ToolOnly:    true,
+		Description: "List of {name, agg} pairs specifying which float signals to aggregate. Each `name` is a signal field name; each `agg` is one of AVG, MED, MAX, MIN, RAND, FIRST, LAST.",
+	})
+}
+
+func overrideLatestSignals(t *mcpserver.ToolDefinition) {
+	t.Description = "Get the most recent value for a named list of float signals. Pass signalNames as an array of strings (e.g. [\"speed\",\"powertrainTractionBatteryStateOfChargeCurrent\"]). For non-float signals (strings, locations) use get_signals_snapshot. Signal names come from get_available_signals or get_data_summary."
+	t.Query = `query($tokenId: Int!, $filter: SignalFilter) { signalsLatest(tokenId: $tokenId, filter: $filter) { __MCPGEN_SELECTION__ } }`
+	t.SelectionTemplate = "lastSeen{{range .signalNames}} {{.}} {timestamp value}{{end}}"
+	t.Args = append(t.Args, mcpserver.ArgDefinition{
+		Name:        "signalNames",
+		Type:        "array",
+		ItemsType:   "string",
+		Required:    true,
+		ToolOnly:    true,
+		Description: "List of float-signal field names to return the latest value for.",
+	})
+}

internal/graph/mcp_overrides_test.go

@@ -0,0 +1,115 @@
+package graph
+
+import (
+	"strings"
+	"testing"
+	"text/template"
+
+	"github.com/DIMO-Network/server-garage/pkg/mcpserver"
+	"github.com/stretchr/testify/require"
+	"github.com/vektah/gqlparser/v2/ast"
+	"github.com/vektah/gqlparser/v2/parser"
+)
+
+func TestOverrideMCPTools_PatchesBothTools(t *testing.T) {
+	out, err := OverrideMCPTools(MCPTools)
+	require.NoError(t, err)
+	require.Len(t, out, len(MCPTools))
+
+	for _, name := range []string{"telemetry_get_signals_time_series", "telemetry_get_latest_signals"} {
+		tool := findTool(t, out, name)
+		require.Contains(t, tool.Query, mcpserver.SelectionPlaceholder, "%s query must contain selection placeholder", name)
+		require.NotEmpty(t, tool.SelectionTemplate, "%s must have a selection template", name)
+		last := tool.Args[len(tool.Args)-1]
+		require.True(t, last.ToolOnly, "%s appended arg must be ToolOnly", name)
+		require.True(t, last.Required, "%s appended arg must be required", name)
+	}
+}
+
+func TestOverrideMCPTools_DoesNotMutateInput(t *testing.T) {
+	_, err := OverrideMCPTools(MCPTools)
+	require.NoError(t, err)
+
+	for _, name := range []string{"telemetry_get_signals_time_series", "telemetry_get_latest_signals"} {
+		orig := findTool(t, MCPTools, name)
+		require.NotContains(t, orig.Query, mcpserver.SelectionPlaceholder, "generated %s must be untouched", name)
+		require.Empty(t, orig.SelectionTemplate)
+		for _, a := range orig.Args {
+			require.False(t, a.ToolOnly)
+		}
+	}
+}
+
+func TestOverrideMCPTools_ErrorsWhenToolMissing(t *testing.T) {
+	var trimmed []mcpserver.ToolDefinition
+	for _, tool := range MCPTools {
+		if tool.Name == "telemetry_get_signals_time_series" {
+			continue
+		}
+		trimmed = append(trimmed, tool)
+	}
+
+	_, err := OverrideMCPTools(trimmed)
+	require.Error(t, err, "must fail when an expected tool is missing from generated output")
+
+	_, err = OverrideMCPTools(nil)
+	require.Error(t, err)
+}
+
+// TestOverrideMCPTools_TemplatesRenderValidGraphQL renders each override's
+// selection template with representative arguments, splices it into the query,
+// and checks the result parses as GraphQL.
+func TestOverrideMCPTools_TemplatesRenderValidGraphQL(t *testing.T) {
+	out, err := OverrideMCPTools(MCPTools)
+	require.NoError(t, err)
+
+	cases := []struct {
+		toolName string
+		args     map[string]any
+	}{
+		{
+			toolName: "telemetry_get_signals_time_series",
+			args: map[string]any{
+				"signalRequests": []any{
+					map[string]any{"name": "speed", "agg": "AVG"},
+					map[string]any{"name": "powertrainTractionBatteryStateOfChargeCurrent", "agg": "LAST"},
+				},
+			},
+		},
+		{
+			toolName: "telemetry_get_latest_signals",
+			args: map[string]any{
+				"signalNames": []any{"speed", "powertrainTractionBatteryStateOfChargeCurrent"},
+			},
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.toolName, func(t *testing.T) {
+			tool := findTool(t, out, tc.toolName)
+
+			tmpl, err := template.New(tc.toolName).Parse(tool.SelectionTemplate)
+			require.NoError(t, err)
+			var buf strings.Builder
+			require.NoError(t, tmpl.Execute(&buf, tc.args))
+
+			query := strings.Replace(tool.Query, mcpserver.SelectionPlaceholder, buf.String(), 1)
+			_, parseErr := parser.ParseQuery(&ast.Source{Input: query})
+			require.Nil(t, parseErr, "rendered query must parse: %s", query)
+
+			require.Contains(t, buf.String(), "speed")
+			require.Contains(t, buf.String(), "powertrainTractionBatteryStateOfChargeCurrent")
+		})
+	}
+}
+
+func findTool(t *testing.T, tools []mcpserver.ToolDefinition, name string) mcpserver.ToolDefinition {
+	t.Helper()
+	for _, tool := range tools {
+		if tool.Name == name {
+			return tool
+		}
+	}
+	t.Fatalf("tool %q not found", name)
+	return mcpserver.ToolDefinition{}
+}