1 files changed, 218 insertions, 0 deletions
diff --git a/internal/executor/agentmcp_test.go b/internal/executor/agentmcp_test.go
new file mode 100644
index 0000000..3973af0
--- /dev/null
+++ b/internal/executor/agentmcp_test.go
@@ -0,0 +1,218 @@
+package executor
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/modelcontextprotocol/go-sdk/mcp"
+)
+
+// recordingChannel is a fake AgentChannel that records tool invocations.
+type recordingChannel struct {
+	asked    string
+	summary  string
+	spawned  []SubtaskSpec
+	progress []string
+	spawnID  string
+}
+
+func (c *recordingChannel) AskUser(_ context.Context, q string) (string, error) {
+	c.asked = q
+	return "", ErrAgentBlocked
+}
+func (c *recordingChannel) ReportSummary(_ context.Context, s string) error {
+	c.summary = s
+	return nil
+}
+func (c *recordingChannel) SpawnSubtask(_ context.Context, spec SubtaskSpec) (string, error) {
+	c.spawned = append(c.spawned, spec)
+	return c.spawnID, nil
+}
+func (c *recordingChannel) RecordProgress(_ context.Context, m string) error {
+	c.progress = append(c.progress, m)
+	return nil
+}
+
+func resultText(t *testing.T, res *mcp.CallToolResult) string {
+	t.Helper()
+	if len(res.Content) == 0 {
+		t.Fatal("expected content in tool result")
+	}
+	tc, ok := res.Content[0].(*mcp.TextContent)
+	if !ok {
+		t.Fatalf("expected TextContent, got %T", res.Content[0])
+	}
+	return tc.Text
+}
+
+func TestRegistry_MintLookupRevoke(t *testing.T) {
+	reg := NewRegistry()
+	tok, err := reg.Mint(&recordingChannel{})
+	if err != nil {
+		t.Fatalf("Mint: %v", err)
+	}
+	if tok == "" {
+		t.Fatal("expected non-empty token")
+	}
+	if _, ok := reg.server(tok); !ok {
+		t.Error("expected server for minted token")
+	}
+	if _, ok := reg.server("nonsense"); ok {
+		t.Error("expected no server for unknown token")
+	}
+	reg.Revoke(tok)
+	if _, ok := reg.server(tok); ok {
+		t.Error("expected no server after revoke")
+	}
+}
+
+func TestRegistry_MintUniqueTokens(t *testing.T) {
+	reg := NewRegistry()
+	a, _ := reg.Mint(&recordingChannel{})
+	b, _ := reg.Mint(&recordingChannel{})
+	if a == b {
+		t.Error("expected unique tokens per mint")
+	}
+}
+
+// connectInMemory wires a client to a server over the SDK's in-memory transport.
+func connectInMemory(t *testing.T, srv *mcp.Server) *mcp.ClientSession {
+	t.Helper()
+	ctx := context.Background()
+	clientT, serverT := mcp.NewInMemoryTransports()
+	if _, err := srv.Connect(ctx, serverT, nil); err != nil {
+		t.Fatalf("server connect: %v", err)
+	}
+	client := mcp.NewClient(&mcp.Implementation{Name: "test", Version: "1"}, nil)
+	cs, err := client.Connect(ctx, clientT, nil)
+	if err != nil {
+		t.Fatalf("client connect: %v", err)
+	}
+	t.Cleanup(func() { _ = cs.Close() })
+	return cs
+}
+
+func TestAgentServer_AskUser_RecordsAndInstructsStop(t *testing.T) {
+	ch := &recordingChannel{}
+	cs := connectInMemory(t, newAgentServer(ch))
+	res, err := cs.CallTool(context.Background(), &mcp.CallToolParams{
+		Name:      "ask_user",
+		Arguments: map[string]any{"question": "Proceed?", "options": []string{"yes", "no"}},
+	})
+	if err != nil {
+		t.Fatalf("CallTool: %v", err)
+	}
+	if !strings.Contains(ch.asked, `"text":"Proceed?"`) || !strings.Contains(ch.asked, "yes") {
+		t.Errorf("channel did not receive question+options: %q", ch.asked)
+	}
+	if txt := resultText(t, res); !strings.Contains(strings.ToLower(txt), "recorded") {
+		t.Errorf("expected stop instruction, got %q", txt)
+	}
+}
+
+func TestAgentServer_ReportSummary(t *testing.T) {
+	ch := &recordingChannel{}
+	cs := connectInMemory(t, newAgentServer(ch))
+	if _, err := cs.CallTool(context.Background(), &mcp.CallToolParams{
+		Name:      "report_summary",
+		Arguments: map[string]any{"summary": "did the work"},
+	}); err != nil {
+		t.Fatalf("CallTool: %v", err)
+	}
+	if ch.summary != "did the work" {
+		t.Errorf("summary not recorded, got %q", ch.summary)
+	}
+}
+
+func TestAgentServer_SpawnSubtask(t *testing.T) {
+	ch := &recordingChannel{spawnID: "sub-1"}
+	cs := connectInMemory(t, newAgentServer(ch))
+	res, err := cs.CallTool(context.Background(), &mcp.CallToolParams{
+		Name:      "spawn_subtask",
+		Arguments: map[string]any{"name": "child", "instructions": "do it", "model": "sonnet"},
+	})
+	if err != nil {
+		t.Fatalf("CallTool: %v", err)
+	}
+	if len(ch.spawned) != 1 || ch.spawned[0].Name != "child" || ch.spawned[0].Instructions != "do it" || ch.spawned[0].Model != "sonnet" {
+		t.Errorf("subtask spec not propagated: %+v", ch.spawned)
+	}
+	if txt := resultText(t, res); !strings.Contains(txt, "sub-1") {
+		t.Errorf("expected returned subtask ID in result, got %q", txt)
+	}
+}
+
+func TestAgentServer_RecordProgress(t *testing.T) {
+	ch := &recordingChannel{}
+	cs := connectInMemory(t, newAgentServer(ch))
+	if _, err := cs.CallTool(context.Background(), &mcp.CallToolParams{
+		Name:      "record_progress",
+		Arguments: map[string]any{"message": "halfway"},
+	}); err != nil {
+		t.Fatalf("CallTool: %v", err)
+	}
+	if len(ch.progress) != 1 || ch.progress[0] != "halfway" {
+		t.Errorf("progress not recorded: %+v", ch.progress)
+	}
+}
+
+type bearerRT struct {
+	token string
+	base  http.RoundTripper
+}
+
+func (b bearerRT) RoundTrip(r *http.Request) (*http.Response, error) {
+	r = r.Clone(r.Context())
+	if b.token != "" {
+		r.Header.Set("Authorization", "Bearer "+b.token)
+	}
+	return b.base.RoundTrip(r)
+}
+
+func TestAgentMCPHandler_HTTP_AuthAndDispatch(t *testing.T) {
+	ch := &recordingChannel{}
+	reg := NewRegistry()
+	tok, _ := reg.Mint(ch)
+
+	httpSrv := httptest.NewServer(NewAgentMCPHandler(reg))
+	defer httpSrv.Close()
+
+	ctx := context.Background()
+	client := mcp.NewClient(&mcp.Implementation{Name: "test", Version: "1"}, nil)
+	cs, err := client.Connect(ctx, &mcp.StreamableClientTransport{
+		Endpoint:   httpSrv.URL,
+		HTTPClient: &http.Client{Transport: bearerRT{token: tok, base: http.DefaultTransport}},
+	}, nil)
+	if err != nil {
+		t.Fatalf("client connect with valid token: %v", err)
+	}
+	defer cs.Close()
+
+	if _, err := cs.CallTool(ctx, &mcp.CallToolParams{
+		Name:      "record_progress",
+		Arguments: map[string]any{"message": "over http"},
+	}); err != nil {
+		t.Fatalf("CallTool over HTTP: %v", err)
+	}
+	if len(ch.progress) != 1 || ch.progress[0] != "over http" {
+		t.Errorf("HTTP dispatch did not reach channel: %+v", ch.progress)
+	}
+}
+
+func TestAgentMCPHandler_HTTP_BadTokenRejected(t *testing.T) {
+	reg := NewRegistry()
+	httpSrv := httptest.NewServer(NewAgentMCPHandler(reg))
+	defer httpSrv.Close()
+
+	client := mcp.NewClient(&mcp.Implementation{Name: "test", Version: "1"}, nil)
+	_, err := client.Connect(context.Background(), &mcp.StreamableClientTransport{
+		Endpoint:   httpSrv.URL,
+		HTTPClient: &http.Client{Transport: bearerRT{token: "invalid", base: http.DefaultTransport}},
+	}, nil)
+	if err == nil {
+		t.Fatal("expected connect to fail with invalid token")
+	}
+}