codespace/internal/acp/client_test.go

package acp

import (
	"bytes"
	"encoding/json"
	"log/slog"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"codespace/internal/fs"
)

// fakeSubscription implements process.Subscription for tests.
type fakeSubscription struct {
	ch   chan []byte
	once sync.Once
}

func newFakeSubscription() *fakeSubscription {
	return &fakeSubscription{ch: make(chan []byte, 16)}
}

func (s *fakeSubscription) Output() <-chan []byte {
	return s.ch
}

func (s *fakeSubscription) send(b []byte) {
	s.ch <- b
}

func (s *fakeSubscription) Close() error {
	s.once.Do(func() { close(s.ch) })
	return nil
}

// recordingWriteCloser captures every byte written to it.
type recordingWriteCloser struct {
	mu     sync.Mutex
	buf    bytes.Buffer
	closed bool
}

func (w *recordingWriteCloser) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.buf.Write(p)
}

func (w *recordingWriteCloser) Close() error {
	w.closed = true
	return nil
}

func (w *recordingWriteCloser) String() string {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.buf.String()
}

// stubFS implements fs.FileSystem for read/write tests.
type stubFS struct {
	mu    sync.Mutex
	files map[string]string
}

func newStubFS() *stubFS {
	return &stubFS{files: make(map[string]string)}
}

func (s *stubFS) Read(path string) ([]byte, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if data, ok := s.files[path]; ok {
		return []byte(data), nil
	}
	return nil, os.ErrNotExist
}

func (s *stubFS) Write(path string, data []byte) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.files[path] = string(data)
	return nil
}

func (s *stubFS) List(path string) ([]fs.FileInfo, error) { return nil, nil }
func (s *stubFS) Mkdir(path string) error                 { return nil }
func (s *stubFS) Remove(path string) error                { return nil }
func (s *stubFS) Rename(oldPath, newPath string) error    { return nil }
func (s *stubFS) Stat(path string) (*fs.FileInfo, error)  { return nil, nil }

func TestNDJSONFramingHandlesSplitLines(t *testing.T) {
	sub := newFakeSubscription()
	stdin := &recordingWriteCloser{}
	tr := newTransport(stdin, sub, slog.Default())
	tr.start()
	defer tr.Close()

	got := make(chan struct {
		method  string
		content string
	}, 2)
	tr.onNotification = func(method string, params json.RawMessage) {
		if method != "session/update" {
			return
		}
		var up SessionUpdateParams
		if err := json.Unmarshal(params, &up); err != nil {
			t.Errorf("unmarshal notification: %v", err)
			return
		}
		got <- struct {
			method  string
			content string
		}{method, up.Update.Content}
	}

	line1 := `{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","content":"hello"}}}`
	line2 := `{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","content":" world"}}}`

	// Send one complete line plus a partial second line.
	sub.send([]byte(line1 + "\n" + line2[:len(line2)-2]))

	select {
	case n := <-got:
		if n.content != "hello" {
			t.Fatalf("first content = %q, want hello", n.content)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for first notification")
	}

	// Complete the second line.
	sub.send([]byte(line2[len(line2)-2:] + "\n"))

	select {
	case n := <-got:
		if n.content != " world" {
			t.Fatalf("second content = %q, want ' world'", n.content)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for second notification")
	}
}

func TestRequestResponseCorrelation(t *testing.T) {
	sub := newFakeSubscription()
	stdin := &recordingWriteCloser{}
	tr := newTransport(stdin, sub, slog.Default())
	tr.start()
	defer tr.Close()

	ch := make(chan *jsonRPCMessage, 1)
	tr.pendingMu.Lock()
	tr.pending[1] = ch
	tr.pendingMu.Unlock()

	sub.send([]byte(`{"jsonrpc":"2.0","id":1,"result":{"sessionId":"sess-1"}}` + "\n"))

	select {
	case resp := <-ch:
		if resp.ID == nil || *resp.ID != 1 {
			t.Fatalf("id = %v, want 1", resp.ID)
		}
		var res SessionNewResult
		if err := json.Unmarshal(resp.Result, &res); err != nil {
			t.Fatal(err)
		}
		if res.SessionID != "sess-1" {
			t.Fatalf("sessionId = %q, want sess-1", res.SessionID)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for response")
	}
}

func TestNotificationDispatch(t *testing.T) {
	sub := newFakeSubscription()
	stdin := &recordingWriteCloser{}
	tr := newTransport(stdin, sub, slog.Default())
	tr.start()
	defer tr.Close()

	got := make(chan string, 1)
	tr.onNotification = func(method string, params json.RawMessage) {
		got <- method
	}

	sub.send([]byte(`{"jsonrpc":"2.0","method":"ping"}` + "\n"))

	select {
	case m := <-got:
		if m != "ping" {
			t.Fatalf("method = %q, want ping", m)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timeout waiting for notification")
	}
}

func TestAgentRequestTerminalReturnsMethodNotFound(t *testing.T) {
	c := &Client{
		workspaceID: "ws",
		fs:          newStubFS(),
		history:     NewHistory(),
		lg:          slog.Default(),
	}
	stdin := &recordingWriteCloser{}
	tr := newTransport(stdin, newFakeSubscription(), slog.Default())
	c.transport = tr
	c.ready = true
	c.sessionID = "s1"

	c.handleRequest("terminal/create", json.RawMessage(`{}`), 7)

	line := strings.TrimSpace(stdin.String())
	if line == "" {
		t.Fatal("expected a response to be written")
	}
	var resp jsonRPCMessage
	if err := json.Unmarshal([]byte(line), &resp); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}
	if resp.ID == nil || *resp.ID != 7 {
		t.Fatalf("id = %v, want 7", resp.ID)
	}
	if resp.Error == nil {
		t.Fatal("expected error response")
	}
	if resp.Error.Code != -32601 {
		t.Fatalf("error code = %d, want -32601", resp.Error.Code)
	}
}

func TestAgentRequestReadFile(t *testing.T) {
	stub := newStubFS()
	stub.files["main.go"] = "package main\n"

	c := &Client{
		workspaceID: "ws",
		fs:          stub,
		history:     NewHistory(),
		lg:          slog.Default(),
	}
	stdin := &recordingWriteCloser{}
	tr := newTransport(stdin, newFakeSubscription(), slog.Default())
	c.transport = tr
	c.ready = true
	c.sessionID = "s1"

	params, _ := json.Marshal(FSReadParams{Path: "main.go", SessionID: "s1"})
	c.handleRequest("fs/read_text_file", params, 3)

	line := strings.TrimSpace(stdin.String())
	if line == "" {
		t.Fatal("expected a response to be written")
	}
	var resp jsonRPCMessage
	if err := json.Unmarshal([]byte(line), &resp); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}
	if resp.ID == nil || *resp.ID != 3 {
		t.Fatalf("id = %v, want 3", resp.ID)
	}
	if resp.Error != nil {
		t.Fatalf("unexpected error: %v", resp.Error)
	}
	var res FSReadResult
	if err := json.Unmarshal(resp.Result, &res); err != nil {
		t.Fatal(err)
	}
	if res.Content != "package main\n" {
		t.Fatalf("content = %q, want 'package main\\n'", res.Content)
	}
}