ollie/cmd/olliesrv/internal/agent/turn.go

package agent

import (
	"context"
	"errors"
	"fmt"
	"os"
	"runtime/debug"
	"strconv"
	"strings"

	"ollie/cmd/olliesrv/internal/backend"
	"ollie/toolsrv"
)

// Submit processes one line of user input: it starts an agent turn that streams
// events to the bus. If a turn is already in progress the prompt is queued as
// an in-stream interruption instead.
//
// Slash commands are NOT handled here — the caller (Session.Submit) dispatches
// those before delegating to Agent.Submit.
//
// Continuations (post-turn hook context, unconsumed inject, FIFO drain) are
// handled via an explicit loop rather than recursion to avoid stack growth.
func (ag *Agent) Submit(ctx context.Context, input string) {
	defer func() {
		if r := recover(); r != nil {
			ag.log.Error("panic: %v\n%s", r, debug.Stack())
			if a := ag.currentAction.Swap(nil); a != nil {
				a.cancel(fmt.Errorf("%v", r))
			}
			ag.SetState("idle")
			ag.emit(Event{Role: "error", Content: fmt.Sprintf("%v", r)})
		}
	}()
	ag.log.Debug("Agent.Submit() input_len=%d running=%v", len(input), ag.IsRunning())
	if input == "" {
		return
	}

	if ag.IsRunning() {
		// While running, slash commands are still dispatched immediately.
		if ag.HandleCommand(ctx, input) {
			return
		}
		ag.fifo.Push(input)
		return
	}

	// Serialize turns so that e.g. a /compact arriving via ctl cannot race
	// with an executeTurn arriving via prompt.
	ag.submitMu.Lock()
	defer ag.submitMu.Unlock()

	if ag.HandleCommand(ctx, input) {
		return
	}

	if ag.IsRunning() {
		ag.fifo.Push(input)
		return
	}

	for input != "" && ctx.Err() == nil {
		input = ag.executeTurn(ctx, input)
	}
}

// executeTurn runs a single agent turn and returns the next prompt to execute,
// or "" if there is nothing more to do.
func (ag *Agent) executeTurn(ctx context.Context, input string) string {
	ag.emit(Event{Role: "user", Content: input})

	// Snapshot session state before this turn modifies it. Restored on failure
	// so the session is clean for the next attempt.
	snapSession := ag.history
	var snapMessages []backend.Message
	if ag.history != nil {
		snapMessages = cloneMessages(ag.history.messages)
	}

	if ag.history == nil {
		for _, msg := range ag.startupMessages {
			ag.log.Debug("startup: %s", msg)
			ag.emit(infoEvent(msg))
		}
		ag.startupMessages = nil
		ag.history = newHistory(input)
		if sc := ag.spawnContext(ctx); sc != "" {
			ag.history.appendUserMessage(sc)
		}
		ag.history.appendUserMessage(input)
	} else {
		ag.history.appendUserMessage(input)
	}

	actCtx, actCancel := context.WithCancelCause(ctx)
	handle := &actionHandle{cancel: actCancel}
	ag.currentAction.Store(handle)
	ag.SetState("thinking")

	ag.log.Debug("turn: start input=%s session=%s", auditTruncate(input), ag.sessionID)

	// Refresh tool definitions only when the registry has changed (tool_load/unload).
	// This avoids ListTools IPC + allocation on every turn in the common case.
	if ag.toolsNeedRefresh() {
		ti, listErr := ag.runtime.ToolServer.ListTools()
		if listErr != nil {
			ag.log.Error("tool server unavailable: %v", listErr)
			ag.emit(Event{Role: "error", Content: fmt.Sprintf("tool server unavailable: %v", listErr)})
			ag.SetState("idle")
			return ""
		}
		ag.runtime.Tools = toolInfosToBackend(ti)
		meta := make(map[string]toolsrv.ToolInfo, len(ti))
		for _, info := range ti {
			meta[info.Name] = info
		}
		ag.runtime.ToolMeta = meta
	}

	turnCtx := TurnCtx{
		Output:      ag.output,
		Save:        ag.save,
		ResultCache: &ag.resultCache,
		BgTracker:   &ag.bgProcs,
		ToolServer:  ag.runtime.ToolServer,
	}

	var replyBuf strings.Builder
	turnCtx.Output = func(ev Event) {
		switch ev.Role {
		case "assistant":
			replyBuf.WriteString(ev.Content)
		case "call":
			ag.SetState("calling: " + ev.Name)
			ag.log.Debug("call: %s %s", ev.Name, auditTruncate(string(ev.Content)))
		case "tool":
			ag.log.Debug("result: %s %s", ev.Name, auditTruncate(ev.Content))
		case "state":
			ag.SetState(ev.Content)
		case "limitretry":
			ag.SetState("limitretry")
		case "error":
			ag.log.Debug("error: %s", ev.Content)
		}
		if ev.Role == "usage" && ag.history != nil {
			var in, out, est, cached, creation int
			var costUSD float64
			fmt.Sscanf(ev.Content, "%d %d %d %g %d %d", &in, &out, &est, &costUSD, &cached, &creation)
			ag.history.addUsage(backend.Usage{
				InputTokens:         in,
				CachedInputTokens:   cached,
				CacheCreationTokens: creation,
				OutputTokens:        out,
				CostUSD:             costUSD,
			}, est != 0)
			ag.notifyChange()
			if maxCostStr := os.Getenv("OLLIE_MAX_SESSION_COST"); maxCostStr != "" {
				if limit, ferr := strconv.ParseFloat(maxCostStr, 64); ferr == nil && limit > 0 && ag.history.SessionCostUSD >= limit {
					ag.emit(infoEvent(fmt.Sprintf("spending cap $%.2f reached — stopping", limit)))
					ag.Interrupt(ErrInterrupted)
				}
			}
		}
		ag.emit(ev)
	}
	turnCtx.PopInject = func() string {
		if p := ag.pendingInject.Swap(nil); p != nil {
			return *p
		}
		return ""
	}
	turnCtx.IncrToolCallCount = func() int64 {
		return ag.toolCallCount.Add(1)
	}
	turnCtx.Save = ag.save
	turnCtx.ResultCache = &ag.resultCache
	turnCtx.AutoCompact = func(ctx context.Context) {
		if ctx.Err() != nil || ag.history == nil {
			return
		}
		limit := ag.autoCompactLimit(ctx)
		if limit <= 0 || ag.history.estimateTokens() < limit {
			return
		}
		ag.emit(Event{Role: "info", Content: "auto-compacting context...\n"})
		ag.SetState("compacting")
		if _, err := ag.runCompact(ctx, "auto"); err != nil {
			panic(fmt.Sprintf("mid-turn auto-compact: %v", err))
		}
		ag.SetState("thinking")
	}

	// Warn once when context usage crosses 60%; compact at 75%.
	if ag.history != nil {
		tokens := ag.history.estimateTokens()
		if compactLimit := ag.autoCompactLimit(ctx); compactLimit > 0 && tokens >= compactLimit {
			ag.emit(Event{Role: "info", Content: "auto-compacting context...\n"})
			ag.SetState("compacting")
			if _, err := ag.runCompact(ctx, "auto"); err != nil {
				panic(fmt.Sprintf("auto-compact: %v", err))
			}
			ag.SetState("thinking")
		} else if warnLimit := ag.autoWarnLimit(ctx); warnLimit > 0 && tokens >= warnLimit && !ag.warnedContext {
			ctxLen := ag.runtime.Backend.ContextLength(ctx)
			if ctxLen <= 0 {
				ctxLen = defaultContextLength
			}
			pct := tokens * 100 / ctxLen
			ag.emit(Event{Role: "info", Content: fmt.Sprintf("context at %d%% — will auto-compact at 75%%\n", pct)})
			ag.warnedContext = true
		}
	}

	// Spending cap: reject before spending more tokens.
	if maxCostStr := os.Getenv("OLLIE_MAX_SESSION_COST"); maxCostStr != "" && ag.history != nil {
		if limit, err := strconv.ParseFloat(maxCostStr, 64); err == nil && limit > 0 {
			if ag.history.SessionCostUSD >= limit {
				ag.emit(Event{Role: "error", Content: fmt.Sprintf("spending cap $%.2f reached (session total $%.4f)", limit, ag.history.SessionCostUSD)})
				ag.SetState("idle")
				actCancel(nil)
				ag.currentAction.CompareAndSwap(handle, nil)
				if snapSession == nil {
					ag.history = nil
				} else {
					ag.history.messages = snapMessages
				}
				return ""
			}
		}
	}

	if ag.history != nil {
		ag.history.resetTurnAccumulators()
	}

	// Run the turn, retrying once after compaction on context overflow.
	var (
		overflowRetried bool
		err             error
	)
	for {
		turnCtx.Context = actCtx
		err = run(ag.runtime, turnCtx, ag.history)
		actCancel(nil)
		ag.currentAction.CompareAndSwap(handle, nil)

		if err == nil {
			break
		}
		if errors.Is(err, context.Canceled) || errors.Is(err, ErrInterrupted) {
			break
		}
		var ctxErr *backend.ContextOverflowError
		if !overflowRetried && errors.As(err, &ctxErr) && ag.history != nil {
			overflowRetried = true
			ag.history.messages = snapMessages
			ag.emit(Event{Role: "info", Content: "context overflow — compacting and retrying...\n"})
			ag.SetState("compacting")
			if _, cerr := ag.runCompact(ctx, "overflow"); cerr != nil {
				break
			}
			ag.history.appendUserMessage(input)
			ag.SetState("thinking")
			ag.history.resetTurnAccumulators()
			replyBuf.Reset()
			actCtx, actCancel = context.WithCancelCause(ctx)
			handle = &actionHandle{cancel: actCancel}
			ag.currentAction.Store(handle)
			continue
		}
		break
	}

	ag.setReply(replyBuf.String())
	replyBuf.Reset()
	ag.SetState("idle")
	ag.flush()

	if err != nil {
		// Keep completed work — only remove cancelled tool results.
		if ag.history != nil {
			ag.history.removeCancelledToolResults()
		}
		if errors.Is(err, context.Canceled) || errors.Is(err, ErrInterrupted) {
			ag.log.Debug("turn: interrupted session=%s", ag.sessionID)
			ag.save()
			return ""
		}
		ag.emit(Event{Role: "error", Content: err.Error()})
		// Drain one FIFO item.
		if next, ok := ag.fifo.Pop(); ok {
			return next
		}
		return ""
	}

	if ag.history != nil {
		ag.history.recordTurnCost(ag.runtime.Backend.Model())
		if ag.history.LastTurnCostUSD > 0 {
			ag.emit(Event{Role: "info", Content: fmt.Sprintf("costLast=$%.4f\n", ag.history.LastTurnCostUSD)})
		}
		ag.log.Debug("turn: end reply=%s cost=$%.4f session_total=$%.4f session=%s",
			auditTruncate(ag.Reply()), ag.history.LastTurnCostUSD, ag.history.SessionCostUSD, ag.sessionID)
		ag.notifyChange()
	}
	ag.save()

	// Inject that was pending but never consumed (text-only response with no
	// tool calls) — treat it as the next user message.
	if p := ag.pendingInject.Swap(nil); p != nil {
		return *p
	}

	// Drain one item from the FIFO; the outer loop handles the rest.
	if next, ok := ag.fifo.Pop(); ok {
		return next
	}

	return ""
}

// autoCompactLimit returns the token threshold for auto-compaction (75%).
func (ag *Agent) autoCompactLimit(ctx context.Context) int {
	ctxLen := ag.runtime.Backend.ContextLength(ctx)
	if ctxLen <= 0 {
		ctxLen = defaultContextLength
	}
	return ctxLen * 3 / 4
}

// autoWarnLimit returns the token threshold for a context-usage warning (60%).
func (ag *Agent) autoWarnLimit(ctx context.Context) int {
	ctxLen := ag.runtime.Backend.ContextLength(ctx)
	if ctxLen <= 0 {
		ctxLen = defaultContextLength
	}
	return ctxLen * 3 / 5
}

// spawnContext assembles the agent context injected at each session refresh
// point (session start, post-clear, post-compaction).
func (ag *Agent) spawnContext(ctx context.Context) string {
	var parts []string
	// Inject AGENTS.md from the working directory if it exists.
	if cwd := ag.effectiveCwd(); cwd != "" {
		if data, err := os.ReadFile(cwd + "/AGENTS.md"); err == nil && len(data) > 0 {
			parts = append(parts, string(data))
		}
	}
	return strings.Join(parts, "\n\n---\n\n")
}

// runCompact executes a full compaction cycle: compact, spawn-context
// re-injection. Returns (n compacted, error). Returns (0, nil) if
// there was nothing to compact. Caller manages setState.
func (ag *Agent) runCompact(ctx context.Context, trigger string) (int, error) {
	// Use a cheaper model for compaction if configured.
	compactModel := resolveCompactionModel(ag.runtime.CompactionModel, ag.runtime.Backend)
	origModel := ag.runtime.Backend.Model()
	if compactModel != "" && compactModel != origModel {
		ag.runtime.Backend.SetModel(compactModel)
		defer ag.runtime.Backend.SetModel(origModel)
	}
	n, _, err := ag.history.compact(ctx, ag.runtime.Backend)
	if err != nil {
		return 0, err
	}
	if n > 0 {
		ag.log.Debug("compact: removed %d messages trigger=%s session=%s", n, trigger, ag.sessionID)
		ag.warnedContext = false
		if sc := ag.spawnContext(ctx); sc != "" {
			ag.history.appendUserMessage(sc)
		}
	}
	return n, nil
}

// ErrInterrupted is returned when the user cancels an agent turn (Ctrl-C).
var ErrInterrupted = errors.New("interrupted")

// defaultContextLength is used when the backend cannot report the model's
// actual context window. 128k tokens is a safe default for modern models.
const defaultContextLength = 128000

// infoEvent wraps a plain-text message as an info Event.
func infoEvent(text string) Event {
	return Event{Role: "info", Content: text + "\n"}
}