agentic-coding-workflow/internal/acp/relay_test.go

// relay_test.go: integration tests for relay.go using io.Pipe pairs to
// simulate the agent subprocess (no real exec). Tests cover:
//   - Server-initiated Call round-trip (int id) + event persistence
//   - Notification fanout to subscribers
//   - Agent-initiated fs/read_text_file dispatched to real FsHandler
//   - Slow subscriber gets disconnected with final envelope
package acp_test

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/yan1h/agent-coding-workflow/internal/acp"
	"github.com/yan1h/agent-coding-workflow/internal/acp/handlers"
)

// fakeEventRepo only implements InsertEvent meaningfully; everything else
// panics. Sufficient for relay tests, which only exercise event persistence.
type fakeEventRepo struct {
	mu     sync.Mutex
	events []*acp.Event
}

func (f *fakeEventRepo) InsertEvent(_ context.Context, e *acp.Event) (*acp.Event, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	cp := *e
	cp.ID = int64(len(f.events) + 1)
	cp.CreatedAt = time.Now()
	f.events = append(f.events, &cp)
	out := cp
	return &out, nil
}

func (f *fakeEventRepo) snapshot() []*acp.Event {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]*acp.Event, len(f.events))
	copy(out, f.events)
	return out
}

// All other Repository methods panic — relay tests must not exercise them.
func (f *fakeEventRepo) CreateAgentKind(context.Context, *acp.AgentKind) (*acp.AgentKind, error) {
	panic("n/a")
}
func (f *fakeEventRepo) GetAgentKindByID(context.Context, uuid.UUID) (*acp.AgentKind, error) {
	panic("n/a")
}
func (f *fakeEventRepo) GetAgentKindByName(context.Context, string) (*acp.AgentKind, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ListAgentKinds(context.Context) ([]*acp.AgentKind, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ListEnabledAgentKinds(context.Context) ([]*acp.AgentKind, error) {
	panic("n/a")
}
func (f *fakeEventRepo) UpdateAgentKind(context.Context, *acp.AgentKind) (*acp.AgentKind, error) {
	panic("n/a")
}
func (f *fakeEventRepo) DeleteAgentKind(context.Context, uuid.UUID) error { panic("n/a") }
func (f *fakeEventRepo) CountAgentKindUsage(context.Context, uuid.UUID) (int64, error) {
	panic("n/a")
}
func (f *fakeEventRepo) InsertSession(context.Context, *acp.Session) (*acp.Session, error) {
	panic("n/a")
}
func (f *fakeEventRepo) GetSessionByID(context.Context, uuid.UUID) (*acp.Session, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ListSessionsByUser(context.Context, uuid.UUID, *uuid.UUID) ([]*acp.Session, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ListAllSessions(context.Context, *uuid.UUID) ([]*acp.Session, error) {
	panic("n/a")
}
func (f *fakeEventRepo) UpdateSessionRunning(context.Context, uuid.UUID, string, int32) error {
	panic("n/a")
}
func (f *fakeEventRepo) UpdateSessionFinished(context.Context, uuid.UUID, acp.SessionStatus, *int32, *string) error {
	panic("n/a")
}
func (f *fakeEventRepo) CountActiveSessions(context.Context) (int64, error) { panic("n/a") }
func (f *fakeEventRepo) CountActiveSessionsByUser(context.Context, uuid.UUID) (int64, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ResetStuckSessionsOnRestart(context.Context) ([]*acp.Session, error) {
	panic("n/a")
}
func (f *fakeEventRepo) AcquireMainWorktree(context.Context, uuid.UUID, uuid.UUID) (bool, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ReleaseMainWorktree(context.Context, uuid.UUID, uuid.UUID) (bool, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ListEventsSince(context.Context, uuid.UUID, int64, int32) ([]*acp.Event, error) {
	panic("n/a")
}
func (f *fakeEventRepo) PurgeEventsBefore(context.Context, time.Time) (int64, error) {
	panic("n/a")
}

func (f *fakeEventRepo) InTx(context.Context, func(context.Context, pgx.Tx) error) error {
	panic("n/a")
}

func (f *fakeEventRepo) WithTx(pgx.Tx) acp.Repository {
	return f
}

// relayTestRig wires up a Relay with two io.Pipe pairs that simulate the agent
// subprocess. Returns:
//   - r:           the relay under test
//   - agentWriter: an io.WriteCloser the test "agent" writes to (relay reads from)
//   - agentReader: an io.ReadCloser the test "agent" reads from (relay writes to)
//   - repo:        the shared event store
//   - sess:        a SessionContext suitable for fs/perm dispatch
type relayTestRig struct {
	r           *acp.Relay
	agentWriter io.WriteCloser
	agentReader io.ReadCloser
	repo        *fakeEventRepo
	sess        handlers.SessionContext
	// keep references so the helper test can close them
	relayStdinR io.ReadCloser
	relayStdinW io.WriteCloser
}

func newRelayTestRig(t *testing.T) *relayTestRig {
	t.Helper()

	// relay → agent: relay encodes to relayStdinW, agent reads from relayStdinR
	relayStdinR, relayStdinW := io.Pipe()
	// agent → relay: agent writes to agentStdoutW, relay decodes from agentStdoutR
	agentStdoutR, agentStdoutW := io.Pipe()

	repo := &fakeEventRepo{}
	fs := handlers.NewFsHandler()
	perm := handlers.NewPermissionHandler(nil)

	cwd := t.TempDir()
	sess := handlers.SessionContext{
		SessionID:   uuid.New(),
		WorkspaceID: uuid.New(),
		UserID:      uuid.New(),
		CwdPath:     cwd,
	}

	r := acp.NewRelay(
		sess.SessionID,
		acp.NewDecoder(agentStdoutR, 0),
		acp.NewEncoder(relayStdinW),
		repo, fs, perm,
		acp.RelayConfig{EventMaxPayload: 65536, EventTruncateField: 4096, WSSendBuffer: 16},
		nil,
	)
	return &relayTestRig{
		r:           r,
		agentWriter: agentStdoutW,
		agentReader: relayStdinR,
		repo:        repo,
		sess:        sess,
		relayStdinR: relayStdinR,
		relayStdinW: relayStdinW,
	}
}

// runRelay starts r.Run in a goroutine and returns a cancel func + done chan.
// The caller must invoke cancel() and wait on doneCh in test cleanup.
func (rig *relayTestRig) runRelay(t *testing.T) (cancel context.CancelFunc, doneCh <-chan struct{}) {
	t.Helper()
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() {
		rig.r.Run(ctx, rig.sess)
		close(done)
	}()
	return cancel, done
}

// shutdown closes the agent-side pipes (causing relay reader to EOF) and waits
// for Run to return. Called via t.Cleanup.
func (rig *relayTestRig) shutdown(t *testing.T, cancel context.CancelFunc, doneCh <-chan struct{}) {
	t.Helper()
	// Close agent → relay so reader sees EOF.
	_ = rig.agentWriter.Close()
	cancel()
	// Drain the relay → agent pipe so writer never blocks.
	_ = rig.agentReader.Close()
	select {
	case <-doneCh:
	case <-time.After(3 * time.Second):
		t.Log("relay.Run did not return within 3s of shutdown")
	}
}

// ---------- Test 1 ----------

func TestRelay_Call_RoundTripIntID(t *testing.T) {
	t.Parallel()
	rig := newRelayTestRig(t)
	cancel, doneCh := rig.runRelay(t)
	t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })

	// Goroutine "agent": decode the request from relay's writer side, emit a
	// matching response.
	agentDec := acp.NewDecoder(rig.agentReader, 0)
	agentEnc := acp.NewEncoder(rig.agentWriter)
	agentDone := make(chan error, 1)
	go func() {
		req, err := agentDec.DecodeMessage()
		if err != nil {
			agentDone <- err
			return
		}
		if !req.IsRequest() || req.Method != "initialize" {
			agentDone <- errors.New("expected initialize request")
			return
		}
		resp, err := acp.NewResponseOK(req.ID, map[string]any{"protocolVersion": 1})
		if err != nil {
			agentDone <- err
			return
		}
		agentDone <- agentEnc.Encode(resp)
	}()

	ctx, cancelCall := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancelCall()
	resp, err := rig.r.Call(ctx, "initialize", map[string]any{"protocolVersion": 1})
	require.NoError(t, err)
	require.NotNil(t, resp)
	require.True(t, resp.IsResponse())

	var result struct {
		ProtocolVersion int `json:"protocolVersion"`
	}
	require.NoError(t, json.Unmarshal(resp.Result, &result))
	assert.Equal(t, 1, result.ProtocolVersion)

	// Confirm agent goroutine finished cleanly.
	select {
	case err := <-agentDone:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("agent goroutine did not finish in time")
	}

	// Both directions must have been persisted: outgoing request + incoming
	// response. Check repo accumulator.
	require.Eventually(t, func() bool {
		return len(rig.repo.snapshot()) >= 2
	}, 2*time.Second, 10*time.Millisecond, "expected at least 2 events persisted")

	events := rig.repo.snapshot()
	var sawReqIn, sawRespOut bool
	for _, ev := range events {
		switch {
		case ev.Direction == acp.DirectionIn && ev.RPCKind == acp.RPCRequest:
			sawReqIn = true
		case ev.Direction == acp.DirectionOut && ev.RPCKind == acp.RPCResponse:
			sawRespOut = true
		}
	}
	assert.True(t, sawReqIn, "expected request event with direction=in")
	assert.True(t, sawRespOut, "expected response event with direction=out")
}

// ---------- Test 2 ----------

func TestRelay_FanoutToSubscribers(t *testing.T) {
	t.Parallel()
	rig := newRelayTestRig(t)
	cancel, doneCh := rig.runRelay(t)
	t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })

	sub := rig.r.Subscribe(rig.sess.UserID)
	require.NotNil(t, sub)
	t.Cleanup(func() { rig.r.Unsubscribe(sub.ID) })

	agentEnc := acp.NewEncoder(rig.agentWriter)
	notif, err := acp.NewNotification("session/update", map[string]any{
		"sessionId": "agent-sess-1",
		"update":    map[string]any{"kind": "agent_message_chunk"},
	})
	require.NoError(t, err)
	require.NoError(t, agentEnc.Encode(notif))

	select {
	case env, ok := <-sub.Send:
		require.True(t, ok, "subscriber channel closed unexpectedly")
		require.NotNil(t, env)
		assert.Equal(t, "out", env.Direction)
		assert.Equal(t, "session/update", env.Method)
		assert.Equal(t, string(acp.RPCNotification), env.Kind)
	case <-time.After(3 * time.Second):
		t.Fatal("did not receive notification envelope within 3s")
	}
}

// ---------- Test 3 ----------

func TestRelay_AgentFsRequest_Echoed(t *testing.T) {
	t.Parallel()
	rig := newRelayTestRig(t)
	cancel, doneCh := rig.runRelay(t)
	t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })

	// Seed file inside the session cwd.
	fpath := filepath.Join(rig.sess.CwdPath, "x.txt")
	require.NoError(t, os.WriteFile(fpath, []byte("x"), 0o644))

	// Agent encoder/decoder for the subprocess side.
	agentDec := acp.NewDecoder(rig.agentReader, 0)
	agentEnc := acp.NewEncoder(rig.agentWriter)

	// Agent sends fs/read_text_file with int ID.
	req, err := acp.NewRequestInt(101, "fs/read_text_file", map[string]any{
		"path":      "x.txt",
		"sessionId": "agent-sess-1",
	})
	require.NoError(t, err)
	require.NoError(t, agentEnc.Encode(req))

	// Read the relay's response from the agent-side reader.
	respCh := make(chan *acp.Message, 1)
	errCh := make(chan error, 1)
	go func() {
		m, err := agentDec.DecodeMessage()
		if err != nil {
			errCh <- err
			return
		}
		respCh <- m
	}()

	select {
	case resp := <-respCh:
		require.True(t, resp.IsResponse(), "expected response from relay")
		require.Nil(t, resp.Error, "expected non-error response, got %+v", resp.Error)

		var result struct {
			Content string `json:"content"`
		}
		require.NoError(t, json.Unmarshal(resp.Result, &result))
		assert.Equal(t, "x", result.Content)
	case err := <-errCh:
		if !errors.Is(err, io.EOF) {
			t.Fatalf("agent decoder error: %v", err)
		}
		t.Fatal("agent saw EOF before relay responded")
	case <-time.After(3 * time.Second):
		t.Fatal("relay did not respond to fs/read_text_file within 3s")
	}
}

// ---------- Test 4 ----------

func TestRelay_SlowSubscriberDropped(t *testing.T) {
	t.Parallel()
	rig := newRelayTestRig(t)
	cancel, doneCh := rig.runRelay(t)
	t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })

	sub := rig.r.Subscribe(rig.sess.UserID)
	require.NotNil(t, sub)

	// Flood ~100 notifications from agent without reading sub.Send. The fanout
	// loop will eventually see sub.Send full (cap = WSSendBuffer = 16), enter
	// the slow-consumer branch, attempt to push a slow_consumer_disconnect
	// envelope (best-effort: select-default may drop it if buffer is still
	// full), and close the channel. Either outcome (final disc envelope or
	// closed channel) signals "subscriber dropped".
	agentEnc := acp.NewEncoder(rig.agentWriter)
	floodDone := make(chan struct{})
	go func() {
		defer close(floodDone)
		for i := 0; i < 100; i++ {
			notif, err := acp.NewNotification("session/update", map[string]any{
				"seq": i,
			})
			if err != nil {
				return
			}
			if err := agentEnc.Encode(notif); err != nil {
				return
			}
		}
	}()

	// Wait for the flood to complete (relay drains the pipe) before draining
	// sub.Send, otherwise the test would race and prevent the buffer from
	// filling up.
	select {
	case <-floodDone:
	case <-time.After(3 * time.Second):
		t.Fatal("flood did not complete within 3s; relay reader stalled?")
	}

	// Now drain sub.Send. Either we see slow_consumer_disconnect, or the
	// channel was closed. We must observe at least one of these within 3s.
	deadline := time.NewTimer(3 * time.Second)
	defer deadline.Stop()

	var sawSlowDisc, sawClose bool
drain:
	for {
		select {
		case env, ok := <-sub.Send:
			if !ok {
				sawClose = true
				break drain
			}
			if env != nil && env.Kind == "slow_consumer_disconnect" {
				sawSlowDisc = true
				// Keep draining until close (or we'll hit the deadline).
			}
		case <-deadline.C:
			t.Fatal("subscriber not dropped within 3s")
		}
	}
	assert.True(t, sawSlowDisc || sawClose,
		"expected slow_consumer_disconnect envelope or closed channel after flood")
}

// ===== PermissionRequest（权限审批框架，测试桩） =====
func (f *fakeEventRepo) InsertPermissionRequest(context.Context, *acp.PermissionRequest) (*acp.PermissionRequest, error) {
	return &acp.PermissionRequest{}, nil
}
func (f *fakeEventRepo) GetPermissionRequestByID(context.Context, uuid.UUID) (*acp.PermissionRequest, error) {
	return nil, nil
}
func (f *fakeEventRepo) ListPendingPermissionRequestsBySession(context.Context, uuid.UUID) ([]*acp.PermissionRequest, error) {
	return nil, nil
}
func (f *fakeEventRepo) ListPendingPermissionRequestsByUser(context.Context, uuid.UUID) ([]*acp.PermissionRequest, error) {
	return nil, nil
}
func (f *fakeEventRepo) DecidePermissionRequest(context.Context, uuid.UUID, acp.PermissionStatus, *string, *uuid.UUID) (*acp.PermissionRequest, bool, error) {
	return nil, false, nil
}
func (f *fakeEventRepo) ExpirePendingPermissionRequestsBySession(context.Context, uuid.UUID) (int64, error) {
	return 0, nil
}
func (f *fakeEventRepo) ExpireStalePermissionRequests(context.Context, time.Time) (int64, error) {
	return 0, nil
}