Agent-Native/runtime/session_agent.go

package runtime

import (
	"context"
	"sync"

	"agent-native/runtime/llm"
	"agent-native/runtime/session"
)

// Delivery selects how a queued steer or follow-up is handed to the loop.
type Delivery string

const (
	DeliverOne Delivery = "one"
	DeliverAll Delivery = "all"
)

// AgentSession is the headless session object: queues, abort, and
// mid-run model changes around Run.
type AgentSession struct {
	cfg Config

	mu         sync.Mutex
	steer      []llm.Message
	follow     []llm.Message
	steerMode  Delivery
	followMode Delivery
	model      string
	thinking   string
	cancel     context.CancelFunc
	runID      int
}

// NewAgentSession copies cfg. Queue modes default to one-at-a-time.
func NewAgentSession(cfg Config) *AgentSession {
	if cfg.Model == "" {
		cfg.Model = "default"
	}
	return &AgentSession{
		cfg:        cfg,
		steerMode:  DeliverOne,
		followMode: DeliverOne,
		model:      cfg.Model,
		thinking:   cfg.ThinkingLevel,
	}
}

// SetQueueMode sets the delivery mode for steer and follow-up.
func (a *AgentSession) SetQueueMode(steer, follow Delivery) {
	a.mu.Lock()
	defer a.mu.Unlock()
	if steer != "" {
		a.steerMode = steer
	}
	if follow != "" {
		a.followMode = follow
	}
}

// Steer queues a message to be delivered after the current tool batch.
func (a *AgentSession) Steer(msg llm.Message) {
	a.mu.Lock()
	a.steer = append(a.steer, msg)
	a.mu.Unlock()
	a.emitQueue("steer")
}

// FollowUp queues a message for after the assistant stops.
func (a *AgentSession) FollowUp(msg llm.Message) {
	a.mu.Lock()
	a.follow = append(a.follow, msg)
	a.mu.Unlock()
	a.emitQueue("follow_up")
}

// ClearQueue drops both queues.
func (a *AgentSession) ClearQueue() {
	a.mu.Lock()
	a.steer = nil
	a.follow = nil
	a.mu.Unlock()
	a.emitQueue("clear")
}

// Abort cancels the in-flight Prompt.
func (a *AgentSession) Abort() {
	a.mu.Lock()
	cancel := a.cancel
	a.mu.Unlock()
	if cancel != nil {
		cancel()
	}
}

// SetModel changes the model and thinking level used on the next turn.
// A session store, when present, records the change.
func (a *AgentSession) SetModel(model, thinking string) error {
	if thinking != "" {
		norm, err := llm.NormalizeThinking(thinking)
		if err != nil {
			return err
		}
		thinking = norm
	}
	a.mu.Lock()
	if model != "" {
		a.model = model
	}
	if thinking != "" {
		a.thinking = thinking
	}
	store := a.cfg.Session
	a.mu.Unlock()
	if store == nil {
		return nil
	}
	if model != "" {
		if _, err := store.AppendEntry(session.Entry{Type: session.TypeModelChange, Model: model}); err != nil {
			return err
		}
	}
	if thinking != "" {
		if _, err := store.AppendEntry(session.Entry{Type: session.TypeThinkingLevelChange, ThinkingLevel: thinking}); err != nil {
			return err
		}
	}
	return nil
}

// Prompt runs one user message until the agent settles.
func (a *AgentSession) Prompt(ctx context.Context, user llm.Message) (Result, error) {
	ctx, cancel := context.WithCancel(ctx)
	a.mu.Lock()
	a.runID++
	id := a.runID
	a.cancel = cancel
	cfg := a.cfg
	a.mu.Unlock()
	defer func() {
		cancel()
		a.mu.Lock()
		if a.runID == id {
			a.cancel = nil
		}
		a.mu.Unlock()
	}()
	cfg.GetSteering = func() []llm.Message { return a.take("steer") }
	cfg.GetFollowUp = func() []llm.Message { return a.take("follow") }
	cfg.Resolve = func() (string, string) {
		a.mu.Lock()
		defer a.mu.Unlock()
		return a.model, a.thinking
	}
	return Run(ctx, cfg, user)
}

// Compact appends a compaction entry and emits compaction events.
func (a *AgentSession) Compact(keep int, summarize func([]llm.Message) (string, error)) (bool, error) {
	if a.cfg.Session == nil {
		return false, nil
	}
	emit(a.cfg, Event{Type: "compaction_start"})
	done, err := a.cfg.Session.Compact(keep, summarize)
	text := "noop"
	if done {
		text = "compacted"
	}
	if err != nil {
		text = err.Error()
	}
	emit(a.cfg, Event{Type: "compaction_end", Text: text})
	return done, err
}

func (a *AgentSession) take(kind string) []llm.Message {
	a.mu.Lock()
	defer a.mu.Unlock()
	q := &a.steer
	mode := a.steerMode
	if kind == "follow" {
		q = &a.follow
		mode = a.followMode
	}
	if len(*q) == 0 {
		return nil
	}
	if mode == DeliverAll {
		out := append([]llm.Message(nil), (*q)...)
		*q = nil
		return out
	}
	out := []llm.Message{(*q)[0]}
	*q = (*q)[1:]
	return out
}

func (a *AgentSession) emitQueue(text string) {
	emit(a.cfg, Event{Type: "queue_update", Text: text})
}