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

	// 回合状态机：relay turn 测试用。
	turns          []*acp.Turn
	lastStopReason *string

	// 成本核算：relay usage 测试用。
	usageLedger  []*acp.SessionUsageRecord
	totPrompt    int64
	totComplete  int64
	totThink     int64
	totCost      float64
	dedupSources map[int64]struct{}
}

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) GetSessionByStepID(context.Context, uuid.UUID) (*acp.Session, error) {
	panic("n/a")
}
func (f *fakeEventRepo) GetCrashedSessionForResume(context.Context, uuid.UUID) (*acp.Session, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ResetSessionForResume(context.Context, uuid.UUID) 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) UpdateSessionSandboxMode(context.Context, uuid.UUID, string) error {
	panic("n/a")
}
func (f *fakeEventRepo) UpdateSessionFinished(context.Context, uuid.UUID, acp.SessionStatus, *int32, *string) error {
	panic("n/a")
}
func (f *fakeEventRepo) MarkSessionFailedIfActive(context.Context, uuid.UUID, string) (bool, 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) InsertSessionUsage(_ context.Context, rec *acp.SessionUsageRecord) (bool, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if rec.SourceEventID != nil {
		if f.dedupSources == nil {
			f.dedupSources = map[int64]struct{}{}
		}
		if _, seen := f.dedupSources[*rec.SourceEventID]; seen {
			return false, nil
		}
		f.dedupSources[*rec.SourceEventID] = struct{}{}
	}
	cp := *rec
	f.usageLedger = append(f.usageLedger, &cp)
	return true, nil
}
func (f *fakeEventRepo) AddSessionUsageTotals(_ context.Context, _ uuid.UUID, dp, dc, dt int64, dCost float64) (int64, int64, int64, float64, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.totPrompt += dp
	f.totComplete += dc
	f.totThink += dt
	f.totCost += dCost
	return f.totPrompt, f.totComplete, f.totThink, f.totCost, nil
}
func (f *fakeEventRepo) SumProjectCostUSD(context.Context, uuid.UUID, time.Time) (float64, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.totCost, nil
}
func (f *fakeEventRepo) MarkSessionTerminatedReason(context.Context, uuid.UUID, string) error {
	return nil
}
func (f *fakeEventRepo) ListSessionsForReaper(context.Context, time.Time, time.Time) ([]acp.ReaperSession, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ListSessionUsage(context.Context, uuid.UUID, int32) ([]*acp.SessionUsageRecord, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]*acp.SessionUsageRecord, len(f.usageLedger))
	copy(out, f.usageLedger)
	return out, nil
}
func (f *fakeEventRepo) SummarizeAcpUsageByUser(context.Context, time.Time, time.Time) ([]acp.AcpUsageUserRow, error) {
	panic("n/a")
}
func (f *fakeEventRepo) SummarizeAcpUsageByModel(context.Context, time.Time, time.Time) ([]acp.AcpUsageModelRow, error) {
	panic("n/a")
}
func (f *fakeEventRepo) SummarizeAcpUsageByDay(context.Context, time.Time, time.Time) ([]acp.AcpUsageDayRow, error) {
	panic("n/a")
}
func (f *fakeEventRepo) SessionRollup(context.Context, acp.DashboardFilter) (acp.SessionRollup, error) {
	panic("n/a")
}
func (f *fakeEventRepo) SessionsByDay(context.Context, acp.DashboardFilter) ([]acp.SessionDayRow, error) {
	panic("n/a")
}
func (f *fakeEventRepo) SessionTimeline(context.Context, uuid.UUID) ([]acp.SessionTimelineBucket, error) {
	panic("n/a")
}
func (f *fakeEventRepo) ListConfigFiles(context.Context, uuid.UUID) ([]*acp.ConfigFile, error) {
	panic("n/a")
}
func (f *fakeEventRepo) UpsertConfigFile(context.Context, *acp.ConfigFile) (*acp.ConfigFile, error) {
	panic("n/a")
}
func (f *fakeEventRepo) DeleteConfigFile(context.Context, uuid.UUID, string) (bool, 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
}

// ----- Turn 方法（relay turn 测试用，功能性实现）-----

func (f *fakeEventRepo) InsertTurn(_ context.Context, t *acp.Turn) (*acp.Turn, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	cp := *t
	cp.ID = int64(len(f.turns) + 1)
	cp.Status = acp.TurnInProgress
	f.turns = append(f.turns, &cp)
	out := cp
	return &out, nil
}
func (f *fakeEventRepo) MarkTurnCompleted(_ context.Context, _ uuid.UUID, turnIndex int, stop string, updateCount int) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, t := range f.turns {
		if t.TurnIndex == turnIndex {
			t.Status = acp.TurnCompleted
			s := stop
			t.StopReason = &s
			t.UpdateCount = updateCount
		}
	}
	return nil
}
func (f *fakeEventRepo) MarkTurnAborted(_ context.Context, _ uuid.UUID, turnIndex int) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, t := range f.turns {
		if t.TurnIndex == turnIndex {
			t.Status = acp.TurnAborted
		}
	}
	return nil
}
func (f *fakeEventRepo) IncrementTurnUpdateCount(context.Context, uuid.UUID, int) error { return nil }
func (f *fakeEventRepo) GetLatestTurnBySession(context.Context, uuid.UUID) (*acp.Turn, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if len(f.turns) == 0 {
		return nil, nil
	}
	out := *f.turns[len(f.turns)-1]
	return &out, nil
}
func (f *fakeEventRepo) ListTurnsBySession(context.Context, uuid.UUID) ([]*acp.Turn, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]*acp.Turn, len(f.turns))
	copy(out, f.turns)
	return out, nil
}
func (f *fakeEventRepo) UpdateSessionLastStopReason(_ context.Context, _ uuid.UUID, reason string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	r := reason
	f.lastStopReason = &r
	return nil
}
func (f *fakeEventRepo) PurgeTurnsBefore(context.Context, time.Time) (int64, error) { return 0, nil }

func (f *fakeEventRepo) turnSnapshot() []*acp.Turn {
	f.mu.Lock()
	defer f.mu.Unlock()
	out := make([]*acp.Turn, len(f.turns))
	for i, t := range f.turns {
		cp := *t
		out[i] = &cp
	}
	return out
}
func (f *fakeEventRepo) lastStop() *string {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.lastStopReason
}

// 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
}

// ---------- Turn detection ----------

// TestRelay_TurnDetection feeds a session/update notification then a string-id
// session/prompt response with stopReason, asserting the tracker hooks fire:
// the turn is created, completed with the stop reason, and update_count tracked.
func TestRelay_TurnDetection(t *testing.T) {
	t.Parallel()
	rig := newRelayTestRig(t)

	bus := &fakeBus{}
	tr := acp.NewTurnTracker(rig.repo, bus, rig.sess, nil)
	rig.r.SetTracker(tr)

	cancel, doneCh := rig.runRelay(t)
	t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })

	// Register the turn (as session_service would before sending the prompt).
	_, err := tr.StartTurn(context.Background(), "client-init")
	require.NoError(t, err)

	agentEnc := acp.NewEncoder(rig.agentWriter)

	// Agent emits two session/update notifications then the prompt response.
	notif, err := acp.NewNotification("session/update", map[string]any{
		"sessionId": "agent-sess-1",
		"update":    map[string]any{"sessionUpdate": "agent_message_chunk"},
	})
	require.NoError(t, err)
	require.NoError(t, agentEnc.Encode(notif))
	require.NoError(t, agentEnc.Encode(notif))

	idJSON, _ := json.Marshal("client-init")
	resp := &acp.Message{
		JSONRPC: "2.0", ID: idJSON,
		Result: json.RawMessage(`{"stopReason":"max_tokens"}`),
	}
	require.NoError(t, agentEnc.Encode(resp))

	// The turn must complete with stop_reason=max_tokens and update_count=2.
	require.Eventually(t, func() bool {
		ts := rig.repo.turnSnapshot()
		if len(ts) != 1 {
			return false
		}
		return ts[0].Status == acp.TurnCompleted &&
			ts[0].StopReason != nil && *ts[0].StopReason == "max_tokens" &&
			ts[0].UpdateCount == 2
	}, 3*time.Second, 20*time.Millisecond, "turn not completed as expected")

	// last_stop_reason persisted on the session.
	require.Eventually(t, func() bool {
		ls := rig.repo.lastStop()
		return ls != nil && *ls == "max_tokens"
	}, 2*time.Second, 20*time.Millisecond)

	// Bus received turn_completed then session_idle.
	require.Eventually(t, func() bool {
		return len(bus.snapshot()) >= 2
	}, 2*time.Second, 20*time.Millisecond)
	evs := bus.snapshot()
	assert.Equal(t, acp.TurnCompletedEvent, evs[0].Kind)
	assert.Equal(t, acp.SessionIdleEvent, evs[1].Kind)
	assert.Equal(t, acp.StopMaxTokens, evs[0].StopReason)
}

// TestRelay_IntResponseStillRoutesToInflight is a regression guard: with a
// tracker installed, an int64-id response must still route to inflight (the
// existing Call mechanism), not the turn path.
func TestRelay_IntResponseStillRoutesToInflight(t *testing.T) {
	t.Parallel()
	rig := newRelayTestRig(t)
	rig.r.SetTracker(acp.NewTurnTracker(rig.repo, &fakeBus{}, rig.sess, nil))
	cancel, doneCh := rig.runRelay(t)
	t.Cleanup(func() { rig.shutdown(t, cancel, doneCh) })

	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
		}
		respMsg, err := acp.NewResponseOK(req.ID, map[string]any{"protocolVersion": 1})
		if err != nil {
			agentDone <- err
			return
		}
		agentDone <- agentEnc.Encode(respMsg)
	}()

	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.True(t, resp.IsResponse())
	select {
	case err := <-agentDone:
		require.NoError(t, err)
	case <-time.After(2 * time.Second):
		t.Fatal("agent goroutine did not finish")
	}
}