Add tests for RegisterToolsWithMiddleware and RegisterAllWithMiddleware

Tests verify that middleware:
- Receives tool name at registration time (outer function called)
- Can read context values passed through from tool handlers
- Can detect IsError on result without Go error (ghErrors pattern)
- Can modify the result before returning
- Is applied to all tools when using RegisterAllWithMiddleware
- Handles nil middleware gracefully

Author: SamMorrowDrums

pkg/inventory/registry_test.go

@@ -1643,3 +1643,262 @@ func TestFilteringOrder(t *testing.T) {
 		}
 	}
 }
+
+// contextKey is a custom type for context keys to avoid collisions
+type contextKey string
+
+const testErrorKey contextKey = "test_error"
+
+func TestRegisterToolsWithMiddleware(t *testing.T) {
+	// Test that middleware can:
+	// 1. Access the tool name
+	// 2. Read context values set by the tool handler
+	// 3. Inspect and modify the result
+	// 4. Access IsError on the result
+
+	t.Run("middleware receives tool name and can wrap handler", func(t *testing.T) {
+		var capturedToolNameAtRegistration string
+		var capturedToolNameAtInvocation string
+		var handlerWasCalled bool
+
+		tool := NewServerToolFromHandler(
+			mcp.Tool{
+				Name:        "test_tool",
+				Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true},
+				InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
+			},
+			testToolsetMetadata("test"),
+			func(_ any) mcp.ToolHandler {
+				return func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+					handlerWasCalled = true
+					return &mcp.CallToolResult{
+						Content: []mcp.Content{&mcp.TextContent{Text: "success"}},
+					}, nil
+				}
+			},
+		)
+
+		middleware := func(toolName string, handler mcp.ToolHandler) mcp.ToolHandler {
+			// This outer function is called at registration time
+			capturedToolNameAtRegistration = toolName
+			return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+				// This inner function is called at invocation time
+				capturedToolNameAtInvocation = toolName
+				return handler(ctx, req)
+			}
+		}
+
+		server := mcp.NewServer(&mcp.Implementation{Name: "test"}, nil)
+		reg := NewBuilder().SetTools([]ServerTool{tool}).WithToolsets([]string{"all"}).Build()
+		reg.RegisterToolsWithMiddleware(context.Background(), server, nil, middleware)
+
+		// Verify the middleware is applied at registration time (wrapping)
+		// by checking that the tool name was captured in the outer function
+		if capturedToolNameAtRegistration != "test_tool" {
+			t.Errorf("Expected middleware to capture tool name at registration 'test_tool', got %q", capturedToolNameAtRegistration)
+		}
+
+		// The inner function (invocation-time capture) should not have been called yet
+		if capturedToolNameAtInvocation != "" {
+			t.Error("Inner middleware function should not be called during registration")
+		}
+
+		// Handler shouldn't be called until tool is invoked
+		if handlerWasCalled {
+			t.Error("Handler should not be called during registration")
+		}
+	})
+
+	t.Run("middleware can read context values set by tool", func(t *testing.T) {
+		var middlewareSeenError string
+
+		tool := NewServerToolFromHandler(
+			mcp.Tool{
+				Name:        "error_tool",
+				Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true},
+				InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
+			},
+			testToolsetMetadata("test"),
+			func(_ any) mcp.ToolHandler {
+				return func(ctx context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+					// Simulate storing an error in context (like ghErrors does)
+					// The context is passed by value, but if we use a pointer in the context
+					// the middleware can see modifications
+					if ptr, ok := ctx.Value(testErrorKey).(*string); ok && ptr != nil {
+						*ptr = "github_api_error: 404 not found"
+					}
+					return &mcp.CallToolResult{
+						Content: []mcp.Content{&mcp.TextContent{Text: "error occurred"}},
+						IsError: true,
+					}, nil // Note: returning nil error, but IsError is true
+				}
+			},
+		)
+
+		middleware := func(_ string, handler mcp.ToolHandler) mcp.ToolHandler {
+			return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+				// Set up a pointer in context that the handler can write to
+				errorHolder := ""
+				ctx = context.WithValue(ctx, testErrorKey, &errorHolder)
+
+				result, err := handler(ctx, req)
+
+				// Read what the handler wrote
+				middlewareSeenError = errorHolder
+
+				return result, err
+			}
+		}
+
+		server := mcp.NewServer(&mcp.Implementation{Name: "test"}, nil)
+		reg := NewBuilder().SetTools([]ServerTool{tool}).WithToolsets([]string{"all"}).Build()
+		reg.RegisterToolsWithMiddleware(context.Background(), server, nil, middleware)
+
+		// Simulate calling the tool - recreate the middleware wrapper
+		handler := tool.Handler(nil)
+		wrappedHandler := middleware("error_tool", handler)
+
+		result, err := wrappedHandler(context.Background(), &mcp.CallToolRequest{})
+
+		if err != nil {
+			t.Errorf("Expected no Go error, got %v", err)
+		}
+		if !result.IsError {
+			t.Error("Expected result.IsError to be true")
+		}
+		if middlewareSeenError != "github_api_error: 404 not found" {
+			t.Errorf("Middleware didn't see context error, got: %q", middlewareSeenError)
+		}
+	})
+
+	t.Run("middleware can detect IsError without Go error", func(t *testing.T) {
+		// This tests the exact pattern used by ghErrors.NewGitHubAPIStatusErrorResponse
+		// which returns (result_with_IsError_true, nil)
+
+		var middlewareDetectedError bool
+		var middlewareDetectedGoError bool
+
+		tool := NewServerToolFromHandler(
+			mcp.Tool{
+				Name:        "status_error_tool",
+				Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true},
+				InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
+			},
+			testToolsetMetadata("test"),
+			func(_ any) mcp.ToolHandler {
+				return func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+					// This is exactly how ghErrors.NewGitHubAPIStatusErrorResponse returns
+					return &mcp.CallToolResult{
+						Content: []mcp.Content{&mcp.TextContent{Text: "unexpected status 404: not found"}},
+						IsError: true,
+					}, nil // No Go error!
+				}
+			},
+		)
+
+		middleware := func(_ string, handler mcp.ToolHandler) mcp.ToolHandler {
+			return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+				result, err := handler(ctx, req)
+
+				// Check both error indicators
+				middlewareDetectedGoError = (err != nil)
+				middlewareDetectedError = (result != nil && result.IsError)
+
+				return result, err
+			}
+		}
+
+		handler := tool.Handler(nil)
+		wrappedHandler := middleware("status_error_tool", handler)
+
+		_, _ = wrappedHandler(context.Background(), &mcp.CallToolRequest{})
+
+		if middlewareDetectedGoError {
+			t.Error("Should NOT detect Go error (it's nil)")
+		}
+		if !middlewareDetectedError {
+			t.Error("SHOULD detect error via result.IsError")
+		}
+	})
+
+	t.Run("middleware can modify result", func(t *testing.T) {
+		tool := NewServerToolFromHandler(
+			mcp.Tool{
+				Name:        "modifiable_tool",
+				Annotations: &mcp.ToolAnnotations{ReadOnlyHint: true},
+				InputSchema: json.RawMessage(`{"type":"object","properties":{}}`),
+			},
+			testToolsetMetadata("test"),
+			func(_ any) mcp.ToolHandler {
+				return func(_ context.Context, _ *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+					return &mcp.CallToolResult{
+						Content: []mcp.Content{&mcp.TextContent{Text: "original"}},
+					}, nil
+				}
+			},
+		)
+
+		middleware := func(_ string, handler mcp.ToolHandler) mcp.ToolHandler {
+			return func(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+				result, err := handler(ctx, req)
+
+				// Middleware can modify the result
+				if result != nil {
+					result.Content = []mcp.Content{&mcp.TextContent{Text: "modified by middleware"}}
+				}
+
+				return result, err
+			}
+		}
+
+		handler := tool.Handler(nil)
+		wrappedHandler := middleware("modifiable_tool", handler)
+
+		result, _ := wrappedHandler(context.Background(), &mcp.CallToolRequest{})
+
+		if result == nil {
+			t.Fatal("Expected result")
+		}
+		textContent, ok := result.Content[0].(*mcp.TextContent)
+		if !ok {
+			t.Fatal("Expected TextContent")
+		}
+		if textContent.Text != "modified by middleware" {
+			t.Errorf("Expected modified text, got: %s", textContent.Text)
+		}
+	})
+
+	t.Run("RegisterAllWithMiddleware applies middleware to tools", func(t *testing.T) {
+		middlewareCallCount := 0
+
+		tools := []ServerTool{
+			mockTool("tool1", "toolset1", true),
+			mockTool("tool2", "toolset1", true),
+		}
+
+		middleware := func(_ string, handler mcp.ToolHandler) mcp.ToolHandler {
+			middlewareCallCount++
+			return handler
+		}
+
+		server := mcp.NewServer(&mcp.Implementation{Name: "test"}, nil)
+		reg := NewBuilder().SetTools(tools).WithToolsets([]string{"all"}).Build()
+		reg.RegisterAllWithMiddleware(context.Background(), server, nil, middleware)
+
+		// Middleware should be called once per tool during registration
+		if middlewareCallCount != 2 {
+			t.Errorf("Expected middleware to be called 2 times, got %d", middlewareCallCount)
+		}
+	})
+
+	t.Run("nil middleware is handled gracefully", func(_ *testing.T) {
+		tool := mockTool("test_tool", "toolset1", true)
+
+		server := mcp.NewServer(&mcp.Implementation{Name: "test"}, nil)
+		reg := NewBuilder().SetTools([]ServerTool{tool}).WithToolsets([]string{"all"}).Build()
+
+		// Should not panic with nil middleware
+		reg.RegisterToolsWithMiddleware(context.Background(), server, nil, nil)
+		reg.RegisterAllWithMiddleware(context.Background(), server, nil, nil)
+	})
+}