path: root/internal/executor/agentmcp_test.go

package executor

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/thepeterstone/claudomator/internal/role"
)

// recordingChannel is a fake AgentChannel that records tool invocations.
type recordingChannel struct {
	asked               string
	summary             string
	spawned             []SubtaskSpec
	progress            []string
	spawnID             string
	proposedEpics       []EpicProposal
	epicID              string
	proposedRoleConfigs []role.RoleConfig
	roleConfigVersion   int
}

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 (c *recordingChannel) ProposeEpic(_ context.Context, spec EpicProposal) (string, error) {
	c.proposedEpics = append(c.proposedEpics, spec)
	return c.epicID, nil
}
func (c *recordingChannel) ProposeRoleConfig(_ context.Context, cfg role.RoleConfig) (int, error) {
	c.proposedRoleConfigs = append(c.proposedRoleConfigs, cfg)
	return c.roleConfigVersion, nil
}
func (c *recordingChannel) ReportVerdict(_ context.Context, approved bool, reasoning string) error {
	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)
	}
}

// TestAgentServer_SpawnSubtask_RolePassthrough proves the MCP transport
// (ContainerRunner-driven claude/gemini agents) can spawn role-typed
// subtasks too, mirroring the native tool-use loop's coverage in
// nativerunner_test.go.
func TestAgentServer_SpawnSubtask_RolePassthrough(t *testing.T) {
	ch := &recordingChannel{spawnID: "sub-2"}
	cs := connectInMemory(t, newAgentServer(ch))
	_, err := cs.CallTool(context.Background(), &mcp.CallToolParams{
		Name:      "spawn_subtask",
		Arguments: map[string]any{"name": "eval", "instructions": "evaluate", "role": "evaluator_quality"},
	})
	if err != nil {
		t.Fatalf("CallTool: %v", err)
	}
	if len(ch.spawned) != 1 || ch.spawned[0].Role != "evaluator_quality" {
		t.Errorf("subtask role not propagated: %+v", ch.spawned)
	}
}

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)
	}
}

func TestAgentServer_ProposeEpic(t *testing.T) {
	ch := &recordingChannel{epicID: "epic-1"}
	cs := connectInMemory(t, newAgentServer(ch))
	res, err := cs.CallTool(context.Background(), &mcp.CallToolParams{
		Name:      "propose_epic",
		Arguments: map[string]any{"name": "Checkout revamp", "description": "relaunch", "story_ids": []string{"s1", "s2"}},
	})
	if err != nil {
		t.Fatalf("CallTool: %v", err)
	}
	if len(ch.proposedEpics) != 1 {
		t.Fatalf("expected 1 proposed epic, got %d", len(ch.proposedEpics))
	}
	spec := ch.proposedEpics[0]
	if spec.Name != "Checkout revamp" || spec.Description != "relaunch" {
		t.Errorf("epic spec not propagated: %+v", spec)
	}
	if len(spec.StoryIDs) != 2 || spec.StoryIDs[0] != "s1" || spec.StoryIDs[1] != "s2" {
		t.Errorf("story_ids not propagated: %+v", spec.StoryIDs)
	}
	if txt := resultText(t, res); !strings.Contains(txt, "epic-1") {
		t.Errorf("expected returned epic ID in result, got %q", txt)
	}
}

// TestAgentServer_ProposeRoleConfig proves the propose_role_config MCP tool
// reaches AgentChannel.ProposeRoleConfig with the full role.RoleConfig shape
// decoded correctly, including a nested escalation_ladder tier -- mirroring
// TestAgentServer_ProposeEpic's coverage above for this phase's new tool.
func TestAgentServer_ProposeRoleConfig(t *testing.T) {
	ch := &recordingChannel{roleConfigVersion: 3}
	cs := connectInMemory(t, newAgentServer(ch))
	res, err := cs.CallTool(context.Background(), &mcp.CallToolParams{
		Name: "propose_role_config",
		Arguments: map[string]any{
			"role":          "builder",
			"system_prompt": "Double-check edge cases before finishing.",
			"escalation_ladder": []map[string]any{
				{
					"candidates":  []map[string]any{{"provider": "anthropic", "model": "claude-sonnet-4-6"}},
					"max_retries": 1,
				},
			},
		},
	})
	if err != nil {
		t.Fatalf("CallTool: %v", err)
	}
	if len(ch.proposedRoleConfigs) != 1 {
		t.Fatalf("expected 1 proposed role config, got %d", len(ch.proposedRoleConfigs))
	}
	cfg := ch.proposedRoleConfigs[0]
	if cfg.Role != "builder" || cfg.SystemPrompt != "Double-check edge cases before finishing." {
		t.Errorf("role config not propagated: %+v", cfg)
	}
	if len(cfg.EscalationLadder) != 1 || len(cfg.EscalationLadder[0].Candidates) != 1 ||
		cfg.EscalationLadder[0].Candidates[0].Provider != "anthropic" || cfg.EscalationLadder[0].Candidates[0].Model != "claude-sonnet-4-6" {
		t.Errorf("escalation_ladder not propagated: %+v", cfg.EscalationLadder)
	}
	if txt := resultText(t, res); !strings.Contains(txt, "builder") || !strings.Contains(txt, "3") {
		t.Errorf("expected role name and version in result, got %q", txt)
	}
}

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")
	}
}