Vectrum
Sign In
216 lines · 7.3 KB · Go
package mcp

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

	"github.com/johannSo/codey/internal/tools"
)

// namePrefix and nameSep spell a tool's registered name as
// mcp__<server>__<tool>. It is the convention Claude Code and Codex share, so
// a prompt or a hook written for either names the same tool here. The double
// underscore is what keeps the three parts apart when a server or tool name
// has a single one in it.
const (
	namePrefix = "mcp__"
	nameSep    = "__"
)

// maxToolNameLen is the API's ceiling on a tool name. A longer one is refused
// at registration rather than sent, because the API's refusal would come back
// against the whole request and take every other tool down with it.
const maxToolNameLen = 64

// emptySchema is the input schema for a tool whose server sent none. The API
// requires an object schema on every tool, and a server that takes no
// arguments may well leave the field out.
var emptySchema = json.RawMessage(`{"type":"object","properties":{}}`)

// mcpTool is one server tool wearing Codey's Tool interface.
type mcpTool struct {
	c      *client
	server string
	remote string
	name   string
	desc   string
	schema json.RawMessage
}

// newTool builds the adapter for info, or returns an error if the name it
// would register under cannot be sent to the API.
func newTool(c *client, server string, info toolInfo) (tools.Tool, error) {
	if info.Name == "" {
		return nil, fmt.Errorf("server listed a tool with no name")
	}
	name := toolName(server, info.Name)
	if len(name) > maxToolNameLen {
		return nil, fmt.Errorf("tool %q: registered name %q is longer than the API's %d-character limit", info.Name, name, maxToolNameLen)
	}
	desc := strings.TrimSpace(info.Description)
	if desc == "" {
		// The API rejects a tool with no description, and a server that sent
		// none has still told us what the tool is called and where it lives,
		// which is something for the model to go on.
		desc = fmt.Sprintf("The %s tool of the %s MCP server.", info.Name, server)
	}
	schema := info.InputSchema
	if len(schema) == 0 || string(schema) == "null" {
		schema = emptySchema
	}
	return &mcpTool{c: c, server: server, remote: info.Name, name: name, desc: desc, schema: schema}, nil
}

// toolName is the registered name for tool on server, with any character
// the API's name grammar does not allow replaced by an underscore. Server
// names come from the config and tool names from the server; neither was
// written with ^[a-zA-Z0-9_-]{1,64}$ in mind.
func toolName(server, tool string) string {
	return namePrefix + sanitize(server) + nameSep + sanitize(tool)
}

func sanitize(s string) string {
	var b strings.Builder
	b.Grow(len(s))
	for _, r := range s {
		switch {
		case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '_', r == '-':
			b.WriteRune(r)
		default:
			b.WriteByte('_')
		}
	}
	return b.String()
}

func (t *mcpTool) Name() string                 { return t.name }
func (t *mcpTool) Description() string          { return t.desc }
func (t *mcpTool) InputSchema() json.RawMessage { return t.schema }

// Execute forwards the call and renders what comes back. The arguments are
// passed through untouched: the schema they were validated against is the
// server's, and the server is the one to say whether they fit it.
func (t *mcpTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
	res, err := t.c.callTool(ctx, t.remote, raw)
	if err != nil {
		return "", fmt.Errorf("%s: %w", t.name, err)
	}
	var result callResult
	if err := json.Unmarshal(res, &result); err != nil {
		return "", fmt.Errorf("%s: server sent a malformed result: %w", t.name, err)
	}
	text := result.render()
	if result.IsError {
		// isError is the server saying the tool ran and failed, as distinct
		// from a protocol error. Either way the model needs to see a failure,
		// and the text is the server's account of it.
		return "", fmt.Errorf("%s: %s", t.name, text)
	}
	return text, nil
}

// callResult is what tools/call returns.
type callResult struct {
	Content           []contentItem   `json:"content"`
	StructuredContent json.RawMessage `json:"structuredContent"`
	IsError           bool            `json:"isError"`
}

// contentItem is one element of the content array. The fields are the
// union of the text, image, audio, resource and resource_link shapes; Type
// says which apply.
type contentItem struct {
	Type     string `json:"type"`
	Text     string `json:"text"`
	Data     string `json:"data"`
	MimeType string `json:"mimeType"`
	URI      string `json:"uri"`
	Name     string `json:"name"`
	Resource *struct {
		URI      string `json:"uri"`
		MimeType string `json:"mimeType"`
		Text     string `json:"text"`
		Blob     string `json:"blob"`
	} `json:"resource"`
}

// render flattens the content array to the string the model reads. Text is
// carried verbatim. Binary content is described rather than carried: a
// screenshot is megabytes of base64 that the model cannot read as text, and
// the transcript, the log and the context window would all pay for it. The
// registry has a mechanism for handing the model real image blocks; wiring
// it up here is deliberately left for later.
func (r callResult) render() string {
	var parts []string
	for _, item := range r.Content {
		switch item.Type {
		case "text":
			parts = append(parts, item.Text)
		case "image":
			parts = append(parts, fmt.Sprintf("[image: %s, about %s; images from MCP tools are described, not shown]",
				orUnknown(item.MimeType), approxSize(item.Data)))
		case "audio":
			parts = append(parts, fmt.Sprintf("[audio: %s, about %s; not played]", orUnknown(item.MimeType), approxSize(item.Data)))
		case "resource":
			if item.Resource == nil {
				parts = append(parts, "[resource with no content]")
				continue
			}
			res := item.Resource
			switch {
			case res.Text != "":
				parts = append(parts, fmt.Sprintf("[resource %s]\n%s", res.URI, res.Text))
			case res.Blob != "":
				parts = append(parts, fmt.Sprintf("[resource %s: %s, about %s of binary data; not shown]",
					res.URI, orUnknown(res.MimeType), approxSize(res.Blob)))
			default:
				parts = append(parts, fmt.Sprintf("[resource %s: empty]", res.URI))
			}
		case "resource_link":
			line := "[resource link: " + item.URI
			if item.Name != "" {
				line += " — " + item.Name
			}
			if item.Text != "" {
				line += ": " + item.Text
			}
			parts = append(parts, line+"]")
		default:
			parts = append(parts, fmt.Sprintf("[content of unsupported type %q]", item.Type))
		}
	}
	// A server that fills structuredContent may leave the content array
	// empty; the spec asks it to mirror the one in the other, but does not
	// make it. The JSON is the fallback rather than the default because a
	// server that sent both meant the text for reading.
	if len(parts) == 0 && len(r.StructuredContent) > 0 && string(r.StructuredContent) != "null" {
		parts = append(parts, string(r.StructuredContent))
	}
	if len(parts) == 0 {
		return "(no output)"
	}
	return strings.Join(parts, "\n")
}

func orUnknown(mime string) string {
	if mime == "" {
		return "unknown type"
	}
	return mime
}

// approxSize is the decoded size of base64 data, in round units. Three
// bytes per four characters is close enough for a caption.
func approxSize(b64 string) string {
	n := len(b64) / 4 * 3
	switch {
	case n >= 1<<20:
		return fmt.Sprintf("%.1f MB", float64(n)/(1<<20))
	case n >= 1<<10:
		return fmt.Sprintf("%d KB", n>>10)
	default:
		return fmt.Sprintf("%d bytes", n)
	}
}