//go:build e2e

// Package protocol contains e2e tests for the MCP JSON-RPC protocol
// infrastructure and the AgentTeam WebSocket bidirectional protocol.
package protocol

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/pingcap/failpoint"

	"github.com/tidbcloud/AgentTeam/desktop/internal/api/mcp"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

// ---------------------------------------------------------------------------
// Mock MCP handler
// ---------------------------------------------------------------------------

type testMCPHandler struct {
	tools  []types.MCPToolDef
	callFn func(ctx context.Context, agentID, name string, args json.RawMessage) *types.MCPToolResult
}

type toolNameResult struct {
	Name string `json:"name"`
}

func (h *testMCPHandler) Tools() []types.MCPToolDef         { return h.tools }
func (h *testMCPHandler) SkillDoc(_ types.ToolNamer) string { return "" }
func (h *testMCPHandler) SkillDocSplit(types.ToolNamer) types.SkillDocParts {
	return types.SkillDocParts{}
}
func (h *testMCPHandler) CallTool(ctx context.Context, agentID, name string, args json.RawMessage) *types.MCPToolResult {
	if h.callFn != nil {
		return h.callFn(ctx, agentID, name, args)
	}
	return types.NewToolText("ok: " + name)
}

func defaultTestTools() []types.MCPToolDef {
	return []types.MCPToolDef{
		{Name: "message_send", Description: "Send a message", InputSchema: types.MCPInputSchema{Type: "object"}},
		{Name: "message_history_read", Description: "Read target history", InputSchema: types.MCPInputSchema{Type: "object"}},
	}
}

// ---------------------------------------------------------------------------
// MCP router setup
// ---------------------------------------------------------------------------

func setupMCPRouter(t *testing.T, handler types.MCPHandler) (*gin.Engine, string) {
	t.Helper()
	cfgDir := t.TempDir()
	failpoint.Enable("agentteam-config-dir", fmt.Sprintf(`return("%s")`, cfgDir))
	t.Cleanup(func() { failpoint.Disable("agentteam-config-dir") })
	_ = token.Default()

	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(gin.Recovery())
	group := r.Group("/mcp/server/:serverId/:agentId", mcp.AccessControl())
	mcp.RegisterHandlerRoutes(group, func(serverID string) (types.MCPHandler, bool) {
		return handler, true
	})

	at, err := token.Default().GenerateAgentTokenForServer("agent-1", "test-server")
	if err != nil {
		t.Fatalf("generate agent token: %v", err)
	}
	return r, at
}

func doMCPPost(t *testing.T, router *gin.Engine, agentToken string, body interface{}) (int, json.RawMessage) {
	t.Helper()
	return doMCPPostWithHeaders(t, router, agentToken, nil, body)
}

func doMCPPostWithHeaders(t *testing.T, router *gin.Engine, agentToken string, headers map[string]string, body interface{}) (int, json.RawMessage) {
	t.Helper()
	var reqBody io.Reader
	if body != nil {
		b, _ := json.Marshal(body)
		reqBody = bytes.NewReader(b)
	}
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-1", reqBody)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+agentToken)
	for key, value := range headers {
		req.Header.Set(key, value)
	}
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	return w.Code, w.Body.Bytes()
}

// ---------------------------------------------------------------------------
// Tests: JSON-RPC lifecycle
// ---------------------------------------------------------------------------

func TestMCP_Initialize(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	code, raw := doMCPPost(t, router, at, map[string]interface{}{"jsonrpc": "2.0", "method": "initialize", "id": 1})
	if code != 200 {
		t.Fatalf("expected 200, got %d: %s", code, raw)
	}
	var resp mcpRPCResponse
	json.Unmarshal(raw, &resp)
	if resp.Error != nil {
		t.Fatalf("unexpected error: %v", resp.Error)
	}
	rm, _ := resp.Result.(map[string]interface{})
	if rm["protocolVersion"] != "2024-11-05" {
		t.Errorf("protocolVersion=%v, want 2024-11-05", rm["protocolVersion"])
	}
}

func TestMCP_Ping(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	code, raw := doMCPPost(t, router, at, map[string]interface{}{"jsonrpc": "2.0", "method": "ping", "id": 2})
	if code != 200 {
		t.Fatalf("expected 200, got %d: %s", code, raw)
	}
	var resp mcpRPCResponse
	json.Unmarshal(raw, &resp)
	if resp.Error != nil {
		t.Fatalf("unexpected error: %v", resp.Error)
	}
}

func TestMCP_ToolsList(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	code, raw := doMCPPost(t, router, at, map[string]interface{}{"jsonrpc": "2.0", "method": "tools/list", "id": 3})
	if code != 200 {
		t.Fatalf("expected 200, got %d: %s", code, raw)
	}
	var resp struct {
		Result struct {
			Tools []toolNameResult `json:"tools"`
		} `json:"result"`
	}
	json.Unmarshal(raw, &resp)
	if len(resp.Result.Tools) != 2 {
		t.Fatalf("expected 2 tools, got %d", len(resp.Result.Tools))
	}
}

func TestMCP_LegacyDisabledToolsHeaderCannotFilterInventoryOrCalls(t *testing.T) {
	callCount := 0
	router, at := setupMCPRouter(t, &testMCPHandler{
		tools: defaultTestTools(),
		callFn: func(_ context.Context, _, _ string, _ json.RawMessage) *types.MCPToolResult {
			callCount++
			return types.NewToolText("ok")
		},
	})
	headers := map[string]string{
		// This retired bridge header is deliberately inert. Runtime/session
		// headers must not change the Desktop-owned MCP inventory or dispatch.
		"X-AgentTeam-MCP-Disabled-Tools": "message_send,message_history_read",
	}

	code, raw := doMCPPostWithHeaders(t, router, at, headers, map[string]interface{}{
		"jsonrpc": "2.0",
		"method":  "tools/list",
		"id":      30,
	})
	if code != http.StatusOK {
		t.Fatalf("expected 200, got %d: %s", code, raw)
	}
	var listResp struct {
		Result struct {
			Tools []toolNameResult `json:"tools"`
		} `json:"result"`
	}
	if err := json.Unmarshal(raw, &listResp); err != nil {
		t.Fatalf("decode tools/list response: %v", err)
	}
	if len(listResp.Result.Tools) != len(defaultTestTools()) {
		t.Fatalf("tools/list returned %d tools, want %d", len(listResp.Result.Tools), len(defaultTestTools()))
	}

	code, raw = doMCPPostWithHeaders(t, router, at, headers, map[string]interface{}{
		"jsonrpc": "2.0",
		"method":  "tools/call",
		"id":      31,
		"params":  map[string]interface{}{"name": "message_send", "arguments": map[string]interface{}{}},
	})
	if code != http.StatusOK {
		t.Fatalf("expected 200, got %d: %s", code, raw)
	}
	if callCount != 1 {
		t.Fatalf("tools/call invoked handler %d time(s), want 1", callCount)
	}
}

func TestMCP_ToolsCall(t *testing.T) {
	called := false
	router, at := setupMCPRouter(t, &testMCPHandler{
		tools: defaultTestTools(),
		callFn: func(_ context.Context, _, name string, _ json.RawMessage) *types.MCPToolResult {
			called = true
			return types.NewToolText("sent: " + name)
		},
	})
	code, _ := doMCPPost(t, router, at, map[string]interface{}{
		"jsonrpc": "2.0", "method": "tools/call", "id": 4,
		"params": map[string]interface{}{"name": "message_send", "arguments": map[string]interface{}{"content": "hello"}},
	})
	if code != 200 {
		t.Fatalf("expected 200, got %d", code)
	}
	if !called {
		t.Error("expected CallTool to be invoked")
	}
}

func TestMCP_MethodNotFound(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	code, raw := doMCPPost(t, router, at, map[string]interface{}{"jsonrpc": "2.0", "method": "nonexistent", "id": 5})
	if code != 200 {
		t.Fatalf("expected 200, got %d: %s", code, raw)
	}
	var resp mcpRPCResponse
	json.Unmarshal(raw, &resp)
	if resp.Error == nil || resp.Error.Code != -32601 {
		t.Fatalf("expected -32601, got %v", resp.Error)
	}
}

func TestMCP_InvalidParams(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	code, raw := doMCPPost(t, router, at, map[string]interface{}{"jsonrpc": "2.0", "method": "tools/call", "id": 6, "params": "not an object"})
	if code != 200 {
		t.Fatalf("expected 200, got %d", code)
	}
	var resp mcpRPCResponse
	json.Unmarshal(raw, &resp)
	if resp.Error == nil || resp.Error.Code != -32602 {
		t.Fatalf("expected -32602, got %v", resp.Error)
	}
}

// ---------------------------------------------------------------------------
// Tests: Notifications and batch
// ---------------------------------------------------------------------------

func TestMCP_NotificationReturns202(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-1",
		strings.NewReader(`{"jsonrpc":"2.0","method":"notifications/initialized"}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 202 {
		t.Fatalf("expected 202, got %d", w.Code)
	}
}

func TestMCP_BatchAllNotificationsReturns202(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	batch := `[{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","method":"notifications/cancelled"}]`
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-1", strings.NewReader(batch))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 202 {
		t.Fatalf("expected 202, got %d", w.Code)
	}
}

func TestMCP_BatchMixed(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	batch := `[{"jsonrpc":"2.0","method":"ping","id":1},{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","method":"ping","id":2}]`
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-1", strings.NewReader(batch))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 200 {
		t.Fatalf("expected 200, got %d", w.Code)
	}
	var responses []mcpRPCResponse
	json.Unmarshal(w.Body.Bytes(), &responses)
	if len(responses) != 2 {
		t.Fatalf("expected 2 responses, got %d", len(responses))
	}
}

// ---------------------------------------------------------------------------
// Tests: Auth middleware
// ---------------------------------------------------------------------------

func TestMCP_Auth_MissingToken(t *testing.T) {
	router, _ := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-1", strings.NewReader(`{}`))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 401 {
		t.Fatalf("expected 401, got %d", w.Code)
	}
}

func TestMCP_Auth_InvalidToken(t *testing.T) {
	router, _ := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-1", strings.NewReader(`{}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer invalid-token")
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 401 {
		t.Fatalf("expected 401, got %d", w.Code)
	}
}

func TestMCP_Auth_WrongAgent(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-wrong",
		strings.NewReader(`{"jsonrpc":"2.0","method":"ping","id":1}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 403 {
		t.Fatalf("expected 403, got %d", w.Code)
	}
}

// ---------------------------------------------------------------------------
// Tests: Malformed input
// ---------------------------------------------------------------------------

func TestMCP_MalformedJSON(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-1", strings.NewReader(`{invalid json`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 200 {
		t.Fatalf("expected 200, got %d", w.Code)
	}
	var resp mcpRPCResponse
	json.Unmarshal(w.Body.Bytes(), &resp)
	if resp.Error == nil || resp.Error.Code != -32700 {
		t.Fatalf("expected -32700, got %v", resp.Error)
	}
}

func TestMCP_EmptyBody(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-1", strings.NewReader(""))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 400 {
		t.Fatalf("expected 400, got %d", w.Code)
	}
}

// ---------------------------------------------------------------------------
// Tests: Nil Properties, error tools, channel not found
// ---------------------------------------------------------------------------

func TestMCP_ToolsList_NilProperties(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{
		tools: []types.MCPToolDef{{Name: "test_tool", Description: "Test", InputSchema: types.MCPInputSchema{Type: "object"}}},
	})
	code, raw := doMCPPost(t, router, at, map[string]interface{}{"jsonrpc": "2.0", "method": "tools/list", "id": 1})
	if code != 200 {
		t.Fatalf("expected 200, got %d: %s", code, raw)
	}
	var resp struct {
		Result struct {
			Tools []struct {
				InputSchema struct {
					Properties json.RawMessage `json:"properties"`
				} `json:"inputSchema"`
			} `json:"tools"`
		} `json:"result"`
	}
	json.Unmarshal(raw, &resp)
	if string(resp.Result.Tools[0].InputSchema.Properties) == "null" {
		t.Error("properties should not be null, should be {}")
	}
}

// TestMCP_ToolsList_FullAgentTeamSchema registers a representative AgentTeam MCP surface and
// validates the tools/list JSON response against the constraints that
// OpenCode's Zod v4 validation enforces:
//   - Every inputSchema.properties MUST be a JSON object (never null).
//   - Every inputSchema.type MUST be "object".
//   - Nested property objects with type "object" must have properties as
//     an object (or omitted), never null.
//   - The "required" field, when present, must be a JSON array.
func TestMCP_ToolsList_FullAgentTeamSchema(t *testing.T) {
	allTools := []types.MCPToolDef{
		{Name: "message_send", Description: "Send a message", InputSchema: types.MCPInputSchema{
			Type: "object",
			Properties: map[string]types.MCPProperty{
				"target":     {Type: "string", Description: "target"},
				"content":    {Type: "string", Description: "content"},
				"file_paths": {Type: "array", Description: "files", Items: &types.MCPProperty{Type: "string"}},
			},
			Required: []string{"content"},
		}},
		{Name: "message_history_read", Description: "Read history", InputSchema: types.MCPInputSchema{
			Type:       "object",
			Properties: map[string]types.MCPProperty{"channel": {Type: "string"}},
			Required:   []string{"channel"},
		}},
		{Name: "message_attachment_view", Description: "View file", InputSchema: types.MCPInputSchema{
			Type:       "object",
			Properties: map[string]types.MCPProperty{"attachment_id": {Type: "string"}},
			Required:   []string{"attachment_id"},
		}},
	}

	router, at := setupMCPRouter(t, &testMCPHandler{tools: allTools})
	code, raw := doMCPPost(t, router, at, map[string]interface{}{"jsonrpc": "2.0", "method": "tools/list", "id": 100})
	if code != 200 {
		t.Fatalf("expected 200, got %d: %s", code, raw)
	}

	// Parse into a generic structure so we can inspect every field.
	var resp struct {
		Result struct {
			Tools []json.RawMessage `json:"tools"`
		} `json:"result"`
	}
	if err := json.Unmarshal(raw, &resp); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}
	if len(resp.Result.Tools) != len(allTools) {
		t.Fatalf("expected %d tools, got %d", len(allTools), len(resp.Result.Tools))
	}

	// Validate each tool's JSON against Zod-like constraints.
	for i, toolRaw := range resp.Result.Tools {
		var tool map[string]json.RawMessage
		if err := json.Unmarshal(toolRaw, &tool); err != nil {
			t.Fatalf("tool[%d] unmarshal: %v", i, err)
		}

		// Must have name.
		var name string
		json.Unmarshal(tool["name"], &name)
		if name == "" {
			t.Errorf("tool[%d] missing name", i)
			continue
		}

		// inputSchema must exist.
		schemaRaw, ok := tool["inputSchema"]
		if !ok {
			t.Errorf("tool %q missing inputSchema", name)
			continue
		}

		validateSchema(t, name, "inputSchema", schemaRaw)
	}
	t.Logf("validated %d tools successfully", len(resp.Result.Tools))
}

// validateSchema recursively validates a JSON schema object:
//   - "type" must be present and be a string
//   - "properties" when present must be a JSON object (not null)
//   - "required" when present must be a JSON array
//   - nested property schemas are validated recursively
//   - "items" (for arrays) is validated recursively
func validateSchema(t *testing.T, toolName, path string, raw json.RawMessage) {
	t.Helper()

	var schema map[string]json.RawMessage
	if err := json.Unmarshal(raw, &schema); err != nil {
		t.Errorf("%s.%s: not a JSON object: %v", toolName, path, err)
		return
	}

	// "type" must be a string.
	typeRaw, hasType := schema["type"]
	if !hasType {
		t.Errorf("%s.%s: missing \"type\" field", toolName, path)
		return
	}
	var typeStr string
	if err := json.Unmarshal(typeRaw, &typeStr); err != nil {
		t.Errorf("%s.%s.type: not a string: %v", toolName, path, err)
		return
	}

	// If "properties" is present, it must be a JSON object (never null).
	if propsRaw, hasProps := schema["properties"]; hasProps {
		trimmed := strings.TrimSpace(string(propsRaw))
		if trimmed == "null" {
			t.Errorf("%s.%s.properties: is null (must be {} or omitted)", toolName, path)
		} else if !strings.HasPrefix(trimmed, "{") {
			t.Errorf("%s.%s.properties: not a JSON object: %s", toolName, path, trimmed)
		} else {
			// Validate each nested property schema.
			var props map[string]json.RawMessage
			if err := json.Unmarshal(propsRaw, &props); err == nil {
				for propName, propRaw := range props {
					validateSchema(t, toolName, path+".properties."+propName, propRaw)
				}
			}
		}
	}

	// If "required" is present, it must be a JSON array.
	if reqRaw, hasReq := schema["required"]; hasReq {
		trimmed := strings.TrimSpace(string(reqRaw))
		if trimmed == "null" {
			t.Errorf("%s.%s.required: is null (must be [] or omitted)", toolName, path)
		} else if !strings.HasPrefix(trimmed, "[") {
			t.Errorf("%s.%s.required: not a JSON array: %s", toolName, path, trimmed)
		}
	}

	// If "items" is present (array type), validate it recursively.
	if itemsRaw, hasItems := schema["items"]; hasItems {
		validateSchema(t, toolName, path+".items", itemsRaw)
	}
}

func TestMCP_UnknownTool_ReturnsError(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{
		tools: []types.MCPToolDef{},
		callFn: func(_ context.Context, _, name string, _ json.RawMessage) *types.MCPToolResult {
			return types.NewToolError("unknown tool: " + name)
		},
	})
	code, raw := doMCPPost(t, router, at, map[string]interface{}{
		"jsonrpc": "2.0", "method": "tools/call", "id": 10,
		"params": map[string]interface{}{"name": "nonexistent_tool", "arguments": map[string]interface{}{}},
	})
	if code != 200 {
		t.Fatalf("expected 200, got %d: %s", code, raw)
	}
	var resp mcpRPCResponse
	json.Unmarshal(raw, &resp)
	rm, _ := resp.Result.(map[string]interface{})
	isError, _ := rm["isError"].(bool)
	if !isError {
		t.Error("expected isError=true for unknown tool")
	}
}

func TestMCP_ServerNotFound(t *testing.T) {
	cfgDir := t.TempDir()
	failpoint.Enable("agentteam-config-dir", fmt.Sprintf(`return("%s")`, cfgDir))
	t.Cleanup(func() { failpoint.Disable("agentteam-config-dir") })
	_ = token.Default()

	gin.SetMode(gin.TestMode)
	r := gin.New()
	r.Use(gin.Recovery())
	group := r.Group("/mcp/server/:serverId/:agentId", mcp.AccessControl())
	mcp.RegisterHandlerRoutes(group, func(serverID string) (types.MCPHandler, bool) { return nil, false })

	at, _ := token.Default().GenerateAgentTokenForServer("agent-1", "unknown-server")
	req, _ := http.NewRequest("POST", "/mcp/server/unknown-server/agent-1",
		strings.NewReader(`{"jsonrpc":"2.0","method":"ping","id":1}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	if w.Code != 404 {
		t.Fatalf("expected 404, got %d", w.Code)
	}
}

// ---------------------------------------------------------------------------
// Tests: DELETE, session header, SSE info, keepalive
// ---------------------------------------------------------------------------

func TestMCP_DeleteDisconnect(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	req, _ := http.NewRequest("DELETE", "/mcp/server/test-server/agent-1", nil)
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 204 {
		t.Fatalf("expected 204, got %d", w.Code)
	}
}

func TestMCP_SessionIDHeader(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	req, _ := http.NewRequest("POST", "/mcp/server/test-server/agent-1",
		strings.NewReader(`{"jsonrpc":"2.0","method":"ping","id":1}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	sid := w.Header().Get("Mcp-Session-Id")
	if !strings.HasPrefix(sid, "agentteam-") {
		t.Errorf("Mcp-Session-Id=%s, want prefix agentteam-", sid)
	}
}

func TestMCP_SSE_JSONInfoWithoutAccept(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	req, _ := http.NewRequest("GET", "/mcp/server/test-server/agent-1", nil)
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Authorization", "Bearer "+at)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)
	if w.Code != 200 {
		t.Fatalf("expected 200, got %d", w.Code)
	}
	var info map[string]interface{}
	json.Unmarshal(w.Body.Bytes(), &info)
	if info["name"] != "agentteam" {
		t.Errorf("name=%v, want agentteam", info["name"])
	}
}

func TestMCP_SSEKeepalive(t *testing.T) {
	router, at := setupMCPRouter(t, &testMCPHandler{tools: defaultTestTools()})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	req, _ := http.NewRequest("GET", "/mcp/server/test-server/agent-1", nil)
	req = req.WithContext(ctx)
	req.Header.Set("Accept", "text/event-stream")
	req.Header.Set("Authorization", "Bearer "+at)
	w := newStreamNotifyRecorder()
	done := make(chan struct{})
	go func() {
		router.ServeHTTP(w, req)
		close(done)
	}()
	select {
	case <-w.ready:
	case <-time.After(2 * time.Second):
		t.Fatal("SSE handler did not write initial ready event")
	}
	cancel()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("SSE handler did not exit after cancel")
	}
}

type streamNotifyRecorder struct {
	*httptest.ResponseRecorder
	once  sync.Once
	ready chan struct{}
}

func newStreamNotifyRecorder() *streamNotifyRecorder {
	return &streamNotifyRecorder{
		ResponseRecorder: httptest.NewRecorder(),
		ready:            make(chan struct{}),
	}
}

func (r *streamNotifyRecorder) Write(p []byte) (int, error) {
	n, err := r.ResponseRecorder.Write(p)
	if bytes.Contains(p, []byte(": ready\n\n")) {
		r.notifyReady()
	}
	return n, err
}

func (r *streamNotifyRecorder) WriteString(s string) (int, error) {
	n, err := r.ResponseRecorder.WriteString(s)
	if strings.Contains(s, ": ready\n\n") {
		r.notifyReady()
	}
	return n, err
}

func (r *streamNotifyRecorder) notifyReady() {
	r.once.Do(func() {
		close(r.ready)
	})
}