Add agent library

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
2026-03-13 14:30:47 +01:00
parent 158c36d9ab
commit 9fd251ee5d
40 changed files with 9231 additions and 38 deletions
--- a/pkg/agent/mcp_test.go
+++ b/pkg/agent/mcp_test.go
@@ -0,0 +1,369 @@
+// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS SOFTWARE.
+
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"sync"
+	"testing"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.probo.inc/probo/pkg/llm"
+)
+
+func TestNewMCPServer(t *testing.T) {
+	t.Parallel()
+
+	s := NewMCPServer("my-server", nil)
+
+	require.NotNil(t, s)
+	assert.Equal(t, "my-server", s.name)
+	assert.Nil(t, s.session)
+	assert.False(t, s.toolsCached)
+	assert.Nil(t, s.cachedTools)
+}
+
+func TestMCPServer_Name(t *testing.T) {
+	t.Parallel()
+
+	s := &MCPServer{name: "test-server"}
+	assert.Equal(t, "test-server", s.Name())
+}
+
+func TestMCPServer_Tools(t *testing.T) {
+	t.Parallel()
+
+	t.Run(
+		"returns cached tools without calling session",
+		func(t *testing.T) {
+			t.Parallel()
+
+			tool := &mcpTool{
+				name:        "cached-tool",
+				description: "A cached tool",
+				inputSchema: json.RawMessage(`{"type":"object"}`),
+			}
+
+			s := &MCPServer{name: "s"}
+			s.cachedTools = []Tool{tool}
+			s.toolsCached = true
+
+			tools, err := s.Tools(context.Background())
+			require.NoError(t, err)
+			require.Len(t, tools, 1)
+			assert.Equal(t, "cached-tool", tools[0].Name())
+		},
+	)
+
+	t.Run(
+		"returns a defensive copy of cached tools",
+		func(t *testing.T) {
+			t.Parallel()
+
+			tool := &mcpTool{
+				name:        "tool-a",
+				description: "Tool A",
+				inputSchema: json.RawMessage(`{}`),
+			}
+
+			s := &MCPServer{name: "s"}
+			s.cachedTools = []Tool{tool}
+			s.toolsCached = true
+
+			tools1, err := s.Tools(context.Background())
+			require.NoError(t, err)
+
+			tools2, err := s.Tools(context.Background())
+			require.NoError(t, err)
+
+			// Mutating one slice must not affect the other.
+			tools1[0] = nil
+			assert.NotNil(t, tools2[0])
+
+			// Underlying cache must be untouched.
+			s.mu.RLock()
+			assert.NotNil(t, s.cachedTools[0])
+			s.mu.RUnlock()
+		},
+	)
+
+	t.Run(
+		"concurrent reads return consistent results",
+		func(t *testing.T) {
+			t.Parallel()
+
+			s := &MCPServer{name: "s"}
+			s.cachedTools = []Tool{
+				&mcpTool{
+					name:        "echo",
+					description: "Echoes the input",
+					inputSchema: json.RawMessage(`{"type":"object"}`),
+				},
+			}
+			s.toolsCached = true
+
+			const goroutines = 50
+			var wg sync.WaitGroup
+			wg.Add(goroutines)
+
+			for range goroutines {
+				go func() {
+					defer wg.Done()
+
+					tools, err := s.Tools(context.Background())
+					require.NoError(t, err)
+					assert.Len(t, tools, 1)
+					assert.Equal(t, "echo", tools[0].Name())
+				}()
+			}
+
+			wg.Wait()
+		},
+	)
+}
+
+func TestMCPServer_ResetCache(t *testing.T) {
+	t.Parallel()
+
+	t.Run(
+		"clears cached tools",
+		func(t *testing.T) {
+			t.Parallel()
+
+			s := &MCPServer{name: "s"}
+			s.cachedTools = []Tool{
+				&mcpTool{
+					name:        "tool1",
+					description: "A tool",
+					inputSchema: json.RawMessage(`{"type":"object"}`),
+				},
+			}
+			s.toolsCached = true
+
+			tools, err := s.Tools(context.Background())
+			require.NoError(t, err)
+			require.Len(t, tools, 1)
+
+			s.ResetCache()
+
+			s.mu.RLock()
+			assert.False(t, s.toolsCached)
+			assert.Nil(t, s.cachedTools)
+			s.mu.RUnlock()
+		},
+	)
+
+	t.Run(
+		"is safe to call on fresh server",
+		func(t *testing.T) {
+			t.Parallel()
+
+			s := &MCPServer{name: "s"}
+			s.ResetCache()
+
+			s.mu.RLock()
+			assert.False(t, s.toolsCached)
+			assert.Nil(t, s.cachedTools)
+			s.mu.RUnlock()
+		},
+	)
+
+	t.Run(
+		"concurrent resets do not race",
+		func(t *testing.T) {
+			t.Parallel()
+
+			s := &MCPServer{name: "s"}
+			s.cachedTools = []Tool{
+				&mcpTool{
+					name:        "tool1",
+					description: "A tool",
+					inputSchema: json.RawMessage(`{"type":"object"}`),
+				},
+			}
+			s.toolsCached = true
+
+			const goroutines = 50
+			var wg sync.WaitGroup
+			wg.Add(goroutines)
+
+			for range goroutines {
+				go func() {
+					defer wg.Done()
+					s.ResetCache()
+				}()
+			}
+
+			wg.Wait()
+
+			s.mu.RLock()
+			assert.False(t, s.toolsCached)
+			assert.Nil(t, s.cachedTools)
+			s.mu.RUnlock()
+		},
+	)
+}
+
+func TestMCPTool_Name(t *testing.T) {
+	t.Parallel()
+
+	tool := &mcpTool{name: "get_weather"}
+	assert.Equal(t, "get_weather", tool.Name())
+}
+
+func TestMCPTool_Definition(t *testing.T) {
+	t.Parallel()
+
+	schema := json.RawMessage(`{"type":"object","properties":{"city":{"type":"string"}}}`)
+
+	tool := &mcpTool{
+		name:        "get_weather",
+		description: "Returns current weather for a city",
+		inputSchema: schema,
+	}
+
+	def := tool.Definition()
+
+	assert.Equal(t, llm.Tool{
+		Name:        "get_weather",
+		Description: "Returns current weather for a city",
+		Parameters:  schema,
+	}, def)
+}
+
+func TestMCPTool_Execute(t *testing.T) {
+	t.Parallel()
+
+	t.Run(
+		"invalid JSON arguments returns error result",
+		func(t *testing.T) {
+			t.Parallel()
+
+			tool := &mcpTool{name: "my_tool"}
+
+			result, err := tool.Execute(context.Background(), "not-json")
+			require.NoError(t, err)
+			assert.True(t, result.IsError)
+			assert.Contains(t, result.Content, "Invalid arguments")
+		},
+	)
+
+	t.Run(
+		"malformed JSON arguments returns error result",
+		func(t *testing.T) {
+			t.Parallel()
+
+			tool := &mcpTool{name: "my_tool"}
+
+			result, err := tool.Execute(context.Background(), `{"key": }`)
+			require.NoError(t, err)
+			assert.True(t, result.IsError)
+			assert.Contains(t, result.Content, "Invalid arguments")
+		},
+	)
+
+	t.Run(
+		"JSON array arguments returns error result",
+		func(t *testing.T) {
+			t.Parallel()
+
+			tool := &mcpTool{name: "my_tool"}
+
+			// Execute expects a JSON object, not an array.
+			result, err := tool.Execute(context.Background(), `[1,2,3]`)
+			require.NoError(t, err)
+			assert.True(t, result.IsError)
+			assert.Contains(t, result.Content, "Invalid arguments")
+		},
+	)
+}
+
+func TestExtractMCPContent(t *testing.T) {
+	t.Parallel()
+
+	t.Run(
+		"nil result returns empty",
+		func(t *testing.T) {
+			t.Parallel()
+			assert.Equal(t, "", extractMCPContent(nil))
+		},
+	)
+
+	t.Run(
+		"empty content returns empty",
+		func(t *testing.T) {
+			t.Parallel()
+			result := &mcp.CallToolResult{Content: []mcp.Content{}}
+			assert.Equal(t, "", extractMCPContent(result))
+		},
+	)
+
+	t.Run(
+		"single text content",
+		func(t *testing.T) {
+			t.Parallel()
+			result := &mcp.CallToolResult{
+				Content: []mcp.Content{
+					&mcp.TextContent{Text: "hello world"},
+				},
+			}
+			assert.Equal(t, "hello world", extractMCPContent(result))
+		},
+	)
+
+	t.Run(
+		"multiple text contents joined by newline",
+		func(t *testing.T) {
+			t.Parallel()
+			result := &mcp.CallToolResult{
+				Content: []mcp.Content{
+					&mcp.TextContent{Text: "line one"},
+					&mcp.TextContent{Text: "line two"},
+				},
+			}
+			assert.Equal(t, "line one\nline two", extractMCPContent(result))
+		},
+	)
+
+	t.Run(
+		"non-text content is skipped",
+		func(t *testing.T) {
+			t.Parallel()
+			result := &mcp.CallToolResult{
+				Content: []mcp.Content{
+					&mcp.TextContent{Text: "text part"},
+					&mcp.ImageContent{Data: []byte("base64data"), MIMEType: "image/png"},
+				},
+			}
+			assert.Equal(t, "text part", extractMCPContent(result))
+		},
+	)
+
+	t.Run(
+		"only non-text content returns empty",
+		func(t *testing.T) {
+			t.Parallel()
+			result := &mcp.CallToolResult{
+				Content: []mcp.Content{
+					&mcp.ImageContent{Data: []byte("base64data"), MIMEType: "image/png"},
+				},
+			}
+			assert.Equal(t, "", extractMCPContent(result))
+		},
+	)
+}