Cyrene_v2/internal/services/gateway/session_orchestrator.go

package gateway

import (
	"context"
	"io"
	"log/slog"
	"sync"
	"time"

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

// This file implements the session-message orchestrator (DevDocs/design/15).
//
// It is deliberately NOT a task engine. design/04 §2 (task/subagent) decides
// whether a heavy job gets pushed to the background; this orchestrator decides
// how NEW user messages queue and resume while a session is busy. The two layers
// are orthogonal and compose: a session runs a task (04), a new message arrives
// mid-run (15), and the orchestrator holds it in a bounded pending queue and
// resumes it with the previous turn's context once the current turn finishes.
//
// Non-blocking is a hard constraint: Submit/Interrupt never wait on a running
// turn. The queue is bounded per session and isolated across sessions, and no
// identity decision is made here — the orchestrator only orders intent/messages,
// identity stays in the channel actor metadata (project red line).

// Mode is the per-session busy-state lane of the orchestrator. It is exposed so
// a loop/subagent hook can raise it to executing/subagent (an injection point
// the composition root may wire later), but the orchestrator itself tracks an
// abstract "busy" via mode != ModeIdle.
type Mode string

const (
	// ModeIdle means the session has no turn in flight and no resume pending.
	ModeIdle Mode = "idle"
	// ModeThinking means a turn is in flight (the initial start state).
	ModeThinking Mode = "thinking"
	// ModeExecuting means a turn is executing tools/substeps.
	ModeExecuting Mode = "executing"
	// ModeSubagent means a turn is delegated to a subagent (design/04 §2).
	ModeSubagent Mode = "subagent"
)

// State event types pushed over the WS channel to the frontend. The data map of
// each carries sessionId so a client can correlate per-session state.
const (
	EventSessionBusy      = "session_busy"
	EventPendingAck       = "pending_ack"
	EventPendingDrain     = "pending_drain"
	EventSessionIdle      = "session_idle"
	EventSessionInterrupt = "session_interrupted"
	EventPendingFull      = "pending_full"
)

const (
	// defaultMaxPending bounds the per-session pending queue (design/15 §4).
	defaultMaxPending = 5
	// defaultMaxTail bounds the conversation tail kept for continuation, so the
	// resumed context is incremental, not a full replay (design/15 §2.2, §4).
	defaultMaxTail = 20
	// defaultSessionID is used when a caller passes an empty session id.
	defaultSessionID = "web"
)

// SubmitResult is the outcome of handing one user turn to the orchestrator. It
// is always returned without blocking: Accepted is false only when the turn was
// rejected (pending queue full, or an empty turn), never because a turn is busy.
type SubmitResult struct {
	// Accepted reports whether the turn was taken. A false value with
	// PendingFull or Empty carries the reject reason; it never means "try later".
	Accepted bool
	// Immediate reports that a turn started now because the session was idle.
	Immediate bool
	// Pending reports that the turn was queued behind a busy turn.
	Pending bool
	// QueueDepth is the pending depth after this insertion (0 when Immediate).
	QueueDepth int
	// PendingFull reports the turn was rejected because the queue is at capacity
	// (design/15 §4: bounded pending, over-capacity is refused with a hint).
	PendingFull bool
	// Empty reports the turn was rejected because it carried no text or image.
	Empty bool
	// Err is set on the synchronous path when the run errored; the REST caller
	// uses it to return a 5xx rather than a fabricated reply.
	Err error
}

// TurnRunner executes one agent turn and returns its result. The gateway wires
// it to the loop plus the streaming broadcast, so the orchestrator stays agnostic
// of how a turn is physically driven (it never re-implements a loop engine).
type TurnRunner func(ctx context.Context, sessionID string, messages []model.Message, env loop.SessionEnv) (*loop.Result, error)

// StateSink reports a per-session state transition (busy/pending/idle/interrupt)
// to the transport. The gateway wires it to the WS hub broadcast.
type StateSink func(sessionID, typ string, data map[string]any)

// OrchestratorOptions configures an Orchestrator. Run and Log are required
// (Log may be a discard logger); Emit is optional (state events are dropped when
// nil). MaxPending/MaxTail fall back to the defaults when non-positive.
type OrchestratorOptions struct {
	Run        TurnRunner
	Emit       StateSink
	Log        *slog.Logger
	MaxPending int
	MaxTail    int
	// History is the optional persistent chat-history store (design/16 §3.3).
	// When nil every Append is skipped, so the orchestrator runs without storage
	// (tests / a gateway built without a history store).
	History HistoryStore
	// OnPersisted, when non-nil, is called with the complete stored message right
	// after a user or assistant message is written, so the transport can broadcast
	// a `message_persisted` event for cross-end realtime sync (design/16 §4.2).
	// Dropped when nil (tests / a gateway built without a transport hook).
	OnPersisted func(history.Message)
}

// sessionState is the per-session message state (design/15 §2). Every mutation
// happens under the Orchestrator mutex; the actual loop run is done outside the
// lock so a slow model call never blocks Submit.
type sessionState struct {
	mode    Mode
	gen     uint64 // generation: bumped on interrupt and on a fresh top-level turn
	cancel  context.CancelFunc
	pending []pendingMsg
	tail    []model.Message // bounded conversation tail for context continuation
}

// pendingMsg is one queued user turn. It carries its own env so a resume runs
// under the same memory-isolation boundary the message arrived with, plus the
// client deviceID of the originating turn (design/16 §4.5) so the resumed
// assistant reply is tagged with the same client that sent the message.
type pendingMsg struct {
	messages []model.Message
	env      loop.SessionEnv
	deviceID string
	at       time.Time
}

// Orchestrator owns the per-session message state machine. It is concurrency-safe
// and non-blocking: Submit/Interrupt lock only to mutate state, then hand the
// actual run off to a goroutine. It is the single owner of "is this session busy"
// for the chat path; the loop/task/subagent engines are reused, not re-created.
type Orchestrator struct {
	mu          sync.Mutex
	states      map[string]*sessionState
	run         TurnRunner
	emit        StateSink
	log         *slog.Logger
	history     HistoryStore
	onPersisted func(history.Message)
	maxPending  int
	maxTail     int
}

// NewOrchestrator builds an orchestrator from opts. Run is required: it is the
// turn-driver the orchestrator reuses (the loop, never a re-implemented engine).
// A nil Run makes every Submit reject rather than run, so a misconfiguration
// cannot silently leave a session busy. Tests inject a fake runner here.
func NewOrchestrator(opts OrchestratorOptions) *Orchestrator {
	maxPending := opts.MaxPending
	if maxPending <= 0 {
		maxPending = defaultMaxPending
	}
	maxTail := opts.MaxTail
	if maxTail <= 0 {
		maxTail = defaultMaxTail
	}
	if opts.Log == nil {
		opts.Log = slog.New(slog.NewTextHandler(io.Discard, nil))
	}
	return &Orchestrator{
		states:      make(map[string]*sessionState),
		run:         opts.Run,
		emit:        opts.Emit,
		log:         opts.Log,
		history:     opts.History,
		onPersisted: opts.OnPersisted,
		maxPending:  maxPending,
		maxTail:     maxTail,
	}
}

// Submit hands one user turn to the session. It never blocks: an idle session
// starts a turn on a background goroutine and returns Immediate; a busy session
// enqueues the turn into the bounded pending queue and returns Pending; a full
// queue refuses with PendingFull. Interrupt words are NOT auto-detected here —
// the caller (frontend/intent classifier) calls Interrupt when the user asks to
// stop, per design/15 §3 (interrupt is an intent, not an identity decision).
func (o *Orchestrator) Submit(ctx context.Context, sessionID string, messages []model.Message, env loop.SessionEnv) SubmitResult {
	sr, _ := o.submitInternal(ctx, o.normalize(sessionID), messages, env, false)
	return sr
}

// SubmitSync is the synchronous variant used by the REST non-streaming path. An
// idle session runs the turn to completion and returns the Result; a busy session
// returns a Pending result without blocking so the caller can reply with a
// non-blocking ack (design/15 §2.1). It never waits on an in-flight turn.
func (o *Orchestrator) SubmitSync(ctx context.Context, sessionID string, messages []model.Message, env loop.SessionEnv) (SubmitResult, *loop.Result) {
	return o.submitInternal(ctx, o.normalize(sessionID), messages, env, true)
}

// submitInternal is the shared Submit/SubmitSync body. sync selects the
// synchronous turn path (REST) vs the background goroutine path (WS).
func (o *Orchestrator) submitInternal(ctx context.Context, sessionID string, messages []model.Message, env loop.SessionEnv, sync bool) (SubmitResult, *loop.Result) {
	if o.run == nil {
		o.log.Error("orchestrator submit", "err", "no runner configured", "session", sessionID)
		return SubmitResult{Accepted: false}, nil
	}
	if o.messagesEmpty(messages) {
		return SubmitResult{Empty: true}, nil
	}

	o.mu.Lock()
	st := o.states[sessionID]
	if st == nil {
		st = &sessionState{mode: ModeIdle}
		o.states[sessionID] = st
	}
	deviceID := deviceIDFrom(ctx)
	if st.mode != ModeIdle {
		if len(st.pending) >= o.maxPending {
			mode := st.mode
			o.mu.Unlock()
			o.emitState(sessionID, EventPendingFull, map[string]any{"sessionId": sessionID, "mode": mode, "max": o.maxPending})
			return SubmitResult{PendingFull: true}, nil
		}
		st.pending = append(st.pending, pendingMsg{messages: messages, env: env, deviceID: deviceID, at: time.Now()})
		depth := len(st.pending)
		mode := st.mode
		o.mu.Unlock()
		o.persistUserMessages(ctx, sessionID, messages, deviceID)
		o.emitState(sessionID, EventPendingAck, map[string]any{"sessionId": sessionID, "mode": mode, "queueDepth": depth})
		return SubmitResult{Accepted: true, Pending: true, QueueDepth: depth}, nil
	}

	st.gen++
	gen := st.gen
	st.mode = ModeThinking
	st.tail = o.boundedTail(messages)
	runCtx, cancel := context.WithCancel(ctx)
	st.cancel = cancel
	o.mu.Unlock()

	o.emitState(sessionID, EventSessionBusy, map[string]any{"sessionId": sessionID, "mode": ModeThinking})
	o.persistUserMessages(ctx, sessionID, messages, deviceID)

	if sync {
		res, err := o.run(runCtx, sessionID, messages, env)
		if err != nil {
			o.log.Error("orchestrator turn", "err", err, "session", sessionID)
		}
		_ = o.completeTurn(sessionID, gen, res, err, deviceID)
		return SubmitResult{Accepted: true, Immediate: true, Err: err}, res
	}

	go o.goTurn(sessionID, gen, messages, env, runCtx, cancel, deviceID)
	return SubmitResult{Accepted: true, Immediate: true}, nil
}

// goTurn runs a single turn on a background goroutine and, when it finishes,
// hands the outcome to completeTurn so the orchestrator can resume the queue.
// deviceID tags the persistent assistant reply with the client that sent the turn
// (design/16 §4.5).
func (o *Orchestrator) goTurn(sessionID string, gen uint64, messages []model.Message, env loop.SessionEnv, ctx context.Context, cancel context.CancelFunc, deviceID string) {
	defer cancel()
	res, err := o.run(ctx, sessionID, messages, env)
	if err != nil && ctx.Err() == nil {
		o.log.Error("orchestrator turn", "err", err, "session", sessionID)
	}
	_ = o.completeTurn(sessionID, gen, res, err, deviceID)
}

// completeTurn runs after a turn settles (either it finished, errored, or was
// canceled). It folds the assistant reply into the session's continuation tail,
// then pops the next pending turn and resumes it carrying that tail — so the next
// message "remembers" the prior context instead of starting from a blank slate
// (design/15 §2.2: incremental, not full replay). A stale generation (an
// interrupt or a newer top-level turn superseded this one) is a no-op.
// deviceID is the originating turn's client id, used to tag the persisted reply.
func (o *Orchestrator) completeTurn(sessionID string, gen uint64, res *loop.Result, runErr error, deviceID string) error {
	o.persistAssistantReply(sessionID, res, deviceID)
	o.mu.Lock()
	st := o.states[sessionID]
	if st == nil || st.gen != gen {
		o.mu.Unlock()
		return nil
	}
	if res != nil && res.Text != "" {
		st.tail = o.boundedTail(concatMessages(st.tail, model.TextMessage(model.RoleAssistant, res.Text)))
	}
	if len(st.pending) > 0 {
		next := st.pending[0]
		st.pending = st.pending[1:]
		st.tail = o.boundedTail(concatMessages(st.tail, next.messages...))
		st.mode = ModeThinking
		st.gen++
		gen = st.gen
		runCtx, cancel := context.WithCancel(context.Background())
		st.cancel = cancel
		depth := len(st.pending)
		resume := append([]model.Message(nil), st.tail...)
		env := next.env
		o.mu.Unlock()

		o.emitState(sessionID, EventPendingDrain, map[string]any{"sessionId": sessionID, "mode": ModeThinking, "queueDepth": depth})
		go o.goTurn(sessionID, gen, resume, env, runCtx, cancel, next.deviceID)
		return nil
	}

	st.mode = ModeIdle
	st.cancel = nil
	o.mu.Unlock()

	o.emitState(sessionID, EventSessionIdle, map[string]any{"sessionId": sessionID, "mode": ModeIdle})
	return nil
}

// Interrupt cancels the in-flight turn, clears the pending queue and returns the
// session to idle (design/15 §3, the "停下/算了/换个话题" path). It is the only
// way a busy session resets without draining the queue; the running goroutine's
// completion becomes a stale generation and does not resume anything. It returns
// the number of pending turns that were dropped.
func (o *Orchestrator) Interrupt(sessionID string) int {
	return o.interrupt(o.normalize(sessionID))
}

// InterruptAll cancels every in-flight session and clears all pending queues. It
// is used on gateway shutdown so no background turn outlives the hub. It returns
// the total number of dropped pending turns.
func (o *Orchestrator) InterruptAll() int {
	o.mu.Lock()
	total := 0
	for id, st := range o.states {
		total += o.interruptLocked(id, st)
	}
	o.mu.Unlock()
	return total
}

func (o *Orchestrator) interrupt(sessionID string) int {
	o.mu.Lock()
	defer o.mu.Unlock()
	return o.interruptLocked(sessionID, o.states[sessionID])
}

func (o *Orchestrator) interruptLocked(sessionID string, st *sessionState) int {
	if st == nil || st.mode == ModeIdle {
		return 0
	}
	if st.cancel != nil {
		st.cancel()
	}
	dropped := len(st.pending)
	st.pending = nil
	st.gen++
	st.mode = ModeIdle
	st.cancel = nil
	o.emitState(sessionID, EventSessionInterrupt, map[string]any{"sessionId": sessionID, "mode": ModeIdle, "dropped": dropped})
	return dropped
}

// SetMode raises the reported busy-state lane of a session (thinking →
// executing → subagent). It is an injection point a loop/subagent hook can call;
// any non-idle mode keeps the session busy, and idle un-busies it. A nil session
// is created lazily so a hook may set a mode before the first message.
func (o *Orchestrator) SetMode(sessionID string, mode Mode) {
	sessionID = o.normalize(sessionID)
	o.mu.Lock()
	st := o.states[sessionID]
	if st == nil {
		st = &sessionState{mode: mode}
		o.states[sessionID] = st
	} else {
		st.mode = mode
	}
	o.mu.Unlock()
	if mode == ModeIdle {
		o.emitState(sessionID, EventSessionIdle, map[string]any{"sessionId": sessionID, "mode": mode})
		return
	}
	o.emitState(sessionID, EventSessionBusy, map[string]any{"sessionId": sessionID, "mode": mode})
}

// Busy reports whether the session currently has a turn in flight.
func (o *Orchestrator) Busy(sessionID string) bool {
	o.mu.Lock()
	defer o.mu.Unlock()
	st := o.states[o.normalize(sessionID)]
	return st != nil && st.mode != ModeIdle
}

// QueueDepth reports the number of pending turns for a session.
func (o *Orchestrator) QueueDepth(sessionID string) int {
	o.mu.Lock()
	defer o.mu.Unlock()
	st := o.states[o.normalize(sessionID)]
	if st == nil {
		return 0
	}
	return len(st.pending)
}

// ModeOf reports the current busy-state lane of a session (idle when unknown).
func (o *Orchestrator) ModeOf(sessionID string) Mode {
	o.mu.Lock()
	defer o.mu.Unlock()
	st := o.states[o.normalize(sessionID)]
	if st == nil {
		return ModeIdle
	}
	return st.mode
}

// normalize maps an empty session id to the default so the state map is stable
// for callers that omit it (the Web main session).
func (o *Orchestrator) normalize(sessionID string) string {
	if sessionID == "" {
		return defaultSessionID
	}
	return sessionID
}

// emitState forwards a state transition to the transport when one is wired.
func (o *Orchestrator) emitState(sessionID, typ string, data map[string]any) {
	if o.emit != nil {
		o.emit(sessionID, typ, data)
	}
}

// notifyPersisted hands a just-persisted message to the transport hook when one is
// wired, so the gateway can broadcast `message_persisted` for cross-end sync.
func (o *Orchestrator) notifyPersisted(msg history.Message) {
	if o.onPersisted != nil {
		o.onPersisted(msg)
	}
}

// deviceIDContextKey is the context key carrying a message's client deviceID
// (design/16 §4.5) from the HTTP/WS boundary into the orchestrator's persistence.
type deviceIDContextKey struct{}

// withDeviceID returns a context carrying deviceID so the orchestrator can tag the
// persisted user message. An empty deviceID is a no-op (no tag is written).
func withDeviceID(ctx context.Context, deviceID string) context.Context {
	return context.WithValue(ctx, deviceIDContextKey{}, deviceID)
}

// deviceIDFrom returns the deviceID carried by ctx, or "" when none was set. It
// returns only well-formed non-empty values; anything else is treated as absent.
func deviceIDFrom(ctx context.Context) string {
	v, _ := ctx.Value(deviceIDContextKey{}).(string)
	return v
}

// boundedTail returns a fresh copy of the last maxTail messages, so continuation
// stays incremental and bounded (design/15 §2.2, §4) and never aliases the input
// slice or another state's tail.
func (o *Orchestrator) boundedTail(msgs []model.Message) []model.Message {
	if len(msgs) <= o.maxTail {
		return append([]model.Message(nil), msgs...)
	}
	out := make([]model.Message, o.maxTail)
	copy(out, msgs[len(msgs)-o.maxTail:])
	return out
}

// messagesEmpty reports whether a turn carries no usable text or image, so the
// orchestrator does not mark a session busy on a vacuous message.
func (o *Orchestrator) messagesEmpty(msgs []model.Message) bool {
	for _, m := range msgs {
		for _, c := range m.Content {
			if c.Type == model.ContentTypeText && c.Text != "" {
				return false
			}
			if c.Type == model.ContentTypeImage && c.ImageURL != "" {
				return false
			}
		}
	}
	return true
}

// concatMessages returns a new slice combining a and b (never reusing a or b
// backing arrays) so the continuation tail is safe to hand to a concurrent run.
func concatMessages(a []model.Message, b ...model.Message) []model.Message {
	out := make([]model.Message, 0, len(a)+len(b))
	out = append(out, a...)
	out = append(out, b...)
	return out
}