diff --git a/go.mod b/go.mod index bce2211..40f2df6 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index e1f5ef1..4bab9a3 100644 --- a/go.sum +++ b/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= diff --git a/internal/app/app.go b/internal/app/app.go index d82c84f..26ddc03 100644 --- a/internal/app/app.go +++ b/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}) } diff --git a/internal/graph/mcp_overrides.go b/internal/graph/mcp_overrides.go new file mode 100644 index 0000000..806a413 --- /dev/null +++ b/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 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.", + }) +} diff --git a/internal/graph/mcp_overrides_test.go b/internal/graph/mcp_overrides_test.go new file mode 100644 index 0000000..a7ec331 --- /dev/null +++ b/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{} +}