codespace/internal/acp/client_test.go

package acp

import (
	"bytes"
	"context"
	"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.Text}
	}

	line1 := `{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"hello"}}}}`
	line2 := `{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":" 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)
	}
}

func TestHistoryAppendOrCreateGroupsByMessageID(t *testing.T) {
	h := NewHistory()

	// 1. First chunk for a new messageId creates an entry.
	time.Sleep(10 * time.Millisecond)
	beforeFirst := time.Now()
	time.Sleep(10 * time.Millisecond)
	firstTime := h.AppendOrCreate("agent", "msg-1", "hello")
	time.Sleep(10 * time.Millisecond)
	afterFirst := time.Now()

	if firstTime.Before(beforeFirst) || firstTime.After(afterFirst) {
		t.Fatalf("first chunk timestamp %v out of range [%v, %v]", firstTime, beforeFirst, afterFirst)
	}
	msgs := h.List()
	if len(msgs) != 1 {
		t.Fatalf("after first chunk: len = %d, want 1", len(msgs))
	}
	if msgs[0].Role != "agent" || msgs[0].Text != "hello" || msgs[0].MessageID != "msg-1" {
		t.Fatalf("first chunk message = %+v, want role=agent text=hello messageId=msg-1", msgs[0])
	}

	// 2. Second chunk for the same messageId appends and preserves timestamp.
	time.Sleep(10 * time.Millisecond)
	secondTime := h.AppendOrCreate("agent", "msg-1", " world")
	if !secondTime.Equal(firstTime) {
		t.Fatalf("second chunk timestamp changed: got %v, want %v", secondTime, firstTime)
	}
	msgs = h.List()
	if len(msgs) != 1 {
		t.Fatalf("after second chunk: len = %d, want 1", len(msgs))
	}
	if msgs[0].Text != "hello world" {
		t.Fatalf("after second chunk: text = %q, want \"hello world\"", msgs[0].Text)
	}

	// 3. A different messageId creates a new entry.
	h.AppendOrCreate("agent", "msg-2", "second")
	msgs = h.List()
	if len(msgs) != 2 {
		t.Fatalf("after different id: len = %d, want 2", len(msgs))
	}
	if msgs[1].Text != "second" || msgs[1].MessageID != "msg-2" {
		t.Fatalf("second message = %+v, want text=second messageId=msg-2", msgs[1])
	}

	// 4. Empty messageId always appends a new entry and never collides.
	h.AppendOrCreate("user", "", "a")
	h.AppendOrCreate("user", "", "b")
	msgs = h.List()
	if len(msgs) != 4 {
		t.Fatalf("after two empty-id chunks: len = %d, want 4", len(msgs))
	}
	if msgs[2].Text != "a" || msgs[2].MessageID != "" {
		t.Fatalf("empty-id message[2] = %+v, want text=a messageId=\"\"", msgs[2])
	}
	if msgs[3].Text != "b" || msgs[3].MessageID != "" {
		t.Fatalf("empty-id message[3] = %+v, want text=b messageId=\"\"", msgs[3])
	}

	// 5. Clear resets the index map so the same id can be reused.
	h.Clear()
	if h.Len() != 0 {
		t.Fatalf("after Clear: len = %d, want 0", h.Len())
	}
	h.AppendOrCreate("agent", "msg-1", "after clear")
	msgs = h.List()
	if len(msgs) != 1 {
		t.Fatalf("after clear + chunk: len = %d, want 1", len(msgs))
	}
	if msgs[0].Text != "after clear" || msgs[0].MessageID != "msg-1" {
		t.Fatalf("after clear message = %+v, want text=\"after clear\" messageId=msg-1", msgs[0])
	}
}

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

	out := make(chan StreamEvent, 16)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	go func() {
		for strings.TrimSpace(stdin.String()) == "" {
			time.Sleep(5 * time.Millisecond)
		}
		chunk1 := `{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","messageId":"m1","content":{"type":"text","text":"hello"}}}}`
		chunk2 := `{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s1","update":{"sessionUpdate":"agent_message_chunk","messageId":"m1","content":{"type":"text","text":" world"}}}}`
		response := `{"jsonrpc":"2.0","id":1,"result":{"stopReason":"end_turn"}}`
		sub.send([]byte(chunk1 + "\n"))
		sub.send([]byte(chunk2 + "\n"))
		sub.send([]byte(response + "\n"))
	}()

	if err := c.Stream(ctx, "hi", out); err != nil {
		t.Fatalf("Stream error: %v", err)
	}
	close(out)

	var chunks int
	var complete *StreamEvent
	for ev := range out {
		switch ev.Type {
		case "chunk":
			chunks++
			if ev.MessageID != "m1" {
				t.Fatalf("chunk messageId = %q, want m1", ev.MessageID)
			}
		case "complete":
			if complete != nil {
				t.Fatal("received more than one complete event")
			}
			cp := ev
			complete = &cp
		case "error":
			t.Fatalf("unexpected error event: %q", ev.Error)
		default:
			t.Fatalf("unexpected event type %q", ev.Type)
		}
	}

	if chunks != 2 {
		t.Fatalf("chunks = %d, want 2", chunks)
	}
	if complete == nil {
		t.Fatal("expected complete event")
	}
	if complete.StopReason != "end_turn" {
		t.Fatalf("stopReason = %q, want end_turn", complete.StopReason)
	}
}