add base url override option

allows dynamic discovery of base url

Author: birdayz

README.md

@@ -114,6 +114,21 @@ testdatamcp.ForwardToConnectTestServiceClient(mcpServer, myConnectClient)
 
 This directly connects the MCP handler to the connectrpc client, requiring zero boilerplate.
 
+### URL Override Support
+
+All generated MCP functions support dynamic URL overrides via functional options. This allows LLMs to specify custom base URLs for different service environments:
+
+```go
+// Enable URL override with custom field name and description
+option := runtime.WithBaseURLProperty("base_url", "Base URL for the API")
+
+// Use with any generated function
+testdatamcp.RegisterTestServiceHandler(mcpServer, &srv, option)
+testdatamcp.ForwardToTestServiceClient(mcpServer, client, option)
+```
+
+When enabled, the tool schema automatically includes the URL field, and LLMs can provide URLs like `https://api.example.com` which get parsed and made available in the request context via `runtime.URLOverrideKey{}`.
+
 ## LLM Provider Compatibility
 
 The generator now creates both standard MCP and OpenAI-compatible handlers automatically. You can choose which to use at runtime:

pkg/generator/base_url_override_integration_test.go

@@ -0,0 +1,101 @@
+// Copyright 2025 Redpanda Data, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package generator
+
+import (
+	"encoding/json"
+	"testing"
+
+	. "github.com/onsi/gomega"
+	"github.com/redpanda-data/protoc-gen-go-mcp/pkg/runtime"
+	testdatamcp "github.com/redpanda-data/protoc-gen-go-mcp/pkg/testdata/gen/go/testdata/testdatamcp"
+)
+
+func TestURLOverrideSchemaModification(t *testing.T) {
+	g := NewWithT(t)
+
+	// Test that the AddURLFieldToTool function properly modifies the CreateItem tool schema
+	originalTool := testdatamcp.TestService_CreateItemTool
+
+	// Parse the original schema to verify it doesn't have base_url field
+	var originalSchema map[string]interface{}
+	err := json.Unmarshal(originalTool.RawInputSchema, &originalSchema)
+	g.Expect(err).ToNot(HaveOccurred())
+
+	originalProperties := originalSchema["properties"].(map[string]interface{})
+	g.Expect(originalProperties).ToNot(HaveKey("base_url"))
+
+	// Add base_url field to the tool
+	modifiedTool := runtime.AddURLFieldToTool(originalTool, "base_url", "Base URL for the API")
+
+	// Parse the modified schema
+	var modifiedSchema map[string]interface{}
+	err = json.Unmarshal(modifiedTool.RawInputSchema, &modifiedSchema)
+	g.Expect(err).ToNot(HaveOccurred())
+
+	// Verify the base_url field was added
+	modifiedProperties := modifiedSchema["properties"].(map[string]interface{})
+	g.Expect(modifiedProperties).To(HaveKey("base_url"))
+
+	urlField := modifiedProperties["base_url"].(map[string]interface{})
+	g.Expect(urlField["type"]).To(Equal("string"))
+	g.Expect(urlField["format"]).To(Equal("uri"))
+	g.Expect(urlField["description"]).To(Equal("Base URL for the API"))
+
+	// Verify original fields are still there
+	g.Expect(modifiedProperties).To(HaveKey("name"))
+	g.Expect(modifiedProperties).To(HaveKey("description"))
+	g.Expect(modifiedProperties).To(HaveKey("labels"))
+	g.Expect(modifiedProperties).To(HaveKey("tags"))
+}
+
+func TestURLOverrideSchemaModificationWithCustomField(t *testing.T) {
+	g := NewWithT(t)
+
+	// Test that the AddURLFieldToTool function properly modifies the CreateItem tool schema with custom field name
+	originalTool := testdatamcp.TestService_CreateItemTool
+
+	// Parse the original schema to verify it doesn't have custom field
+	var originalSchema map[string]interface{}
+	err := json.Unmarshal(originalTool.RawInputSchema, &originalSchema)
+	g.Expect(err).ToNot(HaveOccurred())
+
+	originalProperties := originalSchema["properties"].(map[string]interface{})
+	g.Expect(originalProperties).ToNot(HaveKey("api_url"))
+
+	// Add custom field to the tool
+	modifiedTool := runtime.AddURLFieldToTool(originalTool, "api_url", "Custom API endpoint URL")
+
+	// Parse the modified schema
+	var modifiedSchema map[string]interface{}
+	err = json.Unmarshal(modifiedTool.RawInputSchema, &modifiedSchema)
+	g.Expect(err).ToNot(HaveOccurred())
+
+	// Verify the custom field was added
+	modifiedProperties := modifiedSchema["properties"].(map[string]interface{})
+	g.Expect(modifiedProperties).To(HaveKey("api_url"))
+	g.Expect(modifiedProperties).ToNot(HaveKey("base_url")) // Should not have default base_url field
+
+	customField := modifiedProperties["api_url"].(map[string]interface{})
+	g.Expect(customField["type"]).To(Equal("string"))
+	g.Expect(customField["format"]).To(Equal("uri"))
+	g.Expect(customField["description"]).To(Equal("Custom API endpoint URL"))
+
+	// Verify original fields are still there
+	g.Expect(modifiedProperties).To(HaveKey("name"))
+	g.Expect(modifiedProperties).To(HaveKey("description"))
+	g.Expect(modifiedProperties).To(HaveKey("labels"))
+	g.Expect(modifiedProperties).To(HaveKey("tags"))
+}

pkg/generator/generator.go

@@ -65,6 +65,7 @@ import (
   "connectrpc.com/connect"
   grpc "google.golang.org/grpc"
   "github.com/redpanda-data/protoc-gen-go-mcp/pkg/runtime"
+  "net/url"
 )
 
 
@@ -88,13 +89,35 @@ type {{$serviceName}}Server interface {
 
 {{- range $key, $val := .Services }}
 // Register{{$key}}Handler registers standard MCP handlers for {{$key}}
-func Register{{$key}}Handler(s *mcpserver.MCPServer, srv {{$key}}Server) {
+func Register{{$key}}Handler(s *mcpserver.MCPServer, srv {{$key}}Server, opts ...runtime.Option) {
+  config := runtime.NewConfig()
+  for _, opt := range opts {
+    opt(config)
+  }
+
   {{- range $tool_name, $tool_val := $val }}
-  s.AddTool({{$key}}_{{$tool_name}}Tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+  {{$tool_name}}Tool := {{$key}}_{{$tool_name}}Tool
+  // Add URL field to schema if ExtractURL is enabled
+  if config.ExtractURL {
+    {{$tool_name}}Tool = runtime.AddURLFieldToTool({{$tool_name}}Tool, config.URLFieldName, config.URLDescription)
+  }
+  
+  s.AddTool({{$tool_name}}Tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
     var req {{$tool_val.RequestType}}
 
     message := request.Params.Arguments
 
+    // Extract URL if option is enabled
+    if config.ExtractURL {
+      if urlVal, ok := message[config.URLFieldName]; ok {
+        if urlStr, ok := urlVal.(string); ok {
+          if parsedURL, err := url.Parse(urlStr); err == nil {
+            ctx = context.WithValue(ctx, runtime.URLOverrideKey{}, parsedURL)
+          }
+        }
+      }
+    }
+
     marshaled, err := json.Marshal(message)
     if err != nil {
       return nil, err
@@ -120,12 +143,35 @@ func Register{{$key}}Handler(s *mcpserver.MCPServer, srv {{$key}}Server) {
 }
 
 // Register{{$key}}HandlerOpenAI registers OpenAI-compatible MCP handlers for {{$key}}
-func Register{{$key}}HandlerOpenAI(s *mcpserver.MCPServer, srv {{$key}}Server) {
+func Register{{$key}}HandlerOpenAI(s *mcpserver.MCPServer, srv {{$key}}Server, opts ...runtime.Option) {
+  config := runtime.NewConfig()
+  for _, opt := range opts {
+    opt(config)
+  }
+
   {{- range $tool_name, $tool_val := $val }}
-  s.AddTool({{$key}}_{{$tool_name}}ToolOpenAI, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+  {{$tool_name}}ToolOpenAI := {{$key}}_{{$tool_name}}ToolOpenAI
+  // Add URL field to schema if ExtractURL is enabled
+  if config.ExtractURL {
+    {{$tool_name}}ToolOpenAI = runtime.AddURLFieldToTool({{$tool_name}}ToolOpenAI, config.URLFieldName, config.URLDescription)
+  }
+  
+  s.AddTool({{$tool_name}}ToolOpenAI, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
     var req {{$tool_val.RequestType}}
 
     message := request.Params.Arguments
+
+    // Extract URL if option is enabled
+    if config.ExtractURL {
+      if urlVal, ok := message[config.URLFieldName]; ok {
+        if urlStr, ok := urlVal.(string); ok {
+          if parsedURL, err := url.Parse(urlStr); err == nil {
+            ctx = context.WithValue(ctx, runtime.URLOverrideKey{}, parsedURL)
+          }
+        }
+      }
+    }
+
     runtime.FixOpenAI(req.ProtoReflect().Descriptor(), message)
 
     marshaled, err := json.Marshal(message)
@@ -153,14 +199,14 @@ func Register{{$key}}HandlerOpenAI(s *mcpserver.MCPServer, srv {{$key}}Server) {
 }
 
 // Register{{$key}}HandlerWithProvider registers handlers for the specified LLM provider
-func Register{{$key}}HandlerWithProvider(s *mcpserver.MCPServer, srv {{$key}}Server, provider runtime.LLMProvider) {
+func Register{{$key}}HandlerWithProvider(s *mcpserver.MCPServer, srv {{$key}}Server, provider runtime.LLMProvider, opts ...runtime.Option) {
   switch provider {
   case runtime.LLMProviderOpenAI:
-    Register{{$key}}HandlerOpenAI(s, srv)
+    Register{{$key}}HandlerOpenAI(s, srv, opts...)
   case runtime.LLMProviderStandard:
     fallthrough
   default:
-    Register{{$key}}Handler(s, srv)
+    Register{{$key}}Handler(s, srv, opts...)
   }
 }
 {{- end }}
@@ -186,13 +232,35 @@ type Connect{{$serviceName}}Client interface {
 
 {{- range $key, $val := .Services }}
 // ForwardToConnect{{$key}}Client registers a connectrpc client, to forward MCP calls to it.
-func ForwardToConnect{{$key}}Client(s *mcpserver.MCPServer, client Connect{{$key}}Client) {
+func ForwardToConnect{{$key}}Client(s *mcpserver.MCPServer, client Connect{{$key}}Client, opts ...runtime.Option) {
+  config := runtime.NewConfig()
+  for _, opt := range opts {
+    opt(config)
+  }
+
   {{- range $tool_name, $tool_val := $val }}
-  s.AddTool({{$key}}_{{$tool_name}}Tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+  {{$tool_name}}Tool := {{$key}}_{{$tool_name}}Tool
+  // Add URL field to schema if ExtractURL is enabled
+  if config.ExtractURL {
+    {{$tool_name}}Tool = runtime.AddURLFieldToTool({{$tool_name}}Tool, config.URLFieldName, config.URLDescription)
+  }
+  
+  s.AddTool({{$tool_name}}Tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
     var req {{$tool_val.RequestType}}
 
     message := request.Params.Arguments
 
+    // Extract URL if option is enabled
+    if config.ExtractURL {
+      if urlVal, ok := message[config.URLFieldName]; ok {
+        if urlStr, ok := urlVal.(string); ok {
+          if parsedURL, err := url.Parse(urlStr); err == nil {
+            ctx = context.WithValue(ctx, runtime.URLOverrideKey{}, parsedURL)
+          }
+        }
+      }
+    }
+
     marshaled, err := json.Marshal(message)
     if err != nil {
       return nil, err
@@ -219,13 +287,35 @@ func ForwardToConnect{{$key}}Client(s *mcpserver.MCPServer, client Connect{{$key
 
 {{- range $key, $val := .Services }}
 // ForwardTo{{$key}}Client registers a gRPC client, to forward MCP calls to it.
-func ForwardTo{{$key}}Client(s *mcpserver.MCPServer, client {{$key}}Client) {
+func ForwardTo{{$key}}Client(s *mcpserver.MCPServer, client {{$key}}Client, opts ...runtime.Option) {
+  config := runtime.NewConfig()
+  for _, opt := range opts {
+    opt(config)
+  }
+
   {{- range $tool_name, $tool_val := $val }}
-  s.AddTool({{$key}}_{{$tool_name}}Tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+  {{$tool_name}}Tool := {{$key}}_{{$tool_name}}Tool
+  // Add URL field to schema if ExtractURL is enabled
+  if config.ExtractURL {
+    {{$tool_name}}Tool = runtime.AddURLFieldToTool({{$tool_name}}Tool, config.URLFieldName, config.URLDescription)
+  }
+  
+  s.AddTool({{$tool_name}}Tool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
     var req {{$tool_val.RequestType}}
 
     message := request.Params.Arguments
 
+    // Extract URL if option is enabled
+    if config.ExtractURL {
+      if urlVal, ok := message[config.URLFieldName]; ok {
+        if urlStr, ok := urlVal.(string); ok {
+          if parsedURL, err := url.Parse(urlStr); err == nil {
+            ctx = context.WithValue(ctx, runtime.URLOverrideKey{}, parsedURL)
+          }
+        }
+      }
+    }
+
     marshaled, err := json.Marshal(message)
     if err != nil {
       return nil, err

pkg/runtime/base_url_override.go

@@ -0,0 +1,74 @@
+package runtime
+
+import (
+	"encoding/json"
+
+	"github.com/mark3labs/mcp-go/mcp"
+)
+
+// URLOverrideKey is the context key for storing the parsed URL override
+type URLOverrideKey struct{}
+
+// Option defines functional options for MCP functions
+type Option func(*config)
+
+type config struct {
+	ExtractURL     bool
+	URLFieldName   string
+	URLDescription string
+}
+
+// WithBaseURLProperty enables extracting a URL override from the specified field name
+// in request arguments and setting the parsed URL in context.
+func WithBaseURLProperty(fieldName, description string) Option {
+	return func(c *config) {
+		c.ExtractURL = true
+		c.URLFieldName = fieldName
+		c.URLDescription = description
+	}
+}
+
+// NewConfig creates a new config instance
+func NewConfig() *config {
+	return &config{}
+}
+
+// AddURLFieldToTool modifies a tool's schema to include an optional URL override field
+func AddURLFieldToTool(tool mcp.Tool, fieldName, description string) mcp.Tool {
+	// Parse the existing schema
+	var schema map[string]interface{}
+	if err := json.Unmarshal(tool.RawInputSchema, &schema); err != nil {
+		// If we can't parse the schema, return the original tool
+		return tool
+	}
+
+	// Add URL field to properties
+	if properties, ok := schema["properties"].(map[string]interface{}); ok {
+		properties[fieldName] = map[string]interface{}{
+			"type":        "string",
+			"format":      "uri",
+			"description": description,
+		}
+	} else {
+		// If no properties exist, create them
+		schema["properties"] = map[string]interface{}{
+			fieldName: map[string]interface{}{
+				"type":        "string",
+				"format":      "uri",
+				"description": description,
+			},
+		}
+	}
+
+	// Marshal the modified schema back
+	modifiedSchema, err := json.Marshal(schema)
+	if err != nil {
+		// If marshaling fails, return the original tool
+		return tool
+	}
+
+	// Create a new tool with the modified schema
+	modifiedTool := tool
+	modifiedTool.RawInputSchema = json.RawMessage(modifiedSchema)
+	return modifiedTool
+}