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

package agent

import (
	"context"
	"errors"
	"fmt"
	"os"
	"runtime/debug"
	"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)})
			// Re-submit next FIFO item in a new goroutine so queued
			// prompts aren't orphaned by the panic.
			if next, ok := ag.fifo.Pop(); ok {
				go ag.Submit(ctx, next)
			}
		}
	}()
	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)
	}
	// Drain any remaining FIFO items (covers cases where executeTurn
	// returned "" early: interrupt, toolsrv unavailable, etc.).
	for ctx.Err() == nil {
		next, ok := ag.fifo.Pop()
		if !ok {
			break
		}
		for next != "" && ctx.Err() == nil {
			next = ag.executeTurn(ctx, next)
		}
	}
}

// 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 {
	// Prepend resolved user prompts (active global rules) to the user input.
	if ag.runtime.UserPrompt != "" {
		input = ag.runtime.UserPrompt + "\n\n" + input
	}

	ag.emit(Event{Role: "user", Content: input})

	// Snapshot messages before this turn modifies them. Restored on context
	// overflow so the retry starts from a clean state.
	var preTurnLen int
	if ag.history != nil {
		preTurnLen = len(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)

	// Fetch tools from toolsrv (the sole owner of tool state).
	if ag.runtime.ToolServer != nil {
		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
		ag.runtime.Preamble.Set(SectionTools, RenderTools(ti))
	}

	// Install turn-scoped output handler to intercept events for reply
	// accumulation, state updates, and usage tracking.
	var replyBuf strings.Builder
	origOutput := ag.output
	ag.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 origOutput != nil {
			origOutput(ev)
		}
	}
	defer func() { ag.output = origOutput }()

	// Warn once when context usage crosses 60%; compact at 75%.
	if ag.history != nil {
		tokens := ag.history.estimateTokens()
		if compactLimit := ag.autoCompactLimit(actCtx); compactLimit > 0 && tokens >= compactLimit {
			ag.emit(Event{Role: "info", Content: "auto-compacting context...\n"})
			ag.SetState("compacting")
			if _, err := ag.runCompact(actCtx, "auto"); err != nil {
				ag.log.Error("auto-compact failed: %v", err)
			}
			ag.SetState("thinking")
		} else if warnLimit := ag.autoWarnLimit(actCtx); 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
		}
	}

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

	// Strip cold-zone tool results (one batched LLM call per turn).
	if ag.history != nil {
		ag.history.stripCold(actCtx, ag.runtime.Backend)
	}

	// Inject the preamble system message into history before running.
	if ag.history != nil {
		ag.history.injectPreamble(ag.runtime.PreambleString())
	}

	// Run the turn, retrying once after compaction on context overflow.
	var (
		overflowRetried bool
		err             error
	)
	for {
		err = ag.run(actCtx)
		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 = ag.history.messages[:preTurnLen]
			ag.emit(Event{Role: "info", Content: "context overflow — compacting and retrying...\n"})
			ag.SetState("compacting")
			if _, cerr := ag.runCompact(actCtx, "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)
			// Re-inject preamble — compaction strips it from history.
			ag.history.injectPreamble(ag.runtime.PreambleString())
			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"}
}