Cyrene_v2/internal/services/gateway/session_orchestrator_test.go

package gateway

import (
	"context"
	"fmt"
	"log/slog"
	"sync"
	"testing"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/loop"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// callRecord is one invocation captured by the fake runner.
type callRecord struct {
	sessionID string
	messages  []model.Message
	env       loop.SessionEnv
	ctx       context.Context
}

// gatedRunner is a controllable TurnRunner. Each Run blocks until it consumes one
// ticket from tickets, or until ctx is canceled (the interrupt path). When it
// proceeds it records the call and returns a deterministic reply. A nil err
// produces a reply holding the *incoming* messages as text, so a test can assert
// exactly what context the resumed turn carried.
type gatedRunner struct {
	mu      sync.Mutex
	calls   []callRecord
	tickets chan struct{}
	err     error
}

func newGatedRunner(cap int) *gatedRunner {
	return &gatedRunner{tickets: make(chan struct{}, cap)}
}

func (g *gatedRunner) Run(ctx context.Context, sessionID string, messages []model.Message, env loop.SessionEnv) (*loop.Result, error) {
	select {
	case <-g.tickets:
	case <-ctx.Done():
		return nil, ctx.Err()
	}
	g.mu.Lock()
	n := len(g.calls) + 1
	g.calls = append(g.calls, callRecord{sessionID: sessionID, messages: append([]model.Message(nil), messages...), env: env, ctx: ctx})
	g.mu.Unlock()
	if g.err != nil {
		return nil, g.err
	}
	reply := n + len(messages)
	return &loop.Result{Text: fmt.Sprintf("reply-%d", reply), Steps: 1, StopReason: "stop"}, nil
}

// waitCalls blocks until at least n calls are recorded and returns a copy of the
// first n, so a test can advance the gated chain deterministically.
func (g *gatedRunner) waitCalls(n int) []callRecord {
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		g.mu.Lock()
		if len(g.calls) >= n {
			out := append([]callRecord(nil), g.calls[:n]...)
			g.mu.Unlock()
			return out
		}
		g.mu.Unlock()
		time.Sleep(time.Millisecond)
	}
	panic("timed out waiting for runner calls")
}

// count returns the total recorded calls.
func (g *gatedRunner) count() int {
	g.mu.Lock()
	defer g.mu.Unlock()
	return len(g.calls)
}

// eventRec is one state transition captured by the recording sink.
type eventRec struct {
	sessionID, typ string
	data           map[string]any
}

// recordingSink captures every emitted state transition.
type recordingSink struct {
	mu     sync.Mutex
	events []eventRec
}

func (s *recordingSink) Emit(sessionID, typ string, data map[string]any) {
	s.mu.Lock()
	s.events = append(s.events, eventRec{sessionID: sessionID, typ: typ, data: data})
	s.mu.Unlock()
}

func (s *recordingSink) types() []string {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := make([]string, 0, len(s.events))
	for _, e := range s.events {
		out = append(out, e.typ)
	}
	return out
}

func (s *recordingSink) last() (string, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.events) == 0 {
		return "", false
	}
	return s.events[len(s.events)-1].typ, true
}

// userMsg builds a text-only user message.
func userMsg(text string) model.Message {
	return model.TextMessage(model.RoleUser, text)
}

// messageText returns the concatenated text of a message's text parts.
func messageText(m model.Message) string {
	var out string
	for _, c := range m.Content {
		if c.Type == model.ContentTypeText {
			out += c.Text
		}
	}
	return out
}

// seedOrchestrator builds an orchestrator with a gated runner and a recording
// sink, both returned so the test can drive the chain and assert state events.
func seedOrchestrator(run *gatedRunner, maxPending int) (*Orchestrator, *recordingSink) {
	sink := &recordingSink{}
	o := NewOrchestrator(OrchestratorOptions{
		Run:        run.Run,
		Emit:       sink.Emit,
		Log:        slog.Default(),
		MaxPending: maxPending,
		MaxTail:    20,
	})
	return o, sink
}

// readOnlyEnv returns the group-chat / stranger read-only memory lane.
func readOnlyEnv() loop.SessionEnv {
	return loop.SessionEnv{MemWritable: false}
}

func TestSubmitIdleStartsImmediate(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(1)
	o, _ := seedOrchestrator(run, 5)

	sr := o.Submit(context.Background(), "s1", []model.Message{userMsg("hi")}, loop.OwnerEnv())
	if !sr.Accepted || !sr.Immediate {
		t.Fatalf("submit = %+v, want Accepted+Immediate", sr)
	}
	run.tickets <- struct{}{}
	got := run.waitCalls(1)
	if len(got) != 1 {
		t.Fatalf("calls = %d, want 1", len(got))
	}
	if o.Busy("s1") {
		t.Fatal("session should be idle after a completed turn")
	}
}

func TestSubmitBusyPendingNotDropped(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(4)
	o, sink := seedOrchestrator(run, 5)

	// Turn A starts and blocks (no ticket yet), so the session is busy.
	a := o.Submit(context.Background(), "s1", []model.Message{userMsg("A")}, loop.OwnerEnv())
	if !a.Immediate {
		t.Fatalf("A = %+v, want Immediate", a)
	}
	// B and C arrive while A is in flight: they must be queued, not dropped.
	b := o.Submit(context.Background(), "s1", []model.Message{userMsg("B")}, readOnlyEnv())
	c := o.Submit(context.Background(), "s1", []model.Message{userMsg("C")}, loop.OwnerEnv())
	if !b.Accepted || !b.Pending || b.QueueDepth != 1 {
		t.Fatalf("B = %+v, want Pending depth=1", b)
	}
	if !c.Accepted || !c.Pending || c.QueueDepth != 2 {
		t.Fatalf("C = %+v, want Pending depth=2", c)
	}
	if o.QueueDepth("s1") != 2 {
		t.Fatalf("queue depth = %d, want 2", o.QueueDepth("s1"))
	}

	// Release the chain: A, then B, then C. The emitted ack should include the
	// pending depth, and the sink should have seen pending_ack events.
	types := sink.types()
	foundAck := false
	for _, ty := range types {
		if ty == EventPendingAck {
			foundAck = true
		}
	}
	if !foundAck {
		t.Fatalf("no pending_ack in events: %v", types)
	}

	run.tickets <- struct{}{}
	run.tickets <- struct{}{}
	run.tickets <- struct{}{}
	if n := len(run.waitCalls(3)); n != 3 {
		t.Fatalf("total calls = %d, want 3 (A+B+C)", n)
	}
	if o.QueueDepth("s1") != 0 || o.Busy("s1") {
		t.Fatalf("session should be idle and empty after drain: busy=%v depth=%d", o.Busy("s1"), o.QueueDepth("s1"))
	}
}

func TestSubmitBusyResumeCarriesPreviousContext(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(3)
	o, _ := seedOrchestrator(run, 5)

	if sr := o.Submit(context.Background(), "s1", []model.Message{userMsg("A")}, loop.OwnerEnv()); !sr.Immediate {
		t.Fatalf("A = %+v, want Immediate", sr)
	}
	if sr := o.Submit(context.Background(), "s1", []model.Message{userMsg("B")}, readOnlyEnv()); !sr.Pending {
		t.Fatalf("B = %+v, want Pending", sr)
	}
	if sr := o.Submit(context.Background(), "s1", []model.Message{userMsg("C")}, loop.OwnerEnv()); !sr.Pending {
		t.Fatalf("C = %+v, want Pending", sr)
	}

	run.tickets <- struct{}{} // A
	run.tickets <- struct{}{} // B
	run.tickets <- struct{}{} // C
	got := run.waitCalls(3)

	// A starts from a blank turn.
	if len(got[0].messages) != 1 || messageText(got[0].messages[0]) != "A" {
		t.Fatalf("A messages = %v, want [A]", messageTextList(got[0].messages))
	}
	// B must resume carrying A's assistant reply folded into the tail.
	if !matchesText(got[1].messages, "A", "B") {
		t.Fatalf("B messages = %v, want B to resume carrying the prior tail", messageTextList(got[1].messages))
	}
	// C must resume carrying A's and B's assistant replies.
	if !matchesText(got[2].messages, "A", "C") {
		t.Fatalf("C messages = %v, want C to resume carrying the prior tail", messageTextList(got[2].messages))
	}
	// B's env must be the one the message arrived with (read-only lane).
	if got[1].env.MemWritable {
		t.Fatal("B resumed with MemWritable=true, want the caller's read-only env")
	}
}

// matchesText reports whether a resumed turn begins with the captured first user
// turn (so the prior context was carried over), ends with the expected new user
// turn, and contains at least one assistant reply in between.
func matchesText(msgs []model.Message, first, last string) bool {
	if len(msgs) < 3 {
		return false
	}
	if messageText(msgs[0]) != first {
		return false
	}
	if messageText(msgs[len(msgs)-1]) != last {
		return false
	}
	for i := 1; i < len(msgs)-1; i++ {
		if msgs[i].Role == model.RoleAssistant {
			return true
		}
	}
	return false
}

func messageTextList(msgs []model.Message) []string {
	out := make([]string, 0, len(msgs))
	for _, m := range msgs {
		out = append(out, messageText(m))
	}
	return out
}

func TestSubmitPendingBoundedRejectsOverflow(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(8)
	o, _ := seedOrchestrator(run, 3) // tight queue

	if sr := o.Submit(context.Background(), "s1", []model.Message{userMsg("A")}, loop.OwnerEnv()); !sr.Immediate {
		t.Fatalf("A = %+v, want Immediate", sr)
	}
	for i := 0; i < 3; i++ {
		sr := o.Submit(context.Background(), "s1", []model.Message{userMsg(fmt.Sprintf("p%d", i))}, loop.OwnerEnv())
		if !sr.Pending || sr.QueueDepth != i+1 {
			t.Fatalf("pending %d = %+v, want Pending depth=%d", i, sr, i+1)
		}
	}
	// Queue is at capacity (3), so the next message is refused, not dropped.
	over := o.Submit(context.Background(), "s1", []model.Message{userMsg("overflow")}, loop.OwnerEnv())
	if over.Accepted || !over.PendingFull {
		t.Fatalf("overflow = %+v, want PendingFull reject", over)
	}
	if o.QueueDepth("s1") != 3 {
		t.Fatalf("queue depth = %d, want 3 (unchanged)", o.QueueDepth("s1"))
	}
}

func TestInterruptCancelsTurnAndClearsPending(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(2)
	o, _ := seedOrchestrator(run, 5)

	if sr := o.Submit(context.Background(), "s1", []model.Message{userMsg("A")}, loop.OwnerEnv()); !sr.Immediate {
		t.Fatalf("A = %+v, want Immediate", sr)
	}
	if sr := o.Submit(context.Background(), "s1", []model.Message{userMsg("B")}, loop.OwnerEnv()); !sr.Pending {
		t.Fatalf("B = %+v, want Pending", sr)
	}

	// Release no ticket: A is blocked in flight. Interrupt cancels it and clears B.
	dropped := o.Interrupt("s1")
	if dropped != 1 {
		t.Fatalf("dropped = %d, want 1", dropped)
	}
	if o.Busy("s1") || o.QueueDepth("s1") != 0 {
		t.Fatalf("session not idle after interrupt: busy=%v depth=%d", o.Busy("s1"), o.QueueDepth("s1"))
	}

	// A new message after the interrupt must start fresh (Immediate), and the
	// canceled chain must not have resumed B.
	if sr := o.Submit(context.Background(), "s1", []model.Message{userMsg("C")}, loop.OwnerEnv()); !sr.Immediate {
		t.Fatalf("C = %+v, want Immediate after interrupt", sr)
	}
	run.tickets <- struct{}{}
	got := run.waitCalls(1)
	if len(got) != 1 {
		t.Fatalf("calls = %d, want 1 (only C, not the canceled A or B)", len(got))
	}
	if messageText(got[0].messages[0]) != "C" {
		t.Fatalf("only surviving call = %v, want the C turn", messageTextList(got[0].messages))
	}
}

func TestSubmitConcurrentSameSessionIsSafeAndBounded(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(8)
	o, _ := seedOrchestrator(run, 3)

	const senders = 8
	var wg sync.WaitGroup
	accepted := make(chan SubmitResult, senders)
	for i := 0; i < senders; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			accepted <- o.Submit(context.Background(), "s1", []model.Message{userMsg(fmt.Sprintf("m%d", i))}, loop.OwnerEnv())
		}(i)
	}
	wg.Wait()
	close(accepted)

	immediate, pending, full := 0, 0, 0
	for sr := range accepted {
		switch {
		case sr.Immediate:
			immediate++
		case sr.Pending:
			pending++
		case sr.PendingFull:
			full++
		default:
			t.Fatalf("unexpected submit result: %+v", sr)
		}
	}
	// With maxPending=3: exactly 1 turn runs immediately and up to 3 queue behind
	// it; the remaining senders are refused at capacity rather than dropped.
	if immediate != 1 {
		t.Fatalf("immediate = %d, want 1", immediate)
	}
	if pending != 3 {
		t.Fatalf("pending = %d, want 3", pending)
	}
	if full != senders-4 {
		t.Fatalf("pendingFull = %d, want %d", full, senders-4)
	}
	if o.QueueDepth("s1") != 3 {
		t.Fatalf("queue depth = %d, want 3", o.QueueDepth("s1"))
	}

	// Drain the chain: 1 running + 3 queued = 4 turns. No panic / no lost message.
	for i := 0; i < 4; i++ {
		run.tickets <- struct{}{}
	}
	if n := len(run.waitCalls(4)); n != 4 {
		t.Fatalf("total calls = %d, want 4", n)
	}
	if o.Busy("s1") || o.QueueDepth("s1") != 0 {
		t.Fatalf("session should be idle and empty after drain")
	}
}

func TestSubmitEmptyRejected(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(1)
	o, _ := seedOrchestrator(run, 5)
	sr := o.Submit(context.Background(), "s1", nil, loop.OwnerEnv())
	if sr.Accepted || !sr.Empty {
		t.Fatalf("empty submit = %+v, want Empty reject", sr)
	}
	if run.count() != 0 {
		t.Fatalf("runner invoked for an empty turn")
	}
}

func TestSubmitSyncIdleReturnsResultBusyReturnsPending(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(2)
	o, _ := seedOrchestrator(run, 5)

	// Idle → synchronous run returns the reply immediately. Push a ticket so the
	// gated runner's blocking select resolves during the synchronous wait.
	run.tickets <- struct{}{}
	sr, res := o.SubmitSync(context.Background(), "s1", []model.Message{userMsg("hi")}, loop.OwnerEnv())
	if !sr.Accepted || !sr.Immediate {
		t.Fatalf("sync idle = %+v, want Accepted+Immediate", sr)
	}
	if res == nil {
		t.Fatal("sync idle should return a result")
	}

	// Hold a busy turn in flight via an async Submit (blocked, no ticket), then a
	// sync submit must return a Pending ack without blocking and without a result.
	if sr2 := o.Submit(context.Background(), "s1", []model.Message{userMsg("A")}, loop.OwnerEnv()); !sr2.Immediate {
		t.Fatalf("async A = %+v, want Immediate", sr2)
	}
	sr3, res3 := o.SubmitSync(context.Background(), "s1", []model.Message{userMsg("B")}, loop.OwnerEnv())
	if !sr3.Accepted || !sr3.Pending {
		t.Fatalf("sync busy = %+v, want Pending", sr3)
	}
	if res3 != nil {
		t.Fatal("sync busy should not return a result (must not block)")
	}
	run.tickets <- struct{}{} // let A finish so the orchestrator is idle again
}

func TestInterruptAllClearsEverySession(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(2)
	o, _ := seedOrchestrator(run, 5)

	o.Submit(context.Background(), "s1", []model.Message{userMsg("A")}, loop.OwnerEnv())
	o.Submit(context.Background(), "s1", []model.Message{userMsg("B")}, loop.OwnerEnv())
	o.Submit(context.Background(), "s2", []model.Message{userMsg("X")}, loop.OwnerEnv())

	if got := o.InterruptAll(); got != 1 {
		t.Fatalf("InterruptAll dropped = %d, want 1 (B)", got)
	}
	if o.Busy("s1") || o.Busy("s2") {
		t.Fatal("all sessions should be idle after InterruptAll")
	}
}

func TestSetModeBusyAndIdle(t *testing.T) {
	t.Parallel()
	run := newGatedRunner(1)
	o, sink := seedOrchestrator(run, 5)

	o.SetMode("s1", ModeExecuting)
	if !o.Busy("s1") || o.ModeOf("s1") != ModeExecuting {
		t.Fatalf("after executing: busy=%v mode=%q", o.Busy("s1"), o.ModeOf("s1"))
	}
	o.SetMode("s1", ModeIdle)
	if o.Busy("s1") {
		t.Fatal("session must be idle after SetMode(idle)")
	}
	if ty, ok := sink.last(); !ok || ty != EventSessionIdle {
		t.Fatalf("last event = %q, want session_idle", ty)
	}
}