From 85a8bc0df7a564b630d31664c3b9af3316cdf8ac Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Thu, 7 May 2026 13:45:25 +0800 Subject: [PATCH] test(acp): relay round-trip + fanout + agent request handling + slow subscriber --- internal/acp/relay_test.go | 428 +++++++++++++++++++++++++++++++++++++ 1 file changed, 428 insertions(+) create mode 100644 internal/acp/relay_test.go diff --git a/internal/acp/relay_test.go b/internal/acp/relay_test.go new file mode 100644 index 0000000..8bd1f04 --- /dev/null +++ b/internal/acp/relay_test.go @@ -0,0 +1,428 @@ +// 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/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) ([]*acp.Session, error) { + panic("n/a") +} +func (f *fakeEventRepo) ListAllSessions(context.Context) ([]*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") +} + +// 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() + + 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") +}