Skip to content

Commit 238d844

Browse files
muxammadrezamuxammadreza
committed
Preserve output schemas in direct mode
1 parent b769a55 commit 238d844

12 files changed

Lines changed: 296 additions & 63 deletions

File tree

‎internal/config/config.go‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -239,9 +239,9 @@ type ServerConfig struct {
239239
// when the server is configured with both Command and an HTTP/SSE URL — i.e.,
240240
// mcpproxy starts the process AND connects via network. Stdio servers ignore
241241
// this field. Zero or unset → 30s default.
242-
LauncherWaitTimeout Duration `json:"launcher_wait_timeout,omitempty" mapstructure:"launcher_wait_timeout" swaggertype:"string"`
243-
EnabledTools []string `json:"enabled_tools,omitempty" mapstructure:"enabled_tools"` // Allowlist: only these tools are exposed; mutually exclusive with disabled_tools
244-
DisabledTools []string `json:"disabled_tools,omitempty" mapstructure:"disabled_tools"` // Denylist: these tools are hidden; mutually exclusive with enabled_tools
242+
LauncherWaitTimeout Duration `json:"launcher_wait_timeout,omitempty" mapstructure:"launcher_wait_timeout" swaggertype:"string"`
243+
EnabledTools []string `json:"enabled_tools,omitempty" mapstructure:"enabled_tools"` // Allowlist: only these tools are exposed; mutually exclusive with disabled_tools
244+
DisabledTools []string `json:"disabled_tools,omitempty" mapstructure:"disabled_tools"` // Denylist: these tools are hidden; mutually exclusive with enabled_tools
245245
}
246246

247247
// OAuthConfig represents OAuth configuration for a server
@@ -523,14 +523,15 @@ func ConvertFromCursorFormat(cursorConfig *CursorMCPConfig) []*ServerConfig {
523523

524524
// ToolMetadata represents tool information stored in the index
525525
type ToolMetadata struct {
526-
Name string `json:"name"`
527-
ServerName string `json:"server_name"`
528-
Description string `json:"description"`
529-
ParamsJSON string `json:"params_json"`
530-
Hash string `json:"hash"`
531-
Created time.Time `json:"created"`
532-
Updated time.Time `json:"updated"`
533-
Annotations *ToolAnnotations `json:"annotations,omitempty"`
526+
Name string `json:"name"`
527+
ServerName string `json:"server_name"`
528+
Description string `json:"description"`
529+
ParamsJSON string `json:"params_json"`
530+
OutputSchemaJSON string `json:"output_schema_json,omitempty"`
531+
Hash string `json:"hash"`
532+
Created time.Time `json:"created"`
533+
Updated time.Time `json:"updated"`
534+
Annotations *ToolAnnotations `json:"annotations,omitempty"`
534535
}
535536

536537
// ToolAnnotations represents MCP tool behavior hints

‎internal/config/config_test.go‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -334,13 +334,14 @@ func TestConfigSecurityModes(t *testing.T) {
334334
func TestToolMetadata(t *testing.T) {
335335
now := time.Now()
336336
tool := &ToolMetadata{
337-
Name: "test:tool",
338-
ServerName: "test",
339-
Description: "A test tool",
340-
ParamsJSON: `{"type": "object", "properties": {"param1": {"type": "string"}}}`,
341-
Hash: "abc123",
342-
Created: now,
343-
Updated: now,
337+
Name: "test:tool",
338+
ServerName: "test",
339+
Description: "A test tool",
340+
ParamsJSON: `{"type": "object", "properties": {"param1": {"type": "string"}}}`,
341+
OutputSchemaJSON: `{"type":"object","properties":{"result":{"type":"string"}}}`,
342+
Hash: "abc123",
343+
Created: now,
344+
Updated: now,
344345
}
345346

346347
// Test JSON serialization
@@ -355,6 +356,7 @@ func TestToolMetadata(t *testing.T) {
355356
assert.Equal(t, tool.ServerName, restored.ServerName)
356357
assert.Equal(t, tool.Description, restored.Description)
357358
assert.Equal(t, tool.ParamsJSON, restored.ParamsJSON)
359+
assert.Equal(t, tool.OutputSchemaJSON, restored.OutputSchemaJSON)
358360
assert.Equal(t, tool.Hash, restored.Hash)
359361
assert.True(t, tool.Created.Equal(restored.Created))
360362
assert.True(t, tool.Updated.Equal(restored.Updated))

‎internal/hash/hash.go‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,16 @@ import (
77
"fmt"
88
)
99

10-
// ToolHash computes SHA-256 hash for tool change detection
10+
// ToolHash computes SHA-256 hash for tool change detection.
1111
// Format: sha256(serverName + toolName + description + parametersSchemaJSON)
1212
func ToolHash(serverName, toolName, description string, parametersSchema interface{}) (string, error) {
13+
return ToolHashWithOutputSchema(serverName, toolName, description, parametersSchema, "")
14+
}
15+
16+
// ToolHashWithOutputSchema computes SHA-256 hash for the full tool contract,
17+
// including output schema when present.
18+
// Format: sha256(serverName + toolName + description + parametersSchemaJSON + outputSchemaJSON)
19+
func ToolHashWithOutputSchema(serverName, toolName, description string, parametersSchema interface{}, outputSchemaJSON string) (string, error) {
1320
// Serialize parameters schema to JSON for consistent hashing
1421
var schemaBytes []byte
1522
var err error
@@ -21,8 +28,8 @@ func ToolHash(serverName, toolName, description string, parametersSchema interfa
2128
}
2229
}
2330

24-
// Combine server name, tool name, description, and schema JSON
25-
combined := serverName + toolName + description + string(schemaBytes)
31+
// Combine server name, tool name, description, input schema JSON, and output schema JSON
32+
combined := serverName + toolName + description + string(schemaBytes) + outputSchemaJSON
2633

2734
// Compute SHA-256 hash
2835
hasher := sha256.New()
@@ -60,7 +67,12 @@ func VerifyToolHash(serverName, toolName, description string, parametersSchema i
6067

6168
// ComputeToolHash computes a SHA256 hash for a tool (alias for ToolHash that doesn't return error)
6269
func ComputeToolHash(serverName, toolName, description string, inputSchema interface{}) string {
63-
hash, err := ToolHash(serverName, toolName, description, inputSchema)
70+
return ComputeToolHashWithOutputSchema(serverName, toolName, description, inputSchema, "")
71+
}
72+
73+
// ComputeToolHashWithOutputSchema computes a SHA256 hash for a tool including output schema.
74+
func ComputeToolHashWithOutputSchema(serverName, toolName, description string, inputSchema interface{}, outputSchemaJSON string) string {
75+
hash, err := ToolHashWithOutputSchema(serverName, toolName, description, inputSchema, outputSchemaJSON)
6476
if err != nil {
6577
// If hashing fails, return a default hash based on server and tool name
6678
fallback := StringHash(fmt.Sprintf("%s:%s", serverName, toolName))

‎internal/hash/hash_test.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,17 @@ func TestComputeToolHash_DescriptionOnlyChange(t *testing.T) {
150150
assert.NotEqual(t, hash1, hash2, "Description-only changes must produce different hashes")
151151
}
152152

153+
func TestComputeToolHashWithOutputSchema_OutputSchemaChange(t *testing.T) {
154+
inputSchema := map[string]interface{}{"type": "object"}
155+
outputSchema1 := `{"type":"object","properties":{"result":{"type":"string"}}}`
156+
outputSchema2 := `{"type":"object","properties":{"result":{"type":"number"}}}`
157+
158+
hash1 := ComputeToolHashWithOutputSchema("myserver", "my_tool", "desc", inputSchema, outputSchema1)
159+
hash2 := ComputeToolHashWithOutputSchema("myserver", "my_tool", "desc", inputSchema, outputSchema2)
160+
161+
assert.NotEqual(t, hash1, hash2, "Output schema changes must produce different hashes")
162+
}
163+
153164
func TestVerifyToolHash_Match(t *testing.T) {
154165
schema := map[string]interface{}{"type": "object"}
155166
desc := "Test tool"

‎internal/index/bleve.go‎

Lines changed: 41 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,15 @@ type BleveIndex struct {
2121

2222
// ToolDocument represents a tool document in the index
2323
type ToolDocument struct {
24-
ToolName string `json:"tool_name"` // Just the tool name (without server prefix)
25-
FullToolName string `json:"full_tool_name"` // Complete server:tool format
26-
ServerName string `json:"server_name"`
27-
Description string `json:"description"`
28-
ParamsJSON string `json:"params_json"`
29-
Hash string `json:"hash"`
30-
Tags string `json:"tags"`
31-
SearchableText string `json:"searchable_text"` // Combined searchable content
24+
ToolName string `json:"tool_name"` // Just the tool name (without server prefix)
25+
FullToolName string `json:"full_tool_name"` // Complete server:tool format
26+
ServerName string `json:"server_name"`
27+
Description string `json:"description"`
28+
ParamsJSON string `json:"params_json"`
29+
OutputSchemaJSON string `json:"output_schema_json,omitempty"`
30+
Hash string `json:"hash"`
31+
Tags string `json:"tags"`
32+
SearchableText string `json:"searchable_text"` // Combined searchable content
3233
}
3334

3435
// NewBleveIndex creates a new Bleve index
@@ -147,14 +148,15 @@ func (b *BleveIndex) IndexTool(toolMeta *config.ToolMetadata) error {
147148
toolMeta.ParamsJSON)
148149

149150
doc := &ToolDocument{
150-
ToolName: toolName,
151-
FullToolName: toolMeta.Name,
152-
ServerName: toolMeta.ServerName,
153-
Description: toolMeta.Description,
154-
ParamsJSON: toolMeta.ParamsJSON,
155-
Hash: toolMeta.Hash,
156-
Tags: "", // Can be extended later
157-
SearchableText: searchableText,
151+
ToolName: toolName,
152+
FullToolName: toolMeta.Name,
153+
ServerName: toolMeta.ServerName,
154+
Description: toolMeta.Description,
155+
ParamsJSON: toolMeta.ParamsJSON,
156+
OutputSchemaJSON: toolMeta.OutputSchemaJSON,
157+
Hash: toolMeta.Hash,
158+
Tags: "", // Can be extended later
159+
SearchableText: searchableText,
158160
}
159161

160162
// Use server:tool format as document ID for uniqueness
@@ -249,7 +251,7 @@ func (b *BleveIndex) SearchTools(queryStr string, limit int) ([]*config.SearchRe
249251
// Create search request
250252
searchReq := bleve.NewSearchRequest(boolQuery)
251253
searchReq.Size = limit
252-
searchReq.Fields = []string{"tool_name", "full_tool_name", "server_name", "description", "params_json", "hash"}
254+
searchReq.Fields = []string{"tool_name", "full_tool_name", "server_name", "description", "params_json", "output_schema_json", "hash"}
253255
searchReq.Highlight = bleve.NewHighlight()
254256

255257
b.logger.Debug("Searching tools with enhanced query", zap.String("query", queryStr), zap.Int("limit", limit))
@@ -263,11 +265,12 @@ func (b *BleveIndex) SearchTools(queryStr string, limit int) ([]*config.SearchRe
263265
var results []*config.SearchResult
264266
for _, hit := range searchResult.Hits {
265267
toolMeta := &config.ToolMetadata{
266-
Name: getStringField(hit.Fields, "full_tool_name"),
267-
ServerName: getStringField(hit.Fields, "server_name"),
268-
Description: getStringField(hit.Fields, "description"),
269-
ParamsJSON: getStringField(hit.Fields, "params_json"),
270-
Hash: getStringField(hit.Fields, "hash"),
268+
Name: getStringField(hit.Fields, "full_tool_name"),
269+
ServerName: getStringField(hit.Fields, "server_name"),
270+
Description: getStringField(hit.Fields, "description"),
271+
ParamsJSON: getStringField(hit.Fields, "params_json"),
272+
OutputSchemaJSON: getStringField(hit.Fields, "output_schema_json"),
273+
Hash: getStringField(hit.Fields, "hash"),
271274
}
272275

273276
results = append(results, &config.SearchResult{
@@ -306,14 +309,15 @@ func (b *BleveIndex) BatchIndex(tools []*config.ToolMetadata) error {
306309
toolMeta.ParamsJSON)
307310

308311
doc := &ToolDocument{
309-
ToolName: toolName,
310-
FullToolName: toolMeta.Name,
311-
ServerName: toolMeta.ServerName,
312-
Description: toolMeta.Description,
313-
ParamsJSON: toolMeta.ParamsJSON,
314-
Hash: toolMeta.Hash,
315-
Tags: "",
316-
SearchableText: searchableText,
312+
ToolName: toolName,
313+
FullToolName: toolMeta.Name,
314+
ServerName: toolMeta.ServerName,
315+
Description: toolMeta.Description,
316+
ParamsJSON: toolMeta.ParamsJSON,
317+
OutputSchemaJSON: toolMeta.OutputSchemaJSON,
318+
Hash: toolMeta.Hash,
319+
Tags: "",
320+
SearchableText: searchableText,
317321
}
318322

319323
docID := fmt.Sprintf("%s:%s", toolMeta.ServerName, toolName)
@@ -348,7 +352,7 @@ func (b *BleveIndex) GetToolsByServer(serverName string) ([]*config.ToolMetadata
348352
// Create search request with high limit to get all tools
349353
searchReq := bleve.NewSearchRequest(query)
350354
searchReq.Size = 10000 // Maximum tools per server
351-
searchReq.Fields = []string{"tool_name", "full_tool_name", "server_name", "description", "params_json", "hash"}
355+
searchReq.Fields = []string{"tool_name", "full_tool_name", "server_name", "description", "params_json", "output_schema_json", "hash"}
352356

353357
b.logger.Debug("Querying tools by server", zap.String("server", serverName))
354358

@@ -361,11 +365,12 @@ func (b *BleveIndex) GetToolsByServer(serverName string) ([]*config.ToolMetadata
361365
var tools []*config.ToolMetadata
362366
for _, hit := range searchResult.Hits {
363367
toolMeta := &config.ToolMetadata{
364-
Name: getStringField(hit.Fields, "full_tool_name"),
365-
ServerName: getStringField(hit.Fields, "server_name"),
366-
Description: getStringField(hit.Fields, "description"),
367-
ParamsJSON: getStringField(hit.Fields, "params_json"),
368-
Hash: getStringField(hit.Fields, "hash"),
368+
Name: getStringField(hit.Fields, "full_tool_name"),
369+
ServerName: getStringField(hit.Fields, "server_name"),
370+
Description: getStringField(hit.Fields, "description"),
371+
ParamsJSON: getStringField(hit.Fields, "params_json"),
372+
OutputSchemaJSON: getStringField(hit.Fields, "output_schema_json"),
373+
Hash: getStringField(hit.Fields, "hash"),
369374
}
370375
tools = append(tools, toolMeta)
371376
}

‎internal/index/bleve_test.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,60 @@ import (
1212
"github.com/smart-mcp-proxy/mcpproxy-go/internal/config"
1313
)
1414

15+
func TestBleveIndex_PreservesOutputSchemaJSON(t *testing.T) {
16+
tmpDir, err := os.MkdirTemp("", "bleve_output_schema_test_*")
17+
require.NoError(t, err)
18+
defer os.RemoveAll(tmpDir)
19+
20+
idx, err := NewBleveIndex(tmpDir, zap.NewNop())
21+
require.NoError(t, err)
22+
defer idx.Close()
23+
24+
tool := &config.ToolMetadata{
25+
Name: "github:get_issue",
26+
ServerName: "github",
27+
Description: "Get issue",
28+
ParamsJSON: `{"type":"object","properties":{"id":{"type":"string"}}}`,
29+
OutputSchemaJSON: `{"type":"object","properties":{"title":{"type":"string"}},"required":["title"]}`,
30+
Hash: "hash1",
31+
}
32+
33+
require.NoError(t, idx.IndexTool(tool))
34+
35+
results, err := idx.SearchTools("get issue", 10)
36+
require.NoError(t, err)
37+
require.NotEmpty(t, results)
38+
39+
assert.Equal(t, tool.OutputSchemaJSON, results[0].Tool.OutputSchemaJSON)
40+
}
41+
42+
func TestBleveIndex_BatchIndexPreservesOutputSchemaJSON(t *testing.T) {
43+
tmpDir, err := os.MkdirTemp("", "bleve_output_schema_batch_test_*")
44+
require.NoError(t, err)
45+
defer os.RemoveAll(tmpDir)
46+
47+
idx, err := NewBleveIndex(tmpDir, zap.NewNop())
48+
require.NoError(t, err)
49+
defer idx.Close()
50+
51+
tool := &config.ToolMetadata{
52+
Name: "github:list_issues",
53+
ServerName: "github",
54+
Description: "List issues",
55+
ParamsJSON: `{"type":"object"}`,
56+
OutputSchemaJSON: `{"type":"object","properties":{"issues":{"type":"array"}}}`,
57+
Hash: "hash2",
58+
}
59+
60+
require.NoError(t, idx.BatchIndex([]*config.ToolMetadata{tool}))
61+
62+
results, err := idx.SearchTools("list issues", 10)
63+
require.NoError(t, err)
64+
require.NotEmpty(t, results)
65+
66+
assert.Equal(t, tool.OutputSchemaJSON, results[0].Tool.OutputSchemaJSON)
67+
}
68+
1569
func TestBleveIndex_IndexAndSearch_DeFiLlamaTools(t *testing.T) {
1670
// Create test DeFiLlama tools based on the user's data
1771
defiLlamaTools := createTestDeFiLlamaTools()

‎internal/server/mcp_output_schema.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package server
2+
3+
import (
4+
"encoding/json"
5+
6+
"github.com/mark3labs/mcp-go/mcp"
7+
)
8+
9+
func applyOutputSchemaJSON(tool *mcp.Tool, schemaJSON string) {
10+
if tool == nil || schemaJSON == "" || !json.Valid([]byte(schemaJSON)) {
11+
return
12+
}
13+
14+
tool.RawOutputSchema = json.RawMessage(schemaJSON)
15+
tool.OutputSchema = mcp.ToolOutputSchema{}
16+
}

‎internal/server/mcp_output_schema_test.go‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package server
2+
3+
import (
4+
"encoding/json"
5+
"testing"
6+
7+
"github.com/mark3labs/mcp-go/mcp"
8+
"github.com/stretchr/testify/assert"
9+
"github.com/stretchr/testify/require"
10+
)
11+
12+
func TestApplyOutputSchemaJSON_PreservesRawOutputSchema(t *testing.T) {
13+
var tool mcp.Tool
14+
schemaJSON := `{"type":"object","properties":{"result":{"type":"string"}},"required":["result"],"additionalProperties":false}`
15+
16+
applyOutputSchemaJSON(&tool, schemaJSON)
17+
18+
assert.Empty(t, tool.OutputSchema.Type)
19+
require.NotNil(t, tool.RawOutputSchema)
20+
assert.JSONEq(t, schemaJSON, string(tool.RawOutputSchema))
21+
}
22+
23+
func TestApplyOutputSchemaJSON_PreservesValidSchemaWithoutObjectType(t *testing.T) {
24+
var tool mcp.Tool
25+
schemaJSON := `{"anyOf":[{"type":"string"},{"type":"number"}]}`
26+
27+
applyOutputSchemaJSON(&tool, schemaJSON)
28+
29+
assert.Empty(t, tool.OutputSchema.Type)
30+
require.NotNil(t, tool.RawOutputSchema)
31+
assert.JSONEq(t, schemaJSON, string(tool.RawOutputSchema))
32+
}
33+
34+
func TestApplyOutputSchemaJSON_IgnoresInvalidSchema(t *testing.T) {
35+
tool := mcp.Tool{
36+
OutputSchema: mcp.ToolOutputSchema{Type: "object"},
37+
RawOutputSchema: json.RawMessage(`{"type":"object"}`),
38+
}
39+
40+
applyOutputSchemaJSON(&tool, `{invalid`)
41+
42+
assert.Equal(t, "object", tool.OutputSchema.Type)
43+
assert.JSONEq(t, `{"type":"object"}`, string(tool.RawOutputSchema))
44+
}
45+
46+
func TestApplyOutputSchemaJSON_NilToolNoop(t *testing.T) {
47+
assert.NotPanics(t, func() {
48+
applyOutputSchemaJSON(nil, `{"type":"object"}`)
49+
})
50+
}

‎internal/server/mcp_routing.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,8 @@ func (p *MCPProxyServer) buildDirectModeTools() []mcpserver.ServerTool {
8585

8686
mcpTool := mcp.NewTool(directName, opts...)
8787

88+
applyOutputSchemaJSON(&mcpTool, tool.OutputSchemaJSON)
89+
8890
// Apply input schema from upstream tool
8991
if tool.ParamsJSON != "" {
9092
var schema map[string]interface{}

0 commit comments

Comments
 (0)