19147346684

Committed 06 Nov 2025 07:25PM UTC coverage: 46.443% (-7.6%) from 54.037%

Build # 19147346684

Build Type

Pull #48

github

Committed by

fogfish

Commit Message

add file header

Pull Request Pull Request #48: Use MCP Servers as an integration point instead of commands

Run Details

116 of 148 new or added lines in 4 files covered. (78.38%)

5 existing lines in 1 file now uncovered.

235 of 506 relevant lines covered (46.44%)

0.5 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

72.22

/command/registry.go

//
// Copyright (C) 2025 Dmitry Kolesnikov
//
// This file may be modified and distributed under the terms
// of the MIT license.  See the LICENSE file for details.
// https://github.com/kshard/thinker
//

// Package command implements a registry for Model Context Protocol (MCP) servers.
// It provides integration with MCP tools, allowing agents to dynamically discover
// and invoke tools exposed by MCP servers.
package command

import (
        "context"
        "encoding/json"
        "fmt"
        "strings"

        "github.com/kshard/chatter"
        "github.com/kshard/thinker"
        "github.com/modelcontextprotocol/go-sdk/mcp"
)

// MCP Server interface
type Server interface {
        ListTools(ctx context.Context, params *mcp.ListToolsParams) (*mcp.ListToolsResult, error)
        CallTool(ctx context.Context, params *mcp.CallToolParams) (*mcp.CallToolResult, error)
        Close() error
}

const kSchemaSplit = "_"

type Registry struct {
        servers map[string]Server
        cmds    chatter.Registry
}

var _ thinker.Registry = (*Registry)(nil)

// NewRegistry creates a new registry of MCP servers.
func NewRegistry() *Registry {
        return &Registry{
                servers: make(map[string]Server),
                cmds:    chatter.Registry{},
        }
}

// Attach MCP server to the registry, making its tools available to the agent.
// The server is identified by a unique prefix, which is used to namespace
// tool names (e.g., fs_read). Tool names use underscore separator (prefix_toolname)
// due to AWS Bedrock constraints which only allow [a-zA-Z0-9_-] characters.
// The first token before underscore is always the server prefix.
// Each server runs independently, and its tools are registered with the prefix
// to avoid naming conflicts.
func (r *Registry) Attach(id string, server Server) error {
        if id == "" {
                return fmt.Errorf("server ID cannot be empty")
        }

        r.servers[id] = server
        r.cmds = chatter.Registry{}

        return nil
}

// Context returns the registry as LLM embeddable schema.
// It fetches the list of available tools from all attached MCP servers.
func (r *Registry) Context() chatter.Registry {
        // Return cached if available
        if len(r.cmds) > 0 {
                return r.cmds
        }

        ctx := context.Background()
        seq := make([]chatter.Cmd, 0)

        // Collect tools from all attached servers
        for id, srv := range r.servers {
                tools, err := srv.ListTools(ctx, &mcp.ListToolsParams{})
                if err != nil {
                        continue
                }

                for _, tool := range tools.Tools {
                        cmd := convertTool(*tool, id)
                        seq = append(seq, cmd)
                }
        }

        r.cmds = seq
        return r.cmds
}

// Invoke executes the tools requested by the LLM via the appropriate MCP server.
func (r *Registry) Invoke(reply *chatter.Reply) (thinker.Phase, chatter.Message, error) {
        ctx := context.Background()
        answer, err := reply.Invoke(func(name string, args json.RawMessage) (json.RawMessage, error) {
                seq := strings.SplitN(name, kSchemaSplit, 2)
                if len(seq) != 2 {
                        return nil, thinker.Feedback(
                                fmt.Sprintf("invalid tool name %s, missing the prefix", name),
                        )
                }
                id, tool := seq[0], seq[1]

                // Find which server handles this tool
                srv, exists := r.servers[id]
                if !exists {
                        return nil, thinker.Feedback(
                                fmt.Sprintf("tool %s is not available in any attached MCP server", name),
                        )
                }

                // Unmarshal arguments to pass to MCP
                var arguments map[string]any
                if len(args) > 0 {
                        if err := json.Unmarshal(args, &arguments); err != nil {
                                return nil, thinker.Feedback(
                                        fmt.Sprintf("failed to parse arguments for tool %s: %v", name, err),
                                )
                        }
                }

                // Call the tool via MCP using the actual tool name (without prefix)
                result, err := srv.CallTool(ctx, &mcp.CallToolParams{
                        Name:      tool,
                        Arguments: arguments,
                })
                if err != nil {
                        return nil, err
                }

                // Handle tool execution errors
                if result.IsError {
                        errorMsg := "tool execution failed"
                        if len(result.Content) > 0 {
                                if text, ok := result.Content[0].(*mcp.TextContent); ok {
                                        errorMsg = text.Text
                                }
                        }
                        return nil, thinker.Feedback(errorMsg)
                }

                // Extract and pack the result
                output := extractContent(result)
                return pack(output)
        })

        if err != nil {
                return thinker.AGENT_ABORT, nil, err
        }

        return thinker.AGENT_ASK, &answer, nil
}

// convertTool converts an MCP Tool to a chatter.Cmd format.
func convertTool(tool mcp.Tool, prefix string) chatter.Cmd {
        about := tool.Description
        if about == "" && tool.Title != "" {
                about = tool.Title
        }

        var schema json.RawMessage
        if tool.InputSchema != nil {
                if raw, ok := tool.InputSchema.(json.RawMessage); ok {
                        schema = raw
                } else {
                        // Convert to JSON if not already RawMessage
                        if b, err := json.Marshal(tool.InputSchema); err == nil {
                                schema = json.RawMessage(b)
                        }
                }
        }

        // Apply prefix if set (compact IRI notation: prefix:tool_name)
        name := tool.Name
        if prefix != "" {
                name = prefix + kSchemaSplit + tool.Name
        }

        return chatter.Cmd{
                Cmd:    name,
                About:  about,
                Schema: schema,
        }
}

// extractContent extracts the text content from a CallToolResult.
func extractContent(result *mcp.CallToolResult) []byte {
        if len(result.Content) == 0 {
                return []byte{}
        }

        // Try to extract text content
        for _, content := range result.Content {
                switch c := content.(type) {
                case *mcp.TextContent:
                        return []byte(c.Text)
                case *mcp.ImageContent:
                        // For image content, return a description or empty
                        return []byte("[Image content]")
                }
        }

        // Fallback: try to marshal the content as JSON
        if b, err := json.Marshal(result.Content); err == nil {
                return b
        }

        return []byte{}
}

// pack wraps the tool output in the expected format.
func pack(b []byte) (json.RawMessage, error) {
        pckt := map[string]any{
                "toolOutput": string(b),
        }

        bin, err := json.Marshal(pckt)
        if err != nil {
                return nil, err
        }

        return json.RawMessage(bin), nil
}

1	//
2	// Copyright (C) 2025 Dmitry Kolesnikov
3	//
4	// This file may be modified and distributed under the terms
5	// of the MIT license. See the LICENSE file for details.
6	// https://github.com/kshard/thinker
7	//
8
9	// Package command implements a registry for Model Context Protocol (MCP) servers.
10	// It provides integration with MCP tools, allowing agents to dynamically discover
11	// and invoke tools exposed by MCP servers.
12	package command
13
14	import (
15	"context"
16	"encoding/json"
17	"fmt"
18	"strings"
19
20	"github.com/kshard/chatter"
21	"github.com/kshard/thinker"
22	"github.com/modelcontextprotocol/go-sdk/mcp"
23	)
24
25	// MCP Server interface
26	type Server interface {
27	ListTools(ctx context.Context, params mcp.ListToolsParams) (mcp.ListToolsResult, error)
28	CallTool(ctx context.Context, params mcp.CallToolParams) (mcp.CallToolResult, error)
29	Close() error
30	}
31
32	const kSchemaSplit = "_"
33
34	type Registry struct {
35	servers map[string]Server
36	cmds chatter.Registry
37	}
38
39	var _ thinker.Registry = (*Registry)(nil)
40
41	// NewRegistry creates a new registry of MCP servers.
42	func NewRegistry() *Registry {	1✔
43	return &Registry{	1✔
44	servers: make(map[string]Server),	1✔
45	cmds: chatter.Registry{},	1✔
46	}	1✔
47	}	1✔
48
49	// Attach MCP server to the registry, making its tools available to the agent.
50	// The server is identified by a unique prefix, which is used to namespace
51	// tool names (e.g., fs_read). Tool names use underscore separator (prefix_toolname)
52	// due to AWS Bedrock constraints which only allow [a-zA-Z0-9_-] characters.
53	// The first token before underscore is always the server prefix.
54	// Each server runs independently, and its tools are registered with the prefix
55	// to avoid naming conflicts.
56	func (r *Registry) Attach(id string, server Server) error {	1✔
57	if id == "" {	2✔
58	return fmt.Errorf("server ID cannot be empty")	1✔
59	}	1✔
60
61	r.servers[id] = server	1✔
62	r.cmds = chatter.Registry{}	1✔
63		1✔
64	return nil	1✔
65	}
66
67	// Context returns the registry as LLM embeddable schema.
68	// It fetches the list of available tools from all attached MCP servers.
69	func (r *Registry) Context() chatter.Registry {	1✔
70	// Return cached if available	1✔
71	if len(r.cmds) > 0 {	1✔
NEW 72	return r.cmds	×
UNCOV 73	}	×
74
75	ctx := context.Background()	1✔
76	seq := make([]chatter.Cmd, 0)	1✔
77		1✔
78	// Collect tools from all attached servers	1✔
79	for id, srv := range r.servers {	2✔
80	tools, err := srv.ListTools(ctx, &mcp.ListToolsParams{})	1✔
81	if err != nil {	1✔
NEW 82	continue	×
83	}
84
85	for _, tool := range tools.Tools {	2✔
86	cmd := convertTool(*tool, id)	1✔
87	seq = append(seq, cmd)	1✔
88	}	1✔
89	}
90
91	r.cmds = seq	1✔
92	return r.cmds	1✔
93	}
94
95	// Invoke executes the tools requested by the LLM via the appropriate MCP server.
96	func (r Registry) Invoke(reply chatter.Reply) (thinker.Phase, chatter.Message, error) {	1✔
97	ctx := context.Background()	1✔
98	answer, err := reply.Invoke(func(name string, args json.RawMessage) (json.RawMessage, error) {	2✔
99	seq := strings.SplitN(name, kSchemaSplit, 2)	1✔
100	if len(seq) != 2 {	2✔
101	return nil, thinker.Feedback(	1✔
102	fmt.Sprintf("invalid tool name %s, missing the prefix", name),	1✔
103	)	1✔
104	}	1✔
105	id, tool := seq[0], seq[1]	1✔
106		1✔
107	// Find which server handles this tool	1✔
108	srv, exists := r.servers[id]	1✔
109	if !exists {	1✔
NEW 110	return nil, thinker.Feedback(	×
NEW 111	fmt.Sprintf("tool %s is not available in any attached MCP server", name),	×
NEW 112	)	×
NEW 113	}	×
114
115	// Unmarshal arguments to pass to MCP
116	var arguments map[string]any	1✔
117	if len(args) > 0 {	2✔
118	if err := json.Unmarshal(args, &arguments); err != nil {	2✔
119	return nil, thinker.Feedback(	1✔
120	fmt.Sprintf("failed to parse arguments for tool %s: %v", name, err),	1✔
121	)	1✔
122	}	1✔
123	}
124
125	// Call the tool via MCP using the actual tool name (without prefix)
126	result, err := srv.CallTool(ctx, &mcp.CallToolParams{	1✔
127	Name: tool,	1✔
128	Arguments: arguments,	1✔
129	})	1✔
130	if err != nil {	1✔
NEW 131	return nil, err	×
UNCOV 132	}	×
133
134	// Handle tool execution errors
135	if result.IsError {	1✔
NEW 136	errorMsg := "tool execution failed"	×
NEW 137	if len(result.Content) > 0 {	×
NEW 138	if text, ok := result.Content[0].(*mcp.TextContent); ok {	×
NEW 139	errorMsg = text.Text	×
NEW 140	}	×
141	}
NEW 142	return nil, thinker.Feedback(errorMsg)	×
143	}
144
145	// Extract and pack the result
146	output := extractContent(result)	1✔
147	return pack(output)	1✔
148	})
149
150	if err != nil {	2✔
151	return thinker.AGENT_ABORT, nil, err	1✔
152	}	1✔
153
154	return thinker.AGENT_ASK, &answer, nil	1✔
155	}
156
157	// convertTool converts an MCP Tool to a chatter.Cmd format.
158	func convertTool(tool mcp.Tool, prefix string) chatter.Cmd {	1✔
159	about := tool.Description	1✔
160	if about == "" && tool.Title != "" {	1✔
NEW 161	about = tool.Title	×
NEW 162	}	×
163
164	var schema json.RawMessage	1✔
165	if tool.InputSchema != nil {	1✔
NEW 166	if raw, ok := tool.InputSchema.(json.RawMessage); ok {	×
NEW 167	schema = raw	×
NEW 168	} else {	×
NEW 169	// Convert to JSON if not already RawMessage	×
NEW 170	if b, err := json.Marshal(tool.InputSchema); err == nil {	×
NEW 171	schema = json.RawMessage(b)	×
UNCOV 172	}	×
173	}
174	}
175
176	// Apply prefix if set (compact IRI notation: prefix:tool_name)
177	name := tool.Name	1✔
178	if prefix != "" {	2✔
179	name = prefix + kSchemaSplit + tool.Name	1✔
180	}	1✔
181
182	return chatter.Cmd{	1✔
183	Cmd: name,	1✔
184	About: about,	1✔
185	Schema: schema,	1✔
186	}	1✔
187	}
188
189	// extractContent extracts the text content from a CallToolResult.
190	func extractContent(result *mcp.CallToolResult) []byte {	1✔
191	if len(result.Content) == 0 {	1✔
NEW 192	return []byte{}	×
UNCOV 193	}	×
194
195	// Try to extract text content
196	for _, content := range result.Content {	2✔
197	switch c := content.(type) {	1✔
198	case *mcp.TextContent:	1✔
199	return []byte(c.Text)	1✔
NEW 200	case *mcp.ImageContent:	×
NEW 201	// For image content, return a description or empty	×
NEW 202	return []byte("[Image content]")	×
203	}
204	}
205
206	// Fallback: try to marshal the content as JSON
NEW 207	if b, err := json.Marshal(result.Content); err == nil {	×
NEW 208	return b	×
UNCOV 209	}	×
210
NEW 211	return []byte{}	×
212	}
213
214	// pack wraps the tool output in the expected format.
215	func pack(b []byte) (json.RawMessage, error) {	1✔
216	pckt := map[string]any{	1✔
217	"toolOutput": string(b),	1✔
218	}	1✔
219		1✔
220	bin, err := json.Marshal(pckt)	1✔
221	if err != nil {	1✔
222	return nil, err	×
223	}	×
224
225	return json.RawMessage(bin), nil	1✔
226	}

kshard / thinker / 19147346684

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous