From f1551e4c2f3d66ddc565da2d79c8031bf82f6119 Mon Sep 17 00:00:00 2001 From: triepod-ai Date: Sun, 28 Dec 2025 06:27:40 -0600 Subject: [PATCH] feat: add tool annotations support via runtime options MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add support for MCP tool annotations (readOnlyHint, destructiveHint, idempotentHint, openWorldHint, title) through a runtime options pattern consistent with the existing WithExtraProperties() approach. This enables users to specify annotation hints at tool registration time: ```go tool := runtime.ApplyToolAnnotations( generatedTool, runtime.WithToolAnnotations(runtime.ToolAnnotationConfig{ Title: "Get Item", ReadOnlyHint: boolPtr(true), }), ) ``` Changes: - Add WithToolAnnotations() runtime option - Add ApplyToolAnnotations() helper function - Add ApplyOptions() for combined annotations and extra properties - Add comprehensive unit tests (9 test cases) - Add examples/with-annotations demonstrating usage 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- examples/with-annotations/main.go | 165 +++++++++++++++++ pkg/runtime/annotations.go | 91 +++++++++ pkg/runtime/annotations_test.go | 298 ++++++++++++++++++++++++++++++ pkg/runtime/extra_properties.go | 1 + 4 files changed, 555 insertions(+) create mode 100644 examples/with-annotations/main.go create mode 100644 pkg/runtime/annotations.go create mode 100644 pkg/runtime/annotations_test.go diff --git a/examples/with-annotations/main.go b/examples/with-annotations/main.go new file mode 100644 index 0000000..767df3a --- /dev/null +++ b/examples/with-annotations/main.go @@ -0,0 +1,165 @@ +// 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. + +// This example demonstrates how to add MCP tool annotations to generated tools. +// Tool annotations provide semantic hints to LLMs and MCP clients about tool behavior. +package main + +import ( + "context" + "fmt" + + "github.com/mark3labs/mcp-go/mcp" + "github.com/mark3labs/mcp-go/server" + "github.com/redpanda-data/protoc-gen-go-mcp/pkg/runtime" + testdata "github.com/redpanda-data/protoc-gen-go-mcp/pkg/testdata/gen/go/testdata" + "github.com/redpanda-data/protoc-gen-go-mcp/pkg/testdata/gen/go/testdata/testdatamcp" +) + +// Helper function to create bool pointers +func boolPtr(b bool) *bool { + return &b +} + +func main() { + // Create MCP server + s := server.NewMCPServer( + "Example with tool annotations", + "1.0.0", + ) + + srv := testServer{} + + // Register tools with annotations applied manually + // This shows how to add semantic hints to generated tools + + // GetItem is a read-only operation - mark it accordingly + getItemTool := runtime.ApplyToolAnnotations( + testdatamcp.TestService_GetItemTool, + runtime.WithToolAnnotations(runtime.ToolAnnotationConfig{ + Title: "Get Item", + ReadOnlyHint: boolPtr(true), + }), + ) + s.AddTool(getItemTool, makeGetItemHandler(&srv)) + + // CreateItem modifies state - mark as destructive + createItemTool := runtime.ApplyToolAnnotations( + testdatamcp.TestService_CreateItemTool, + runtime.WithToolAnnotations(runtime.ToolAnnotationConfig{ + Title: "Create Item", + DestructiveHint: boolPtr(true), + }), + ) + s.AddTool(createItemTool, makeCreateItemHandler(&srv)) + + // ProcessWellKnownTypes - read-only processing + processTypesTool := runtime.ApplyToolAnnotations( + testdatamcp.TestService_ProcessWellKnownTypesTool, + runtime.WithToolAnnotations(runtime.ToolAnnotationConfig{ + Title: "Process Well-Known Types", + ReadOnlyHint: boolPtr(true), + }), + ) + s.AddTool(processTypesTool, makeProcessTypesHandler(&srv)) + + // TestValidation - read-only validation + validationTool := runtime.ApplyToolAnnotations( + testdatamcp.TestService_TestValidationTool, + runtime.WithToolAnnotations(runtime.ToolAnnotationConfig{ + Title: "Test Validation", + ReadOnlyHint: boolPtr(true), + }), + ) + s.AddTool(validationTool, makeValidationHandler(&srv)) + + fmt.Println("Starting server with annotated tools...") + if err := server.ServeStdio(s); err != nil { + fmt.Printf("Server error: %v\n", err) + } +} + +// Handler functions that wrap the service methods + +func makeGetItemHandler(srv *testServer) server.ToolHandlerFunc { + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + id, _ := request.GetArguments()["id"].(string) + resp, err := srv.GetItem(ctx, &testdata.GetItemRequest{Id: id}) + if err != nil { + return nil, err + } + return mcp.NewToolResultText(fmt.Sprintf("Item: %s - %s", resp.Item.Id, resp.Item.Name)), nil + } +} + +func makeCreateItemHandler(srv *testServer) server.ToolHandlerFunc { + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + name, _ := request.GetArguments()["name"].(string) + resp, err := srv.CreateItem(ctx, &testdata.CreateItemRequest{Name: name}) + if err != nil { + return nil, err + } + return mcp.NewToolResultText(fmt.Sprintf("Created item: %s", resp.Id)), nil + } +} + +func makeProcessTypesHandler(srv *testServer) server.ToolHandlerFunc { + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + resp, err := srv.ProcessWellKnownTypes(ctx, &testdata.ProcessWellKnownTypesRequest{}) + if err != nil { + return nil, err + } + return mcp.NewToolResultText(resp.Message), nil + } +} + +func makeValidationHandler(srv *testServer) server.ToolHandlerFunc { + return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) { + resp, err := srv.TestValidation(ctx, &testdata.TestValidationRequest{}) + if err != nil { + return nil, err + } + return mcp.NewToolResultText(fmt.Sprintf("Success: %v - %s", resp.Success, resp.Message)), nil + } +} + +type testServer struct{} + +func (t *testServer) CreateItem(ctx context.Context, in *testdata.CreateItemRequest) (*testdata.CreateItemResponse, error) { + return &testdata.CreateItemResponse{ + Id: "item-123", + }, nil +} + +func (t *testServer) GetItem(ctx context.Context, in *testdata.GetItemRequest) (*testdata.GetItemResponse, error) { + return &testdata.GetItemResponse{ + Item: &testdata.Item{ + Id: in.GetId(), + Name: "Retrieved item", + }, + }, nil +} + +func (t *testServer) ProcessWellKnownTypes(ctx context.Context, in *testdata.ProcessWellKnownTypesRequest) (*testdata.ProcessWellKnownTypesResponse, error) { + return &testdata.ProcessWellKnownTypesResponse{ + Message: "Processed well-known types", + }, nil +} + +func (t *testServer) TestValidation(ctx context.Context, in *testdata.TestValidationRequest) (*testdata.TestValidationResponse, error) { + return &testdata.TestValidationResponse{ + Success: true, + Message: "Validation test completed", + }, nil +} diff --git a/pkg/runtime/annotations.go b/pkg/runtime/annotations.go new file mode 100644 index 0000000..ab6aa44 --- /dev/null +++ b/pkg/runtime/annotations.go @@ -0,0 +1,91 @@ +package runtime + +import ( + "github.com/mark3labs/mcp-go/mcp" +) + +// ToolAnnotationConfig defines annotation hints for a tool. +// These annotations help LLMs and MCP clients understand tool behavior +// and make better decisions about when and how to use tools. +type ToolAnnotationConfig struct { + // Title is a human-readable title for the tool + Title string + + // ReadOnlyHint indicates that the tool does not modify its environment. + // If true, the tool only reads data without making changes. + ReadOnlyHint *bool + + // DestructiveHint indicates that the tool may perform destructive updates + // to its environment. If false, the tool performs only additive updates. + DestructiveHint *bool + + // IdempotentHint indicates that calling the tool repeatedly with the same + // arguments will have no additional effect on its environment. + IdempotentHint *bool + + // OpenWorldHint indicates that the tool may interact with an "open world" + // of external entities. If false, the tool's domain of interaction is closed. + OpenWorldHint *bool +} + +// WithToolAnnotations adds tool annotations that provide semantic hints +// about tool behavior to MCP clients. +func WithToolAnnotations(annotations ToolAnnotationConfig) Option { + return func(c *config) { + c.Annotations = &annotations + } +} + +// ApplyToolAnnotations applies annotation options to a tool, returning +// a modified tool with the annotations set. +func ApplyToolAnnotations(tool mcp.Tool, opts ...Option) mcp.Tool { + cfg := &config{} + for _, opt := range opts { + opt(cfg) + } + + if cfg.Annotations == nil { + return tool + } + + // Create a copy of the tool to avoid modifying the original + modifiedTool := tool + + // Build the ToolAnnotation from config + modifiedTool.Annotations = mcp.ToolAnnotation{ + Title: cfg.Annotations.Title, + ReadOnlyHint: cfg.Annotations.ReadOnlyHint, + DestructiveHint: cfg.Annotations.DestructiveHint, + IdempotentHint: cfg.Annotations.IdempotentHint, + OpenWorldHint: cfg.Annotations.OpenWorldHint, + } + + return modifiedTool +} + +// ApplyOptions applies all options (including annotations and extra properties) +// to a tool in a single call. +func ApplyOptions(tool mcp.Tool, opts ...Option) mcp.Tool { + cfg := &config{} + for _, opt := range opts { + opt(cfg) + } + + // Apply annotations if present + if cfg.Annotations != nil { + tool.Annotations = mcp.ToolAnnotation{ + Title: cfg.Annotations.Title, + ReadOnlyHint: cfg.Annotations.ReadOnlyHint, + DestructiveHint: cfg.Annotations.DestructiveHint, + IdempotentHint: cfg.Annotations.IdempotentHint, + OpenWorldHint: cfg.Annotations.OpenWorldHint, + } + } + + // Apply extra properties if present + if len(cfg.ExtraProperties) > 0 { + tool = AddExtraPropertiesToTool(tool, cfg.ExtraProperties) + } + + return tool +} diff --git a/pkg/runtime/annotations_test.go b/pkg/runtime/annotations_test.go new file mode 100644 index 0000000..e6d83f8 --- /dev/null +++ b/pkg/runtime/annotations_test.go @@ -0,0 +1,298 @@ +// 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 runtime + +import ( + "encoding/json" + "testing" + + "github.com/mark3labs/mcp-go/mcp" + . "github.com/onsi/gomega" +) + +// Helper function to create bool pointers +func boolPtr(b bool) *bool { + return &b +} + +func TestWithToolAnnotations_Creation(t *testing.T) { + g := NewWithT(t) + + // Test with all fields set + cfg := ToolAnnotationConfig{ + Title: "Test Tool", + ReadOnlyHint: boolPtr(true), + DestructiveHint: boolPtr(false), + IdempotentHint: boolPtr(true), + OpenWorldHint: boolPtr(false), + } + + opt := WithToolAnnotations(cfg) + g.Expect(opt).ToNot(BeNil()) + + // Apply the option to a config + c := &config{} + opt(c) + + g.Expect(c.Annotations).ToNot(BeNil()) + g.Expect(c.Annotations.Title).To(Equal("Test Tool")) + g.Expect(*c.Annotations.ReadOnlyHint).To(BeTrue()) + g.Expect(*c.Annotations.DestructiveHint).To(BeFalse()) + g.Expect(*c.Annotations.IdempotentHint).To(BeTrue()) + g.Expect(*c.Annotations.OpenWorldHint).To(BeFalse()) +} + +func TestWithToolAnnotations_PartialFields(t *testing.T) { + g := NewWithT(t) + + // Test with only Title set + cfg := ToolAnnotationConfig{ + Title: "Read Only Query", + } + + opt := WithToolAnnotations(cfg) + c := &config{} + opt(c) + + g.Expect(c.Annotations).ToNot(BeNil()) + g.Expect(c.Annotations.Title).To(Equal("Read Only Query")) + g.Expect(c.Annotations.ReadOnlyHint).To(BeNil()) + g.Expect(c.Annotations.DestructiveHint).To(BeNil()) +} + +func TestApplyToolAnnotations_BasicApplication(t *testing.T) { + g := NewWithT(t) + + // Create a tool without annotations + originalSchema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, + } + schemaBytes, err := json.Marshal(originalSchema) + g.Expect(err).ToNot(HaveOccurred()) + + tool := mcp.Tool{ + Name: "test_tool", + Description: "A test tool", + RawInputSchema: json.RawMessage(schemaBytes), + } + + // Apply annotations + result := ApplyToolAnnotations(tool, WithToolAnnotations(ToolAnnotationConfig{ + Title: "Test Tool", + ReadOnlyHint: boolPtr(true), + })) + + // Verify annotations are set + g.Expect(result.Annotations.Title).To(Equal("Test Tool")) + g.Expect(result.Annotations.ReadOnlyHint).ToNot(BeNil()) + g.Expect(*result.Annotations.ReadOnlyHint).To(BeTrue()) +} + +func TestApplyToolAnnotations_NoOptions(t *testing.T) { + g := NewWithT(t) + + // Create a tool + originalSchema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, + } + schemaBytes, err := json.Marshal(originalSchema) + g.Expect(err).ToNot(HaveOccurred()) + + tool := mcp.Tool{ + Name: "test_tool", + Description: "A test tool", + RawInputSchema: json.RawMessage(schemaBytes), + } + + // Apply no options + result := ApplyToolAnnotations(tool) + + // Verify tool is unchanged (no annotations added - empty struct) + g.Expect(result.Annotations.Title).To(BeEmpty()) + g.Expect(result.Annotations.ReadOnlyHint).To(BeNil()) + g.Expect(result.Name).To(Equal("test_tool")) + g.Expect(result.Description).To(Equal("A test tool")) +} + +func TestApplyToolAnnotations_PreservesOtherFields(t *testing.T) { + g := NewWithT(t) + + // Create a tool with all fields set + originalSchema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{ + "type": "string", + }, + }, + "required": []string{"query"}, + } + schemaBytes, err := json.Marshal(originalSchema) + g.Expect(err).ToNot(HaveOccurred()) + + tool := mcp.Tool{ + Name: "search_tool", + Description: "Search for items", + RawInputSchema: json.RawMessage(schemaBytes), + } + + // Apply annotations + result := ApplyToolAnnotations(tool, WithToolAnnotations(ToolAnnotationConfig{ + Title: "Search Tool", + ReadOnlyHint: boolPtr(true), + })) + + // Verify original fields are preserved + g.Expect(result.Name).To(Equal("search_tool")) + g.Expect(result.Description).To(Equal("Search for items")) + + // Verify schema is unchanged + var resultSchema map[string]interface{} + err = json.Unmarshal(result.RawInputSchema, &resultSchema) + g.Expect(err).ToNot(HaveOccurred()) + g.Expect(resultSchema["type"]).To(Equal("object")) + + properties := resultSchema["properties"].(map[string]interface{}) + g.Expect(properties).To(HaveKey("query")) +} + +func TestApplyToolAnnotations_AllHintTypes(t *testing.T) { + g := NewWithT(t) + + tool := mcp.Tool{ + Name: "full_annotated_tool", + Description: "A tool with all annotations", + } + + // Apply all annotation types + result := ApplyToolAnnotations(tool, WithToolAnnotations(ToolAnnotationConfig{ + Title: "Full Annotated Tool", + ReadOnlyHint: boolPtr(false), + DestructiveHint: boolPtr(true), + IdempotentHint: boolPtr(false), + OpenWorldHint: boolPtr(true), + })) + + // Verify all annotations are set correctly + g.Expect(result.Annotations.Title).To(Equal("Full Annotated Tool")) + g.Expect(*result.Annotations.ReadOnlyHint).To(BeFalse()) + g.Expect(*result.Annotations.DestructiveHint).To(BeTrue()) + g.Expect(*result.Annotations.IdempotentHint).To(BeFalse()) + g.Expect(*result.Annotations.OpenWorldHint).To(BeTrue()) +} + +func TestApplyOptions_CombinedAnnotationsAndExtraProperties(t *testing.T) { + g := NewWithT(t) + + // Create a tool with a basic schema + originalSchema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "query": map[string]interface{}{ + "type": "string", + }, + }, + } + schemaBytes, err := json.Marshal(originalSchema) + g.Expect(err).ToNot(HaveOccurred()) + + tool := mcp.Tool{ + Name: "api_tool", + Description: "A tool that calls an API", + RawInputSchema: json.RawMessage(schemaBytes), + } + + // Apply both annotations and extra properties + result := ApplyOptions(tool, + WithToolAnnotations(ToolAnnotationConfig{ + Title: "API Tool", + OpenWorldHint: boolPtr(true), + }), + WithExtraProperties(ExtraProperty{ + Name: "base_url", + Description: "API base URL", + Required: true, + }), + ) + + // Verify annotations are set + g.Expect(result.Annotations.Title).To(Equal("API Tool")) + g.Expect(*result.Annotations.OpenWorldHint).To(BeTrue()) + + // Verify extra properties are added to schema + var resultSchema map[string]interface{} + err = json.Unmarshal(result.RawInputSchema, &resultSchema) + g.Expect(err).ToNot(HaveOccurred()) + + properties := resultSchema["properties"].(map[string]interface{}) + g.Expect(properties).To(HaveKey("base_url")) + g.Expect(properties).To(HaveKey("query")) +} + +func TestApplyOptions_OnlyAnnotations(t *testing.T) { + g := NewWithT(t) + + tool := mcp.Tool{ + Name: "simple_tool", + Description: "A simple tool", + } + + // Apply only annotations + result := ApplyOptions(tool, WithToolAnnotations(ToolAnnotationConfig{ + Title: "Simple Tool", + ReadOnlyHint: boolPtr(true), + })) + + // Verify annotations are set + g.Expect(result.Annotations.Title).To(Equal("Simple Tool")) +} + +func TestApplyOptions_OnlyExtraProperties(t *testing.T) { + g := NewWithT(t) + + originalSchema := map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{}, + } + schemaBytes, err := json.Marshal(originalSchema) + g.Expect(err).ToNot(HaveOccurred()) + + tool := mcp.Tool{ + Name: "props_tool", + Description: "A tool with extra properties", + RawInputSchema: json.RawMessage(schemaBytes), + } + + // Apply only extra properties + result := ApplyOptions(tool, WithExtraProperties(ExtraProperty{ + Name: "api_key", + Description: "API key", + Required: false, + })) + + // Verify annotations are NOT set (empty struct) + g.Expect(result.Annotations.Title).To(BeEmpty()) + g.Expect(result.Annotations.ReadOnlyHint).To(BeNil()) + + // Verify extra properties are added + var resultSchema map[string]interface{} + err = json.Unmarshal(result.RawInputSchema, &resultSchema) + g.Expect(err).ToNot(HaveOccurred()) + + properties := resultSchema["properties"].(map[string]interface{}) + g.Expect(properties).To(HaveKey("api_key")) +} diff --git a/pkg/runtime/extra_properties.go b/pkg/runtime/extra_properties.go index 1ca95fb..c0c1da6 100644 --- a/pkg/runtime/extra_properties.go +++ b/pkg/runtime/extra_properties.go @@ -19,6 +19,7 @@ type ExtraProperty struct { type config struct { ExtraProperties []ExtraProperty + Annotations *ToolAnnotationConfig } // WithExtraProperties adds extra properties to tool schemas and extracts them from request arguments