robinhood-agentic-mcp/client/session_test.go

package client

import (
	"encoding/json"
	"errors"
	"strings"
	"testing"

	mcp "github.com/modelcontextprotocol/go-sdk/mcp"
)

func TestToolJSON(t *testing.T) {
	t.Parallel()
	tests := []struct {
		name string
		res  *mcp.CallToolResult
		want string
	}{
		{
			name: "structured-only",
			res: &mcp.CallToolResult{
				StructuredContent: map[string]any{"symbol": "MU"},
			},
			want: `{"symbol":"MU"}`,
		},
		{
			name: "text JSON",
			res: &mcp.CallToolResult{
				Content: []mcp.Content{&mcp.TextContent{Text: `{"symbol":"MU"}`}},
			},
			want: `{"symbol":"MU"}`,
		},
		{
			name: "text non-JSON + structured",
			res: &mcp.CallToolResult{
				StructuredContent: json.RawMessage(`{"symbol":"MU"}`),
				Content:           []mcp.Content{&mcp.TextContent{Text: "not json"}},
			},
			want: `{"symbol":"MU"}`,
		},
		{
			name: "empty",
			res:  &mcp.CallToolResult{},
			want: `{}`,
		},
		{
			name: "wrapped JSON in prose",
			res: &mcp.CallToolResult{
				Content: []mcp.Content{&mcp.TextContent{Text: `here: {"symbol":"MU"} ok`}},
			},
			want: `{"symbol":"MU"}`,
		},
		{
			name: "one TextContent is JSON",
			res: &mcp.CallToolResult{
				Content: []mcp.Content{
					&mcp.TextContent{Text: "note: "},
					&mcp.TextContent{Text: `{"symbol":"MU"}`},
				},
			},
			want: `{"symbol":"MU"}`,
		},
		{
			name: "concatenated text JSON",
			res: &mcp.CallToolResult{
				Content: []mcp.Content{
					&mcp.TextContent{Text: `{"symbol":`},
					&mcp.TextContent{Text: `"MU"}`},
				},
			},
			want: `{"symbol":"MU"}`,
		},
		{
			name: "wrapped JSON array",
			res: &mcp.CallToolResult{
				Content: []mcp.Content{&mcp.TextContent{Text: `items: [1,{"a":2}] done`}},
			},
			want: `[1,{"a":2}]`,
		},
		{
			name: "wrapped JSON with brace in string",
			res: &mcp.CallToolResult{
				Content: []mcp.Content{&mcp.TextContent{Text: `here {"msg":"say } hi"} x`}},
			},
			want: `{"msg":"say } hi"}`,
		},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()
			got, err := toolJSON(tc.res)
			if err != nil {
				t.Fatal(err)
			}
			if string(got) != tc.want {
				t.Fatalf("got %s want %s", got, tc.want)
			}
		})
	}
}

func TestToolJSON_nilRes(t *testing.T) {
	t.Parallel()
	_, err := toolJSON(nil)
	if err == nil || err.Error() != "empty tool result" {
		t.Fatalf("%v", err)
	}
}

func TestToolJSON_getError(t *testing.T) {
	t.Parallel()
	inner := errors.New("tool failed")
	res := &mcp.CallToolResult{}
	res.SetError(inner)
	_, err := toolJSON(res)
	if !errors.Is(err, inner) {
		t.Fatalf("%v", err)
	}
}

func TestToolJSON_isErrorText(t *testing.T) {
	t.Parallel()
	// Jerry dump: IsError=true, StructuredContent=null, Text=`end must be after start`, GetError=nil.
	res := &mcp.CallToolResult{
		IsError:           true,
		StructuredContent: nil,
		Content:           []mcp.Content{&mcp.TextContent{Text: "end must be after start"}},
	}
	if res.GetError() != nil {
		t.Fatal("GetError should be nil on a client-shaped result")
	}
	_, err := toolJSON(res)
	if err == nil {
		t.Fatal("expected error")
	}
	if err.Error() != "end must be after start" {
		t.Fatalf("%v", err)
	}
	if strings.Contains(err.Error(), "non-json tool result") {
		t.Fatalf("got non-json wrapping: %v", err)
	}
}

func TestToolJSON_isErrorOverStructured(t *testing.T) {
	t.Parallel()
	_, err := toolJSON(&mcp.CallToolResult{
		IsError:           true,
		StructuredContent: map[string]any{"symbol": "MU"},
		Content:           []mcp.Content{&mcp.TextContent{Text: "end must be after start"}},
	})
	if err == nil {
		t.Fatal("expected error")
	}
	if err.Error() != "end must be after start" {
		t.Fatalf("%v", err)
	}
	if strings.Contains(err.Error(), "non-json tool result") {
		t.Fatalf("got non-json wrapping: %v", err)
	}
}

func TestToolJSON_isErrorEmpty(t *testing.T) {
	t.Parallel()
	tests := []*mcp.CallToolResult{
		{IsError: true},
		{IsError: true, Content: []mcp.Content{&mcp.TextContent{Text: "  \n"}}},
	}
	for _, res := range tests {
		_, err := toolJSON(res)
		if err == nil {
			t.Fatal("expected error")
		}
		if err.Error() != "tool error" {
			t.Fatalf("%v", err)
		}
		if strings.Contains(err.Error(), "non-json tool result") {
			t.Fatalf("got non-json wrapping: %v", err)
		}
	}
}

func TestToolJSON_nonJSONText(t *testing.T) {
	t.Parallel()
	_, err := toolJSON(&mcp.CallToolResult{
		Content: []mcp.Content{&mcp.TextContent{Text: "not json"}},
	})
	if err == nil {
		t.Fatal("expected error")
	}
	msg := err.Error()
	if !strings.Contains(msg, "non-json tool result") {
		t.Fatalf("missing non-json tool result: %v", err)
	}
	for _, want := range []string{`isError=false`, `types=[text]`, `structured=nil`, `textLen=8`, `prefix="not json"`} {
		if !strings.Contains(msg, want) {
			t.Fatalf("missing %q in %v", want, err)
		}
	}
}