507 lines
16 KiB
Go
507 lines
16 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"runtime/debug"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"ollie/backend"
|
|
)
|
|
|
|
// 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})
|
|
|
|
hookResult := ag.runtime.Hooks.Run(ctx, HookPreTurn, map[string]string{
|
|
"session_id": ag.sessionID,
|
|
"cwd": ag.effectiveCwd(),
|
|
"prompt": input,
|
|
}, ag.log)
|
|
if hookResult.Blocked {
|
|
ag.emit(infoEvent("hook blocked prompt"))
|
|
return ""
|
|
}
|
|
if hookResult.Warning != "" {
|
|
ag.emit(infoEvent(hookResult.Warning))
|
|
}
|
|
if sum := hookResult.Summary(); sum != "" {
|
|
ag.emit(infoEvent("preTurn: " + sum))
|
|
}
|
|
if hookResult.Context != "" {
|
|
input += "\n" + hookResult.Context
|
|
}
|
|
|
|
// 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.auditLog.Debug("turn: start input=%s session=%s", auditTruncate(input), ag.sessionID)
|
|
|
|
// Build per-turn agentConfig from the current runtime.
|
|
ag.cfg = agentConfig{
|
|
Backend: ag.runtime.Backend,
|
|
preamble: ag.runtime.Preamble,
|
|
Tools: ag.runtime.Tools,
|
|
Exec: ag.runtime.Exec,
|
|
ClassifyTool: ag.runtime.ClassifyTool,
|
|
ClassifyTier: ag.runtime.ClassifyTier,
|
|
GenerationParams: ag.runtime.GenParams,
|
|
MaxSteps: ag.runtime.MaxSteps,
|
|
ReadPlanStep: ag.readPlanStep,
|
|
TurnError: ag.turnError,
|
|
}
|
|
|
|
var replyBuf strings.Builder
|
|
ag.cfg.Output = func(ev Event) {
|
|
switch ev.Role {
|
|
case "assistant":
|
|
replyBuf.WriteString(ev.Content)
|
|
case "call":
|
|
ag.SetState("calling: " + ev.Name)
|
|
ag.auditLog.Debug("call: %s %s", ev.Name, auditTruncate(string(ev.Content)))
|
|
case "tool":
|
|
ag.auditLog.Debug("result: %s %s", ev.Name, auditTruncate(ev.Content))
|
|
case "state":
|
|
ag.SetState(ev.Content)
|
|
case "limitretry":
|
|
ag.SetState("limitretry")
|
|
case "error":
|
|
ag.auditLog.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)
|
|
}
|
|
ag.cfg.PopInject = func() string {
|
|
if p := ag.pendingInject.Swap(nil); p != nil {
|
|
return *p
|
|
}
|
|
return ""
|
|
}
|
|
ag.cfg.PreTool = func(ctx context.Context, name string, args json.RawMessage) HookResult {
|
|
return ag.runtime.Hooks.Run(ctx, HookPreTool, map[string]string{
|
|
"session_id": ag.sessionID,
|
|
"cwd": ag.effectiveCwd(),
|
|
"tool": name,
|
|
"args": string(args),
|
|
}, ag.log)
|
|
}
|
|
ag.cfg.PostTool = func(ctx context.Context, name string, args json.RawMessage, result string) HookResult {
|
|
return ag.runtime.Hooks.Run(ctx, HookPostTool, map[string]string{
|
|
"session_id": ag.sessionID,
|
|
"cwd": ag.effectiveCwd(),
|
|
"tool": name,
|
|
"args": string(args),
|
|
"result": result,
|
|
}, ag.log)
|
|
}
|
|
ag.cfg.IncrToolCallCount = func() int64 {
|
|
return ag.toolCallCount.Add(1)
|
|
}
|
|
ag.cfg.SaveSession = func() { ag.saveSession() }
|
|
ag.cfg.ResultCache = &ag.resultCache
|
|
ag.cfg.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.cfg.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 {
|
|
err = run(actCtx, ag.cfg, 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")
|
|
if ag.flushSave != nil {
|
|
ag.flushSave()
|
|
}
|
|
|
|
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.auditLog.Debug("turn: interrupted session=%s", ag.sessionID)
|
|
ag.saveSession()
|
|
return ""
|
|
}
|
|
ag.emit(Event{Role: "error", Content: err.Error()})
|
|
// Drain one FIFO item — the turnError hook may have queued a recovery prompt.
|
|
if next, ok := ag.fifo.Pop(); ok {
|
|
return next
|
|
}
|
|
return ""
|
|
}
|
|
|
|
stopResult := ag.runtime.Hooks.Run(ctx, HookPostTurn, map[string]string{
|
|
"session_id": ag.sessionID,
|
|
"cwd": ag.effectiveCwd(),
|
|
}, ag.log)
|
|
if stopResult.Warning != "" {
|
|
ag.emit(infoEvent(stopResult.Warning))
|
|
}
|
|
if sum := stopResult.Summary(); sum != "" {
|
|
ag.emit(infoEvent("postTurn: " + sum))
|
|
}
|
|
if !stopResult.Blocked && stopResult.Context != "" && ag.history != nil {
|
|
ag.history.appendUserMessage(stopResult.Context)
|
|
}
|
|
|
|
if ag.history != nil {
|
|
ag.history.recordTurnCost(ag.cfg.Backend.Model())
|
|
appendUsageLog(ag.sessionID, ag.cfg.Backend.Name(), ag.cfg.Backend.Model(), ag.history)
|
|
if ag.history.LastTurnCostUSD > 0 {
|
|
ag.emit(Event{Role: "info", Content: fmt.Sprintf("costLast=$%.4f\n", ag.history.LastTurnCostUSD)})
|
|
}
|
|
ag.auditLog.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.saveSession()
|
|
|
|
// Post-turn hook said "continue" — its context becomes the next prompt.
|
|
if stopResult.Blocked && stopResult.Context != "" {
|
|
return stopResult.Context
|
|
}
|
|
|
|
// 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 ""
|
|
}
|
|
|
|
// turnError is called on the first backend error of a turn. If the hook handles
|
|
// the error (exit 0), the turn exits immediately.
|
|
// This is the default implementation; it can be overridden via the turnError field.
|
|
func (ag *Agent) defaultTurnError(ctx context.Context, errType, errMsg string) HookResult {
|
|
hookCtx, cancel := context.WithTimeout(context.Background(), time.Duration(hookTimeout)*time.Second)
|
|
defer cancel()
|
|
return ag.runtime.Hooks.Run(hookCtx, HookTurnError, map[string]string{
|
|
"session_id": ag.sessionID,
|
|
"cwd": ag.effectiveCwd(),
|
|
"model": ag.runtime.Backend.Model(),
|
|
"error_type": errType,
|
|
"error": errMsg,
|
|
}, ag.log)
|
|
}
|
|
|
|
// 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). It combines the
|
|
// agent-specific prompt with any agentSpawn hook output.
|
|
func (ag *Agent) spawnContext(ctx context.Context) string {
|
|
result := ag.runtime.Hooks.Run(ctx, HookAgentSpawn, map[string]string{
|
|
"session_id": ag.sessionID,
|
|
"agent": ag.agentName,
|
|
"cwd": ag.effectiveCwd(),
|
|
"model": ag.runtime.Backend.Model(),
|
|
}, ag.log)
|
|
if result.Warning != "" {
|
|
ag.emit(infoEvent(result.Warning))
|
|
}
|
|
if sum := result.Summary(); sum != "" {
|
|
ag.emit(infoEvent("agentSpawn: " + sum))
|
|
}
|
|
var parts []string
|
|
if result.Context != "" {
|
|
parts = append(parts, result.Context)
|
|
}
|
|
// 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: pre-hook, compact, spawn-context
|
|
// re-injection, post-hook. Returns (n compacted, error). Returns (0, nil) if
|
|
// the pre-hook blocked or there was nothing to compact. Caller manages setState.
|
|
func (ag *Agent) runCompact(ctx context.Context, trigger string) (int, error) {
|
|
payload := map[string]string{"session_id": ag.sessionID, "trigger": trigger, "cwd": ag.effectiveCwd()}
|
|
pre := ag.runtime.Hooks.Run(ctx, HookPreCompact, payload, ag.log)
|
|
if pre.Warning != "" {
|
|
ag.emit(infoEvent(pre.Warning))
|
|
}
|
|
if sum := pre.Summary(); sum != "" {
|
|
ag.emit(infoEvent("preCompact: " + sum))
|
|
}
|
|
if pre.Blocked {
|
|
ag.emit(infoEvent("compact cancelled by hook"))
|
|
return 0, nil
|
|
}
|
|
if pre.Context != "" {
|
|
ag.history.appendUserMessage(pre.Context)
|
|
}
|
|
// 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.auditLog.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)
|
|
}
|
|
}
|
|
post := ag.runtime.Hooks.Run(ctx, HookPostCompact, payload, ag.log)
|
|
if post.Warning != "" {
|
|
ag.emit(infoEvent(post.Warning))
|
|
}
|
|
if sum := post.Summary(); sum != "" {
|
|
ag.emit(infoEvent("postCompact: " + sum))
|
|
}
|
|
if post.Context != "" {
|
|
ag.history.appendUserMessage(post.Context)
|
|
}
|
|
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"}
|
|
}
|
|
|
|
// firstSentence returns the first sentence of s (up to the first period or
|
|
// newline), trimmed. Falls back to s truncated at 80 chars if no sentence end found.
|
|
func firstSentence(s string) string {
|
|
for i, r := range s {
|
|
if r == '.' || r == '\n' {
|
|
return strings.TrimSpace(s[:i+1])
|
|
}
|
|
}
|
|
if len(s) > 80 {
|
|
return s[:77] + "..."
|
|
}
|
|
return s
|
|
}
|