cogito/mcp.go

package cogito

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

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/mudler/xlog"
	"github.com/sashabaranov/go-openai"
	"github.com/tmc/langchaingo/jsonschema"
)

type mcpTool struct {
	name, description string
	inputSchema       toolInputSchema
	session           *mcp.ClientSession
	ctx               context.Context
	props             map[string]jsonschema.Definition
}

func (t *mcpTool) Tool() openai.Tool {
	return openai.Tool{
		Type: openai.ToolTypeFunction,
		Function: &openai.FunctionDefinition{
			Name:        t.name,
			Description: t.description,
			Parameters: jsonschema.Definition{
				Type:       jsonschema.Object,
				Properties: t.props,
				Required:   t.inputSchema.Required,
			},
		},
	}
}

func (t *mcpTool) Execute(args map[string]any) (string, any, error) {

	// Call a tool on the server.
	params := &mcp.CallToolParams{
		Name:      t.name,
		Arguments: args,
	}
	res, err := t.session.CallTool(t.ctx, params)
	if err != nil {
		xlog.Error("CallTool failed: %v", err)
		return "", nil, err
	}

	result := contentToString(res.Content)

	if res.IsError {
		xlog.Error("tool failed", "result", result)
		return result, nil, errors.New("tool failed:  " + result)
	}

	return result, res, nil
}

// contentToString flattens the content blocks of an MCP tool result into a
// single textual representation that can be fed back to the model. Non-text
// blocks (images, audio, resources) are summarized with a descriptive marker
// instead of being asserted to *mcp.TextContent, which would panic and crash
// the host process when a tool returns media (see mudler/LocalAI#10101).
func contentToString(content []mcp.Content) string {
	result := ""
	for _, c := range content {
		switch v := c.(type) {
		case *mcp.TextContent:
			result += v.Text
		case *mcp.ImageContent:
			result += fmt.Sprintf("[image content (%s), %d bytes]", v.MIMEType, len(v.Data))
		case *mcp.AudioContent:
			result += fmt.Sprintf("[audio content (%s), %d bytes]", v.MIMEType, len(v.Data))
		case *mcp.ResourceLink:
			result += fmt.Sprintf("[resource link: %s]", v.URI)
		case *mcp.EmbeddedResource:
			switch {
			case v.Resource == nil:
				result += "[embedded resource]"
			case v.Resource.Text != "":
				result += v.Resource.Text
			default:
				result += fmt.Sprintf("[embedded resource: %s]", v.Resource.URI)
			}
		default:
			xlog.Warn("Unhandled MCP content type", "type", fmt.Sprintf("%T", c))
		}
	}
	return result
}

func (t *mcpTool) Close() {
	if err := t.session.Close(); err != nil {
		xlog.Warn("Failed to close MCP session", "error", err)
	}
}

type toolInputSchema struct {
	Type       string                 `json:"type"`
	Properties map[string]interface{} `json:"properties,omitempty"`
	Required   []string               `json:"required,omitempty"`
}

// CoerceNullableTypes is an exported alias for the same workaround so
// downstream tests can verify their own MCP servers stay compatible
// with the import path. Most callers won't need it.
func CoerceNullableTypes(props map[string]any) { coerceNullableTypes(props) }

// coerceNullableTypes recursively walks a property bag and normalizes the
// JSON-Schema 2020-12 constructs that langchaingo/jsonschema.Definition
// cannot represent. Both rewrites exist for the same reason: Definition is
// the struct we unmarshal every listed tool's properties into, and ANY
// failure there drops the whole tool from the model's tool list with a
// single log line — the server stays healthy and its tool simply never
// reaches the model.
//
//  1. "type": ["null", "X"] becomes "type": "X". Definition.Type is a single
//     string. Picks the first non-null member; falls back to the first
//     member if all are null.
//  2. A boolean schema becomes its object equivalent: `true` (allow
//     anything) becomes {}, `false` (allow nothing) becomes {"not": {}}.
//     2020-12 permits a boolean wherever a schema is allowed, and
//     google/jsonschema-go emits exactly that for a Go `any` — an empty
//     schema marshals as `true`, so a [][]any field yields "items": true.
//     Definition models nested schemas as *Definition, so a bare bool fails
//     the unmarshal.
//
// Recurses into every nested schema location: properties, items,
// oneOf/anyOf/allOf members, prefixItems, $defs/definitions,
// additionalProperties, patternProperties, contains, not, if/then/else,
// propertyNames.
func coerceNullableTypes(props map[string]any) {
	if props == nil {
		return
	}
	// Range with the key so a property that is ITSELF a boolean schema can be
	// replaced in place; the value-only loop could not rewrite it.
	for name, raw := range props {
		if b, ok := raw.(bool); ok {
			props[name] = boolSchema(b)
			continue
		}
		coerceSchema(raw)
	}
}

// boolSchema returns the object form of a JSON-Schema boolean schema,
// preserving its meaning: `true` allows anything, `false` allows nothing.
// Collapsing both to {} would silently widen a deliberately-closed schema.
func boolSchema(allow bool) map[string]any {
	if allow {
		return map[string]any{}
	}
	return map[string]any{"not": map[string]any{}}
}

// schemaValuedKeys are the keywords whose value is a single schema, so a
// boolean there is a boolean schema rather than an ordinary flag.
var schemaValuedKeys = []string{
	"items", "additionalProperties", "contains", "not",
	"if", "then", "else", "propertyNames",
}

// schemaListKeys are the keywords whose value is an array of schemas.
var schemaListKeys = []string{"oneOf", "anyOf", "allOf", "prefixItems", "items"}

// normalizeBoolSchemas replaces every boolean schema directly under obj with
// its object equivalent. Nested bags (properties, $defs, …) are reached by
// coerceNullableTypes, which performs the same replacement for their members.
func normalizeBoolSchemas(obj map[string]any) {
	for _, key := range schemaValuedKeys {
		if b, ok := obj[key].(bool); ok {
			obj[key] = boolSchema(b)
		}
	}
	for _, key := range schemaListKeys {
		arr, ok := obj[key].([]any)
		if !ok {
			continue
		}
		for i, member := range arr {
			if b, ok := member.(bool); ok {
				arr[i] = boolSchema(b)
			}
		}
	}
}

// coerceSchema applies the type-array → string and boolean-schema →
// object rewrites to a single schema node, then recurses into every
// nested schema location.
func coerceSchema(node any) {
	obj, ok := node.(map[string]any)
	if !ok {
		return
	}

	// Run first, so the objects it creates are recursed into below like any
	// other nested schema.
	normalizeBoolSchemas(obj)

	if t, ok := obj["type"].([]any); ok {
		pick := ""
		for _, m := range t {
			s, ok := m.(string)
			if !ok || s == "null" {
				continue
			}
			pick = s
			break
		}
		if pick == "" && len(t) > 0 {
			if s, ok := t[0].(string); ok {
				pick = s
			}
		}
		if pick != "" {
			obj["type"] = pick
		}
	}

	// Object properties — map of name → schema.
	if nested, ok := obj["properties"].(map[string]any); ok {
		coerceNullableTypes(nested)
	}
	// patternProperties — same shape as properties, just regex-keyed.
	if nested, ok := obj["patternProperties"].(map[string]any); ok {
		coerceNullableTypes(nested)
	}
	// $defs / definitions — JSON-Schema named schema bag.
	if nested, ok := obj["$defs"].(map[string]any); ok {
		coerceNullableTypes(nested)
	}
	if nested, ok := obj["definitions"].(map[string]any); ok {
		coerceNullableTypes(nested)
	}
	// Single nested schema fields.
	coerceSchema(obj["items"])
	coerceSchema(obj["additionalProperties"])
	coerceSchema(obj["contains"])
	coerceSchema(obj["not"])
	coerceSchema(obj["if"])
	coerceSchema(obj["then"])
	coerceSchema(obj["else"])
	coerceSchema(obj["propertyNames"])
	// Composition keywords — arrays of schemas.
	for _, key := range []string{"oneOf", "anyOf", "allOf"} {
		arr, ok := obj[key].([]any)
		if !ok {
			continue
		}
		for _, member := range arr {
			coerceSchema(member)
		}
	}
	// "items" can also be an array (tuple validation in older drafts).
	if arr, ok := obj["items"].([]any); ok {
		for _, member := range arr {
			coerceSchema(member)
		}
	}
	// "prefixItems" (2020-12 tuple).
	if arr, ok := obj["prefixItems"].([]any); ok {
		for _, member := range arr {
			coerceSchema(member)
		}
	}
}

func mcpPromptsFromTransport(ctx context.Context, session *mcp.ClientSession, arguments map[string]string) ([]openai.ChatCompletionMessage, error) {
	prompts, err := session.ListPrompts(ctx, nil)
	if err != nil {
		return nil, err
	}

	promptsList := []openai.ChatCompletionMessage{}

	for _, prompt := range prompts.Prompts {
		p, err := session.GetPrompt(ctx, &mcp.GetPromptParams{Name: prompt.Name, Arguments: arguments})
		if err != nil {
			return nil, err
		}
		for _, message := range p.Messages {
			switch message.Content.(type) {
			case *mcp.TextContent:
				promptsList = append(promptsList, openai.ChatCompletionMessage{
					Role:    string(message.Role),
					Content: message.Content.(*mcp.TextContent).Text,
				})
			}
		}
	}

	return promptsList, nil
}

// MCPToolFilter is invoked once per (session, tool) pair during the
// initial tool-discovery pass. Return false to drop the tool from the
// agent's discovered set (the LLM never sees it). A nil filter is
// equivalent to "always allow".
type MCPToolFilter = func(session *mcp.ClientSession, toolName string) bool

// probe the MCP remote and generate tools that are compliant with cogito
func mcpToolsFromTransport(ctx context.Context, session *mcp.ClientSession, filter MCPToolFilter) ([]ToolDefinitionInterface, error) {
	allTools := []ToolDefinitionInterface{}

	tools, err := session.ListTools(ctx, nil)
	if err != nil {
		xlog.Error("Error listing tools: %v", err)
		return nil, err
	}

	for _, tool := range tools.Tools {
		if filter != nil && !filter(session, tool.Name) {
			continue
		}
		dat, err := json.Marshal(tool.InputSchema)
		if err != nil {
			xlog.Error("Error marshalling input schema: %v", err)
			continue
		}

		var inputSchema toolInputSchema
		err = json.Unmarshal(dat, &inputSchema)
		if err != nil {
			xlog.Error("Error unmarshalling input schema: %v", err)
			continue
		}

		// Some MCP servers (e.g. modelcontextprotocol/go-sdk v1.4+)
		// emit JSON Schema 2020-12 "type": ["null", "array"] for
		// nullable fields like Go []string slices. langchaingo's
		// jsonschema.Definition.Type is a single string, so the
		// unmarshal would fail and silently drop the entire tool.
		// Coerce any type-array to its non-null string member before
		// unmarshalling so those tools stay discoverable.
		coerceNullableTypes(inputSchema.Properties)

		props := map[string]jsonschema.Definition{}
		dat, err = json.Marshal(inputSchema.Properties)
		if err != nil {
			xlog.Error("Error marshalling input schema: %v", err)
			continue
		}
		err = json.Unmarshal(dat, &props)
		if err != nil {
			xlog.Error("Error unmarshalling input schema properties: %v", err)
			continue
		}

		allTools = append(allTools, &mcpTool{
			name:        tool.Name,
			description: tool.Description,
			session:     session,
			ctx:         ctx,
			props:       props,
			inputSchema: inputSchema,
		})
	}

	return allTools, nil
}