c82668c789
Co-authored-by: Vojtěch Biberle <vojtech.biberle@gmail.com> Co-authored-by: nova28 <17953305+nova28@users.noreply.github.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Co-authored-by: Carlos Florencio <carlosmflorencio@gmail.com>
1193 lines
33 KiB
Go
1193 lines
33 KiB
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/agentctl/types/streams"
|
|
"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(nil)
|
|
|
|
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(nil)
|
|
}()
|
|
wg.Wait()
|
|
|
|
c.pendingMu.Lock()
|
|
count := len(c.pendingRequests)
|
|
c.pendingMu.Unlock()
|
|
|
|
if count != 0 {
|
|
t.Fatalf("expected 0 pending requests, got %d", count)
|
|
}
|
|
}
|
|
|
|
func TestCleanupPendingRequests_ScopesToRetiredStream(t *testing.T) {
|
|
c := &Client{
|
|
pendingRequests: make(map[string]chan *ws.Message),
|
|
pendingRequestConns: make(map[string]*websocket.Conn),
|
|
}
|
|
oldConn := &websocket.Conn{}
|
|
newConn := &websocket.Conn{}
|
|
oldCh := make(chan *ws.Message, 1)
|
|
newCh := make(chan *ws.Message, 1)
|
|
c.pendingMu.Lock()
|
|
c.pendingRequests["old"] = oldCh
|
|
c.pendingRequests["new"] = newCh
|
|
c.pendingRequestConns["old"] = oldConn
|
|
c.pendingRequestConns["new"] = newConn
|
|
c.pendingMu.Unlock()
|
|
|
|
c.cleanupPendingRequests(oldConn)
|
|
|
|
if _, ok := <-oldCh; ok {
|
|
t.Fatal("expected retired stream request channel to be closed")
|
|
}
|
|
select {
|
|
case _, ok := <-newCh:
|
|
if !ok {
|
|
t.Fatal("replacement stream request channel was closed")
|
|
}
|
|
t.Fatal("replacement stream request channel received an unexpected response")
|
|
default:
|
|
}
|
|
c.pendingMu.Lock()
|
|
defer c.pendingMu.Unlock()
|
|
if _, ok := c.pendingRequests["new"]; !ok {
|
|
t.Fatal("replacement stream request was removed")
|
|
}
|
|
}
|
|
|
|
// --- 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)
|
|
}
|
|
}
|
|
|
|
func TestListPendingPermissions(t *testing.T) {
|
|
c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
|
|
if msg.Action != "agent.permissions.list" {
|
|
t.Errorf("action = %q, want agent.permissions.list", msg.Action)
|
|
}
|
|
resp, _ := ws.NewResponse(msg.ID, msg.Action, streams.PermissionListResponse{
|
|
Permissions: []streams.PendingAgentPermission{{
|
|
RequestID: "request-1",
|
|
PendingID: "pending-1",
|
|
Status: streams.PermissionStatusPending,
|
|
}},
|
|
Total: 1,
|
|
})
|
|
return resp
|
|
})
|
|
defer ts.Close()
|
|
defer c.Close()
|
|
|
|
permissions, err := c.ListPendingPermissions(t.Context())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(permissions) != 1 || permissions[0].RequestID != "request-1" {
|
|
t.Fatalf("unexpected permissions: %+v", permissions)
|
|
}
|
|
}
|
|
|
|
func TestResolvePermissionPreservesStableErrorCode(t *testing.T) {
|
|
c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
|
|
if msg.Action != "agent.permissions.resolve" {
|
|
t.Errorf("action = %q, want agent.permissions.resolve", msg.Action)
|
|
}
|
|
var request streams.PermissionResolveRequest
|
|
if err := msg.ParsePayload(&request); err != nil {
|
|
t.Error(err)
|
|
}
|
|
if request.RequestID != "request-1" || request.PendingID != "pending-1" || request.OptionID != "unknown" {
|
|
t.Errorf("unexpected request: %+v", request)
|
|
}
|
|
resp, _ := ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, streams.PermissionErrorOptionNotOffered, map[string]any{
|
|
"permission_code": streams.PermissionErrorOptionNotOffered,
|
|
})
|
|
return resp
|
|
})
|
|
defer ts.Close()
|
|
defer c.Close()
|
|
|
|
_, err := c.ResolvePermission(t.Context(), "request-1", "pending-1", "unknown")
|
|
var operationErr *PermissionOperationError
|
|
if !errors.As(err, &operationErr) {
|
|
t.Fatalf("error = %v, want PermissionOperationError", err)
|
|
}
|
|
if operationErr.Code != streams.PermissionErrorOptionNotOffered {
|
|
t.Fatalf("code = %q, want %q", operationErr.Code, streams.PermissionErrorOptionNotOffered)
|
|
}
|
|
}
|
|
|
|
func TestCancelPermissionSendsExactGeneration(t *testing.T) {
|
|
c, ts := newTestClientWithStream(t, func(msg ws.Message) *ws.Message {
|
|
if msg.Action != "agent.permissions.cancel" {
|
|
t.Errorf("action = %q, want agent.permissions.cancel", msg.Action)
|
|
}
|
|
var request streams.PermissionCancelRequest
|
|
if err := msg.ParsePayload(&request); err != nil {
|
|
t.Error(err)
|
|
}
|
|
if request.RequestID != "request-1" || request.PendingID != "pending-1" {
|
|
t.Errorf("unexpected cancellation tuple: %+v", request)
|
|
}
|
|
resp, _ := ws.NewResponse(msg.ID, msg.Action, streams.PermissionCancelResponse{
|
|
RequestID: request.RequestID, PendingID: request.PendingID, Status: "cancelled",
|
|
})
|
|
return resp
|
|
})
|
|
defer ts.Close()
|
|
defer c.Close()
|
|
|
|
result, err := c.CancelPermission(t.Context(), "request-1", "pending-1")
|
|
if err != nil || result.Status != "cancelled" {
|
|
t.Fatalf("result=%+v err=%v", result, 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.pendingRequestConns = make(map[string]*websocket.Conn)
|
|
c.pendingRequestConns["pending-test"] = c.agentStreamConn
|
|
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)
|
|
}
|
|
}
|