kandev/apps/backend/internal/agent/runtime/agentctl/agent_test.go

package client

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

	"github.com/gorilla/websocket"
	"github.com/kandev/kandev/internal/common/logger"
	ws "github.com/kandev/kandev/pkg/websocket"
)

func newTestLogger() *logger.Logger {
	log, _ := logger.NewLogger(logger.LoggingConfig{
		Level:  "error",
		Format: "json",
	})
	return log
}

func TestConfigureAgent_SerializesStructuredArgsAndOmitsAbsentFallback(t *testing.T) {
	var bodies [][]byte
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body, err := io.ReadAll(r.Body)
		if err != nil {
			t.Fatalf("read configure request: %v", err)
		}
		bodies = append(bodies, body)
		_, _ = w.Write([]byte(`{"success":true}`))
	}))
	defer server.Close()

	client := &Client{baseURL: server.URL, httpClient: server.Client()}
	args := []string{"runner", "two words", "", `C:\tools\agent.exe`}
	continueArgs := []string{"runner", "continue", "thread id"}
	if err := client.ConfigureAgent(context.Background(), "runner two words  C:\\tools\\agent.exe", args, nil, "", "runner continue thread id", continueArgs); err != nil {
		t.Fatalf("configure structured args: %v", err)
	}
	if err := client.ConfigureAgent(context.Background(), "legacy --flag", nil, nil, "", "", nil); err != nil {
		t.Fatalf("configure legacy fallback: %v", err)
	}

	var structured map[string]json.RawMessage
	if err := json.Unmarshal(bodies[0], &structured); err != nil {
		t.Fatalf("decode structured request: %v", err)
	}
	for _, field := range []string{"agent_args", "continue_args"} {
		if _, ok := structured[field]; !ok {
			t.Fatalf("structured request omitted %s: %s", field, bodies[0])
		}
	}
	var legacy map[string]json.RawMessage
	if err := json.Unmarshal(bodies[1], &legacy); err != nil {
		t.Fatalf("decode legacy request: %v", err)
	}
	if _, ok := legacy["agent_args"]; ok {
		t.Fatalf("legacy fallback unexpectedly sent agent_args: %s", bodies[1])
	}
}

// wsTestServer creates a test WebSocket server that echoes back response messages.
// The handler receives each request and should return a response message.
type wsTestServer struct {
	server  *httptest.Server
	handler func(msg ws.Message) *ws.Message
}

func newWSTestServer(t *testing.T, handler func(msg ws.Message) *ws.Message) *wsTestServer {
	t.Helper()
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}
	ts := &wsTestServer{handler: handler}

	ts.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			t.Logf("upgrade error: %v", err)
			return
		}
		defer func() { _ = conn.Close() }()

		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				return
			}

			var msg ws.Message
			if err := json.Unmarshal(message, &msg); err != nil {
				continue
			}

			if msg.Type == ws.MessageTypeRequest && ts.handler != nil {
				resp := ts.handler(msg)
				if resp != nil {
					data, _ := json.Marshal(resp)
					if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
						return
					}
				}
			}
		}
	}))

	return ts
}

func (ts *wsTestServer) Close() {
	ts.server.Close()
}

// newTestClientWithStream creates a Client connected to a test WebSocket server.
// The handler processes incoming requests and returns responses.
func newTestClientWithStream(t *testing.T, handler func(msg ws.Message) *ws.Message) (*Client, *wsTestServer) {
	t.Helper()

	ts := newWSTestServer(t, handler)

	// Parse the test server URL to get host and port
	// httptest URL is like http://127.0.0.1:PORT
	url := ts.server.URL
	log := newTestLogger()

	// Create client with the test server URL
	c := &Client{
		baseURL:               url,
		httpClient:            &http.Client{Timeout: 5 * time.Second},
		longRunningHTTPClient: &http.Client{Timeout: 5 * time.Minute},
		logger:                log,
		pendingRequests:       make(map[string]chan *ws.Message),
	}

	// Connect the stream
	ctx := context.Background()
	err := c.StreamUpdates(ctx, func(event AgentEvent) {
		// no-op handler for agent events
	}, nil, nil)
	if err != nil {
		ts.Close()
		t.Fatalf("failed to connect stream: %v", err)
	}

	// Give the goroutine time to start
	time.Sleep(50 * time.Millisecond)

	return c, ts
}

// --- Tests for infrastructure methods ---

func TestSendStreamRequest_NotConnected(t *testing.T) {
	log := newTestLogger()
	c := &Client{
		baseURL:         "http://localhost:0",
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	ctx := context.Background()
	_, err := c.sendStreamRequest(ctx, "test.action", nil)
	if err == nil {
		t.Fatal("expected error when stream not connected")
	}
	if !errors.Is(err, ErrAgentStreamNotConnected) {
		t.Fatalf("expected ErrAgentStreamNotConnected, got: %v", err)
	}
	if !strings.Contains(err.Error(), "agent stream not connected") {
		t.Fatalf("expected 'agent stream not connected' error, got: %v", err)
	}
}

func TestSendStreamRequest_ContextCancelled(t *testing.T) {
	// Server that never responds
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		return nil // don't respond
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
	defer cancel()

	_, err := c.sendStreamRequest(ctx, "test.action", nil)
	if err == nil {
		t.Fatal("expected error on timeout")
	}
	if !strings.Contains(err.Error(), "context deadline exceeded") {
		t.Fatalf("expected deadline exceeded error, got: %v", err)
	}
}

func TestResolvePendingRequest_NoMatch(t *testing.T) {
	log := newTestLogger()
	c := &Client{
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	msg := &ws.Message{ID: "nonexistent", Type: ws.MessageTypeResponse}
	if c.resolvePendingRequest(msg) {
		t.Fatal("expected false for unmatched request")
	}
}

func TestResolvePendingRequest_EmptyID(t *testing.T) {
	log := newTestLogger()
	c := &Client{
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	msg := &ws.Message{ID: "", Type: ws.MessageTypeResponse}
	if c.resolvePendingRequest(msg) {
		t.Fatal("expected false for empty ID")
	}
}

func TestResolvePendingRequest_Match(t *testing.T) {
	log := newTestLogger()
	c := &Client{
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	ch := make(chan *ws.Message, 1)
	c.pendingMu.Lock()
	c.pendingRequests["test-id"] = ch
	c.pendingMu.Unlock()

	msg := &ws.Message{ID: "test-id", Type: ws.MessageTypeResponse, Action: "test.action"}
	if !c.resolvePendingRequest(msg) {
		t.Fatal("expected true for matched request")
	}

	select {
	case resp := <-ch:
		if resp.ID != "test-id" {
			t.Fatalf("expected response ID 'test-id', got %q", resp.ID)
		}
	default:
		t.Fatal("expected response in channel")
	}
}

func TestCleanupPendingRequests(t *testing.T) {
	log := newTestLogger()
	c := &Client{
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	ch1 := make(chan *ws.Message, 1)
	ch2 := make(chan *ws.Message, 1)
	c.pendingMu.Lock()
	c.pendingRequests["req-1"] = ch1
	c.pendingRequests["req-2"] = ch2
	c.pendingMu.Unlock()

	c.cleanupPendingRequests()

	c.pendingMu.Lock()
	count := len(c.pendingRequests)
	c.pendingMu.Unlock()

	if count != 0 {
		t.Fatalf("expected 0 pending requests after cleanup, got %d", count)
	}

	// Channels should be closed
	_, ok := <-ch1
	if ok {
		t.Fatal("expected ch1 to be closed")
	}
	_, ok = <-ch2
	if ok {
		t.Fatal("expected ch2 to be closed")
	}
}

func TestCleanupPendingRequests_ConcurrentSafe(t *testing.T) {
	log := newTestLogger()
	c := &Client{
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	// Add many pending requests
	for i := 0; i < 50; i++ {
		ch := make(chan *ws.Message, 1)
		c.pendingMu.Lock()
		c.pendingRequests[strings.Repeat("x", i+1)] = ch
		c.pendingMu.Unlock()
	}

	// Cleanup concurrently
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		c.cleanupPendingRequests()
	}()
	wg.Wait()

	c.pendingMu.Lock()
	count := len(c.pendingRequests)
	c.pendingMu.Unlock()

	if count != 0 {
		t.Fatalf("expected 0 pending requests, got %d", count)
	}
}

// --- Tests for the 7 migrated methods ---

func TestInitialize_Success(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		if msg.Action != "agent.initialize" {
			t.Errorf("expected action 'agent.initialize', got %q", msg.Action)
		}
		resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
			"success": true,
			"agent_info": map[string]string{
				"name":    "test-agent",
				"version": "1.0.0",
			},
		})
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	info, err := c.Initialize(ctx, "kandev", "1.0.0")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if info == nil {
		t.Fatal("expected agent info, got nil")
	} else {
		if info.Name != "test-agent" {
			t.Errorf("expected name 'test-agent', got %q", info.Name)
		}
		if info.Version != "1.0.0" {
			t.Errorf("expected version '1.0.0', got %q", info.Version)
		}
	}
}

func TestInitialize_ServerError(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		resp, _ := ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, "agent not running", nil)
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	_, err := c.Initialize(ctx, "kandev", "1.0.0")
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "agent not running") {
		t.Fatalf("expected 'agent not running' error, got: %v", err)
	}
}

func TestInitialize_SuccessFalse(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
			"success": false,
			"error":   "init failed internally",
		})
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	_, err := c.Initialize(ctx, "kandev", "1.0.0")
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "init failed internally") {
		t.Fatalf("expected 'init failed internally' error, got: %v", err)
	}
}

func TestNewSession_Success(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		if msg.Action != "agent.session.new" {
			t.Errorf("expected action 'agent.session.new', got %q", msg.Action)
		}
		resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
			"success":    true,
			"session_id": "sess-123",
		})
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	sessionID, err := c.NewSession(ctx, "/workspace", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if sessionID != "sess-123" {
		t.Errorf("expected session ID 'sess-123', got %q", sessionID)
	}
}

func TestNewSession_Error(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		resp, _ := ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, "agent not running", nil)
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	_, err := c.NewSession(ctx, "/workspace", nil)
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "agent not running") {
		t.Fatalf("expected 'agent not running' error, got: %v", err)
	}
}

func TestLoadSession_Success(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		if msg.Action != "agent.session.load" {
			t.Errorf("expected action 'agent.session.load', got %q", msg.Action)
		}
		resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
			"success":    true,
			"session_id": "sess-456",
		})
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err := c.LoadSession(ctx, "sess-456", nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestLoadSession_Error(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		resp, _ := ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, "Method not found", nil)
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err := c.LoadSession(ctx, "sess-456", nil)
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "Method not found") {
		t.Fatalf("expected 'Method not found' error, got: %v", err)
	}
}

func TestPrompt_Success(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		if msg.Action != "agent.prompt" {
			t.Errorf("expected action 'agent.prompt', got %q", msg.Action)
		}
		// Verify payload
		var payload struct {
			Text             string `json:"text"`
			PromptGeneration uint64 `json:"prompt_generation"`
		}
		if err := msg.ParsePayload(&payload); err != nil {
			t.Errorf("failed to parse payload: %v", err)
		}
		if payload.Text != "hello agent" {
			t.Errorf("expected text 'hello agent', got %q", payload.Text)
		}
		if payload.PromptGeneration != 42 {
			t.Errorf("expected prompt generation 42, got %d", payload.PromptGeneration)
		}
		resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
			"success": true,
		})
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err := c.Prompt(ctx, "hello agent", nil, 42)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestPrompt_NoActiveSession(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		resp, _ := ws.NewError(msg.ID, msg.Action, ws.ErrorCodeBadRequest, "no active session - call new_session first", nil)
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err := c.Prompt(ctx, "hello agent", nil, 0)
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "no active session") {
		t.Fatalf("expected 'no active session' error, got: %v", err)
	}
}

func TestCancel_Success(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		if msg.Action != "agent.cancel" {
			t.Errorf("expected action 'agent.cancel', got %q", msg.Action)
		}
		resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
			"success": true,
		})
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err := c.Cancel(ctx)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestCancel_Error(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		resp, _ := ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, "agent not running", nil)
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err := c.Cancel(ctx)
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "agent not running") {
		t.Fatalf("expected 'agent not running' error, got: %v", err)
	}
}

func TestGetAgentStderr_Success(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		if msg.Action != "agent.stderr" {
			t.Errorf("expected action 'agent.stderr', got %q", msg.Action)
		}
		resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
			"lines": []string{"error line 1", "error line 2"},
		})
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	lines, err := c.GetAgentStderr(ctx)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if len(lines) != 2 {
		t.Fatalf("expected 2 lines, got %d", len(lines))
	}
	if lines[0] != "error line 1" {
		t.Errorf("expected 'error line 1', got %q", lines[0])
	}
}

func TestGetAgentStderr_NotConnected(t *testing.T) {
	log := newTestLogger()
	c := &Client{
		baseURL:         "http://localhost:0",
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	ctx := context.Background()
	_, err := c.GetAgentStderr(ctx)
	if err == nil {
		t.Fatal("expected error when stream not connected")
	}
	if !strings.Contains(err.Error(), "agent stream not connected") {
		t.Fatalf("expected 'agent stream not connected' error, got: %v", err)
	}
}

func TestRespondToPermission_Success(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		if msg.Action != "agent.permissions.respond" {
			t.Errorf("expected action 'agent.permissions.respond', got %q", msg.Action)
		}
		// Verify payload
		var payload struct {
			PendingID string `json:"pending_id"`
			OptionID  string `json:"option_id"`
			Cancelled bool   `json:"cancelled"`
		}
		if err := msg.ParsePayload(&payload); err != nil {
			t.Errorf("failed to parse payload: %v", err)
		}
		if payload.PendingID != "perm-1" {
			t.Errorf("expected pending_id 'perm-1', got %q", payload.PendingID)
		}
		if payload.OptionID != "allow" {
			t.Errorf("expected option_id 'allow', got %q", payload.OptionID)
		}
		resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
			"success": true,
		})
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err := c.RespondToPermission(ctx, "perm-1", "allow", false)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestRespondToPermission_NotFound(t *testing.T) {
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		resp, _ := ws.NewError(msg.ID, msg.Action, ws.ErrorCodeNotFound, "permission request not found", nil)
		return resp
	})
	defer ts.Close()
	defer c.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	err := c.RespondToPermission(ctx, "nonexistent", "allow", false)
	if err == nil {
		t.Fatal("expected error")
	}
	if !strings.Contains(err.Error(), "permission request not found") {
		t.Fatalf("expected 'permission request not found' error, got: %v", err)
	}
}

// --- Tests for StreamUpdates response routing ---

func TestStreamUpdates_RoutesResponsesToPending(t *testing.T) {
	// This test verifies that StreamUpdates correctly routes response messages
	// to pending requests instead of treating them as agent events.
	var receivedEvents []AgentEvent
	var mu sync.Mutex

	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer func() { _ = conn.Close() }()

		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				return
			}
			var msg ws.Message
			if err := json.Unmarshal(message, &msg); err != nil {
				continue
			}
			if msg.Type == ws.MessageTypeRequest {
				// Send response back
				resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
					"success": true,
				})
				data, _ := json.Marshal(resp)
				_ = conn.WriteMessage(websocket.TextMessage, data)

				// Also send an agent event
				event := AgentEvent{Type: "message_chunk", Text: "hello"}
				eventData, _ := json.Marshal(event)
				_ = conn.WriteMessage(websocket.TextMessage, eventData)
			}
		}
	}))
	defer server.Close()

	log := newTestLogger()
	c := &Client{
		baseURL:         server.URL,
		httpClient:      &http.Client{Timeout: 5 * time.Second},
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	ctx := context.Background()
	err := c.StreamUpdates(ctx, func(event AgentEvent) {
		mu.Lock()
		receivedEvents = append(receivedEvents, event)
		mu.Unlock()
	}, nil, nil)
	if err != nil {
		t.Fatalf("failed to connect: %v", err)
	}

	time.Sleep(50 * time.Millisecond)

	// Send a request and get a response
	reqCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	resp, err := c.sendStreamRequest(reqCtx, "test.action", nil)
	if err != nil {
		t.Fatalf("sendStreamRequest failed: %v", err)
	}
	if resp.Type != ws.MessageTypeResponse {
		t.Errorf("expected response type, got %q", resp.Type)
	}

	// Wait for the agent event to be processed
	time.Sleep(100 * time.Millisecond)

	mu.Lock()
	eventCount := len(receivedEvents)
	mu.Unlock()

	if eventCount != 1 {
		t.Errorf("expected 1 agent event, got %d", eventCount)
	}

	c.Close()
}

func TestStreamUpdates_DisconnectCleansPending(t *testing.T) {
	// This test verifies that pending requests are cleaned up on stream disconnect.
	upgrader := websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		// Close immediately to trigger disconnect
		time.Sleep(100 * time.Millisecond)
		_ = conn.Close()
	}))
	defer server.Close()

	log := newTestLogger()
	c := &Client{
		baseURL:         server.URL,
		httpClient:      &http.Client{Timeout: 5 * time.Second},
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	disconnected := make(chan struct{})
	ctx := context.Background()
	err := c.StreamUpdates(ctx, func(event AgentEvent) {}, nil, func(err error) {
		close(disconnected)
	})
	if err != nil {
		t.Fatalf("failed to connect: %v", err)
	}

	time.Sleep(50 * time.Millisecond)

	// Add a pending request
	ch := make(chan *ws.Message, 1)
	c.pendingMu.Lock()
	c.pendingRequests["pending-test"] = ch
	c.pendingMu.Unlock()

	// Wait for disconnect
	select {
	case <-disconnected:
	case <-time.After(5 * time.Second):
		t.Fatal("timeout waiting for disconnect")
	}

	// The channel should be closed
	_, ok := <-ch
	if ok {
		t.Fatal("expected channel to be closed after disconnect")
	}

	c.pendingMu.Lock()
	count := len(c.pendingRequests)
	c.pendingMu.Unlock()
	if count != 0 {
		t.Fatalf("expected 0 pending requests after disconnect, got %d", count)
	}
}

// TestReadUpdatesStream_BlockedHandlerDoesNotStarveResponse is the regression
// test for the production stream-reader deadlock (task 544afdae). An agent
// event handler (handleAgentReady) blocked on the per-session cancelInFlight
// guard, and that guard's holder — Service.CancelAgent — was itself blocked in
// sendStreamRequest waiting for the agent.cancel response frame. Because the
// read loop used to run handlers inline, it never looped back to deliver that
// response, so both sides waited forever.
//
// This pins the fix: the read loop offloads event handling to an ordered worker
// goroutine, so even while a handler is blocked mid-event, a concurrent
// sendStreamRequest still receives its response frame.
func TestReadUpdatesStream_BlockedHandlerDoesNotStarveResponse(t *testing.T) {
	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}

	// The server emits an agent event first, then answers the request that the
	// test fires while that event's handler is still blocked.
	emitEvent := make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer func() { _ = conn.Close() }()

		// Push an agent event whose handler will block on the client side.
		<-emitEvent
		eventData, _ := json.Marshal(AgentEvent{Type: "complete"})
		_ = conn.WriteMessage(websocket.TextMessage, eventData)

		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				return
			}
			var msg ws.Message
			if err := json.Unmarshal(message, &msg); err != nil {
				continue
			}
			if msg.Type == ws.MessageTypeRequest {
				resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{"success": true})
				data, _ := json.Marshal(resp)
				_ = conn.WriteMessage(websocket.TextMessage, data)
			}
		}
	}))
	defer server.Close()

	handlerEntered := make(chan struct{})
	releaseHandler := make(chan struct{})
	log := newTestLogger()
	c := &Client{
		baseURL:         server.URL,
		httpClient:      &http.Client{Timeout: 5 * time.Second},
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	ctx := context.Background()
	var once sync.Once
	if err := c.StreamUpdates(ctx, func(_ AgentEvent) {
		// Stand in for handleAgentReady blocking on the cancelInFlight guard.
		once.Do(func() { close(handlerEntered) })
		<-releaseHandler
	}, nil, nil); err != nil {
		t.Fatalf("failed to connect: %v", err)
	}
	defer func() {
		close(releaseHandler)
		c.Close()
	}()

	// Trigger the event and wait until its handler is blocked in-flight.
	close(emitEvent)
	select {
	case <-handlerEntered:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for the blocking event handler to be entered")
	}

	// With the handler still blocked, a request must still get its response —
	// the read loop is not wedged behind the handler.
	reqCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	resp, err := c.sendStreamRequest(reqCtx, "agent.cancel", nil)
	if err != nil {
		t.Fatalf("sendStreamRequest was starved by the blocked event handler: %v", err)
	}
	if resp.Type != ws.MessageTypeResponse {
		t.Fatalf("expected response type, got %q", resp.Type)
	}
}

// TestReadUpdatesStream_EventBurstDoesNotStarveResponse pins the fix for the
// bounded-queue regression (cubic review on task 544afdae): even if the first
// event's handler blocks on the cancelInFlight guard and a large burst of
// further events arrives behind it, the read loop must stay free to deliver the
// agent.cancel response frame. A fixed buffered channel would fill and
// backpressure the read loop on send, re-wedging the cancel; the unbounded
// reader-side queue must not.
func TestReadUpdatesStream_EventBurstDoesNotStarveResponse(t *testing.T) {
	const burst = 1000 // comfortably beyond the old 256 buffer

	upgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}
	emitEvents := make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			return
		}
		defer func() { _ = conn.Close() }()

		<-emitEvents
		// Blast a burst of events; the first handler will block on the client
		// side, so all of these queue up behind it.
		for i := 0; i < burst; i++ {
			eventData, _ := json.Marshal(AgentEvent{Type: "message_chunk"})
			if err := conn.WriteMessage(websocket.TextMessage, eventData); err != nil {
				return
			}
		}

		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				return
			}
			var msg ws.Message
			if err := json.Unmarshal(message, &msg); err != nil {
				continue
			}
			if msg.Type == ws.MessageTypeRequest {
				resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{"success": true})
				data, _ := json.Marshal(resp)
				_ = conn.WriteMessage(websocket.TextMessage, data)
			}
		}
	}))
	defer server.Close()

	handlerEntered := make(chan struct{})
	releaseHandler := make(chan struct{})
	log := newTestLogger()
	c := &Client{
		baseURL:         server.URL,
		httpClient:      &http.Client{Timeout: 5 * time.Second},
		logger:          log,
		pendingRequests: make(map[string]chan *ws.Message),
	}

	ctx := context.Background()
	var once sync.Once
	if err := c.StreamUpdates(ctx, func(_ AgentEvent) {
		// Only the first handler blocks; the rest queue up behind it.
		once.Do(func() {
			close(handlerEntered)
			<-releaseHandler
		})
	}, nil, nil); err != nil {
		t.Fatalf("failed to connect: %v", err)
	}
	defer func() {
		close(releaseHandler)
		c.Close()
	}()

	close(emitEvents)
	select {
	case <-handlerEntered:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for the blocking event handler to be entered")
	}

	// With the first handler blocked and a full burst queued behind it, a
	// request must still get its response — the read loop is not wedged behind
	// a full event buffer.
	reqCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	resp, err := c.sendStreamRequest(reqCtx, "agent.cancel", nil)
	if err != nil {
		t.Fatalf("sendStreamRequest was starved by the event burst: %v", err)
	}
	if resp.Type != ws.MessageTypeResponse {
		t.Fatalf("expected response type, got %q", resp.Type)
	}
}

// --- Concurrent request tests ---

func TestConcurrentRequests(t *testing.T) {
	// Multiple concurrent requests should all resolve correctly.
	c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
		// Simulate some processing time
		time.Sleep(10 * time.Millisecond)
		resp, _ := ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{
			"success": true,
			"action":  msg.Action,
		})
		return resp
	})
	defer ts.Close()
	defer c.Close()

	var wg sync.WaitGroup
	errors := make(chan error, 10)

	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()

			_, err := c.sendStreamRequest(ctx, "test.concurrent", nil)
			if err != nil {
				errors <- err
			}
		}()
	}

	wg.Wait()
	close(errors)

	for err := range errors {
		t.Errorf("concurrent request failed: %v", err)
	}
}