diff --git a/pkg/authz/middleware.go b/pkg/authz/middleware.go index 2f20e0c021..343f655d12 100644 --- a/pkg/authz/middleware.go +++ b/pkg/authz/middleware.go @@ -19,6 +19,7 @@ import ( "github.com/stacklok/toolhive/pkg/mcp" "github.com/stacklok/toolhive/pkg/transport/ssecommon" "github.com/stacklok/toolhive/pkg/transport/types" + "github.com/stacklok/toolhive/pkg/vmcp/optimizer" "github.com/stacklok/toolhive/pkg/vmcp/session/optimizerdec" ) @@ -287,7 +288,8 @@ func authorizeAndServe( // It always fully handles the request (authorization, unauthorized response, or serving). // // For pass-through meta-tools (find_tool, call_tool): -// - call_tool: authorizes the real inner tool name from arguments["tool_name"]. +// - call_tool: authorizes the real inner tool name from arguments["tool_name"], +// or from arguments["parameters"]["tool_name"] when the caller nested it there. // - find_tool (and other pass-through tools without a tool_name): allowed through // as a discovery operation with no policy check. // @@ -303,9 +305,13 @@ func handleToolsCall( next http.Handler, ) { if _, isPassThrough := passThroughTools[parsedRequest.ResourceID]; isPassThrough { - if toolName, ok := parsedRequest.Arguments[optimizerdec.CallToolArgToolName].(string); ok && toolName != "" { + rawName, _ := parsedRequest.Arguments[optimizerdec.CallToolArgToolName].(string) + rawArgs, _ := parsedRequest.Arguments[optimizerdec.CallToolArgParameters].(map[string]interface{}) + // Resolve the same way dispatch does, so a nested tool_name is authorized + // under the name that will actually run rather than skipping the check. + toolName, innerArgs := optimizer.ResolveCallToolTarget(rawName, rawArgs) + if toolName != "" { // call_tool: authorize the real backend tool name. - innerArgs, _ := parsedRequest.Arguments[optimizerdec.CallToolArgParameters].(map[string]interface{}) authorizeAndServe(w, r, a, annotationCache, featureOp.Feature, featureOp.Operation, parsedRequest.ID, toolName, innerArgs, next) diff --git a/pkg/authz/middleware_test.go b/pkg/authz/middleware_test.go index 6d8d93099b..99bbe8a2da 100644 --- a/pkg/authz/middleware_test.go +++ b/pkg/authz/middleware_test.go @@ -1149,6 +1149,33 @@ func TestMiddlewareOptimizerMetaTools(t *testing.T) { expectStatus: http.StatusForbidden, expectHandlerHit: false, }, + { + // Without hoisting, the top-level lookup misses and the request falls + // through the pass-through branch entirely, reaching the backend with + // no policy check at all. + name: "call_tool with nested unauthorized inner tool is blocked", + toolName: "call_tool", + arguments: map[string]interface{}{ + "parameters": map[string]interface{}{ + "tool_name": "forbidden_backend", + "query": "x", + }, + }, + expectStatus: http.StatusForbidden, + expectHandlerHit: false, + }, + { + name: "call_tool with nested authorized inner tool passes through", + toolName: "call_tool", + arguments: map[string]interface{}{ + "parameters": map[string]interface{}{ + "tool_name": "allowed_backend", + "query": "x", + }, + }, + expectStatus: http.StatusOK, + expectHandlerHit: true, + }, { name: "find_tool request reaches handler (response filtering applied separately)", toolName: "find_tool", diff --git a/pkg/vmcp/optimizer/optimizer.go b/pkg/vmcp/optimizer/optimizer.go index a28bafc3b8..eb1d670880 100644 --- a/pkg/vmcp/optimizer/optimizer.go +++ b/pkg/vmcp/optimizer/optimizer.go @@ -14,9 +14,12 @@ package optimizer import ( "context" + "encoding/json" "fmt" "log/slog" + "maps" "os" + "slices" "strconv" "strings" "time" @@ -229,6 +232,43 @@ type CallToolInput struct { Parameters map[string]any `json:"parameters" description:"Dictionary of arguments required by the tool. The structure must match the tool's input schema as returned by find_tool."` } +// ResolveCallToolTarget resolves the tool a call_tool invocation targets, +// accepting the common LLM malformation where tool_name is nested inside +// parameters instead of sitting alongside it. A top-level name always wins, so a +// backend tool with its own tool_name argument still works. +// +// Every consumer of a call_tool payload must resolve the name through this +// function. Authorization reads the name from the raw request while dispatch +// reads it from the decoded struct, and a target the two disagree on is a tool +// executing under a policy decision made for a different name. +// +// params is never modified; a copy is returned when a nested name is hoisted. +func ResolveCallToolTarget(name string, params map[string]any) (string, map[string]any) { + if name != "" { + return name, params + } + nested, ok := params["tool_name"].(string) + if !ok || nested == "" { + return name, params + } + hoisted := maps.Clone(params) + delete(hoisted, "tool_name") + return nested, hoisted +} + +// UnmarshalJSON hoists a nested tool_name so dispatch targets the same tool +// authorization approved. See ResolveCallToolTarget. +func (in *CallToolInput) UnmarshalJSON(data []byte) error { + type rawCallToolInput CallToolInput // drops the method set to avoid recursion + var raw rawCallToolInput + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + raw.ToolName, raw.Parameters = ResolveCallToolTarget(raw.ToolName, raw.Parameters) + *in = CallToolInput(raw) + return nil +} + // NewOptimizerFactory creates the embedding client and SQLite tool store from // the given OptimizerConfig, then returns an OptimizerFactory and a cleanup // function that closes the store. The caller must invoke the cleanup function @@ -380,7 +420,10 @@ func (d *toolOptimizer) FindTool(ctx context.Context, input FindToolInput) (*Fin // is invoked directly with the given parameters. func (d *toolOptimizer) CallTool(ctx context.Context, input CallToolInput) (*mcp.CallToolResult, error) { if input.ToolName == "" { - return nil, fmt.Errorf("tool_name is required") + return nil, fmt.Errorf( + `tool_name is required: call_tool expects {"tool_name": "", `+ + `"parameters": {}}, got parameters keys %v`, + slices.Sorted(maps.Keys(input.Parameters))) } // Verify the tool exists diff --git a/pkg/vmcp/optimizer/optimizer_test.go b/pkg/vmcp/optimizer/optimizer_test.go index e80c17cec6..9fe12a7243 100644 --- a/pkg/vmcp/optimizer/optimizer_test.go +++ b/pkg/vmcp/optimizer/optimizer_test.go @@ -7,6 +7,7 @@ import ( "context" "encoding/json" "fmt" + "maps" "strings" "testing" @@ -20,6 +21,7 @@ import ( "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/tokencounter" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types" "github.com/stacklok/toolhive/pkg/vmcp/optimizer/internal/types/mocks" + "github.com/stacklok/toolhive/pkg/vmcp/schema" ) func TestGetAndValidateConfig(t *testing.T) { @@ -845,7 +847,7 @@ func TestOptimizer_CallTool(t *testing.T) { Parameters: map[string]any{}, }, expectedError: true, - errorContains: "tool_name is required", + errorContains: `call_tool expects {"tool_name"`, }, } @@ -877,3 +879,86 @@ func TestOptimizer_CallTool(t *testing.T) { }) } } + +func TestResolveCallToolTarget(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + toolName string + params map[string]any + expectedName string + expectedParams map[string]any + }{ + { + name: "top-level name is returned unchanged", + toolName: "search", + params: map[string]any{"query": "weather"}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "nested name is hoisted and removed from params", + params: map[string]any{"tool_name": "search", "query": "weather"}, + expectedName: "search", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "top-level name wins and params are left untouched", + toolName: "search", + params: map[string]any{"tool_name": "other", "query": "weather"}, + expectedName: "search", + expectedParams: map[string]any{"tool_name": "other", "query": "weather"}, + }, + { + name: "nested empty name is not hoisted", + params: map[string]any{"tool_name": ""}, + expectedName: "", + expectedParams: map[string]any{"tool_name": ""}, + }, + { + name: "nested non-string name is not hoisted", + params: map[string]any{"tool_name": 123}, + expectedName: "", + expectedParams: map[string]any{"tool_name": 123}, + }, + { + name: "no name anywhere", + params: map[string]any{"query": "weather"}, + expectedName: "", + expectedParams: map[string]any{"query": "weather"}, + }, + { + name: "nil params", + expectedName: "", + expectedParams: nil, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + // Callers share this map with downstream consumers, so it must survive intact. + original := maps.Clone(tc.params) + + gotName, gotParams := ResolveCallToolTarget(tc.toolName, tc.params) + require.Equal(t, tc.expectedName, gotName) + require.Equal(t, tc.expectedParams, gotParams) + require.Equal(t, original, tc.params, "input params must not be mutated") + }) + } +} + +// Both call_tool handlers decode via schema.Translate, so the hoist must survive +// that round-trip and not just a direct json.Unmarshal. +func TestCallToolInput_TranslateHoistsNestedToolName(t *testing.T) { + t.Parallel() + + got, err := schema.Translate[CallToolInput](map[string]any{ + "parameters": map[string]any{"tool_name": "search", "query": "weather"}, + }) + require.NoError(t, err) + require.Equal(t, "search", got.ToolName) + require.Equal(t, map[string]any{"query": "weather"}, got.Parameters) +} diff --git a/test/e2e/vmcp_optimizer_test.go b/test/e2e/vmcp_optimizer_test.go index 01fd2e46dd..a4d13f7ddf 100644 --- a/test/e2e/vmcp_optimizer_test.go +++ b/test/e2e/vmcp_optimizer_test.go @@ -8,7 +8,8 @@ // This file adds deeper coverage for RFC THV-0059 Phase 4: // // - Tier-1 find→call round-trip: verifies that find_tool locates the yardstick -// echo tool by description and call_tool invokes it end-to-end. +// echo tool by description and call_tool invokes it end-to-end, both with +// tool_name alongside parameters and with it nested inside them. // - Tier-1 two-backend with conflict resolution: verifies that optimizer // discovers tools from both backends when prefix conflict resolution is active. // - Tier-1 composite + optimizer: verifies that composite tools are indexed by @@ -93,7 +94,7 @@ var _ = Describe("vMCP optimizer", Label("vmcp", "e2e", "optimizer"), func() { BeforeEach(func() { fx.setup("vmcp-opt-roundtrip", "") }) AfterEach(func() { fx.teardown() }) - It("find_tool locates the echo tool and call_tool invokes it end-to-end", func() { + It("find_tool locates the echo tool and call_tool invokes it, flat and nested", func() { By("starting thv vmcp serve with --optimizer") fx.vMCPCmd = e2e.StartLongRunningTHVCommand(fx.cfg, "vmcp", "serve", @@ -133,6 +134,20 @@ var _ = Describe("vMCP optimizer", Label("vmcp", "e2e", "optimizer"), func() { Expect(callResult.Content).ToNot(BeEmpty(), "call_tool must return content") Expect(mcp.GetTextFromContent(callResult.Content[0])).To(ContainSubstring("hellooptimizer"), "echo tool must return the input message") + + By("invoking the same tool with tool_name nested inside parameters") + nestedResult, err := mcpClient.CallTool(ctx, "call_tool", map[string]any{ + "parameters": map[string]any{ + "tool_name": echoToolName, + "input": "hellonested", + }, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(nestedResult.IsError).To(BeFalse(), + "call_tool must accept a tool_name nested inside parameters") + Expect(nestedResult.Content).ToNot(BeEmpty(), "call_tool must return content") + Expect(mcp.GetTextFromContent(nestedResult.Content[0])).To(ContainSubstring("hellonested"), + "the nested tool_name must be hoisted and the remaining parameters passed through") }) })