diff --git a/pkg/agent/core.go b/pkg/agent/core.go index d5b7fda..6bcdbe1 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -796,6 +796,12 @@ func (s *agent) executeTurn(ctx context.Context, input string, handler EventHand fmt.Sscanf(ev.Content, "%d %d %d %g", &in, &out, &est, &costUSD) s.session.addUsage(backend.Usage{InputTokens: in, OutputTokens: out, CostUSD: costUSD}, est != 0) s.notifyChange() + if maxCostStr := os.Getenv("OLLIE_MAX_SESSION_COST"); maxCostStr != "" { + if limit, ferr := strconv.ParseFloat(maxCostStr, 64); ferr == nil && limit > 0 && s.session.SessionCostUSD >= limit { + handler(infoEvent(fmt.Sprintf("spending cap $%.2f reached — stopping", limit))) + s.Interrupt(ErrInterrupted) + } + } } handler(ev) } @@ -860,12 +866,64 @@ func (s *agent) executeTurn(ctx context.Context, input string, handler EventHand } } + // Spending cap: reject before spending more tokens. + if maxCostStr := os.Getenv("OLLIE_MAX_SESSION_COST"); maxCostStr != "" && s.session != nil { + if limit, err := strconv.ParseFloat(maxCostStr, 64); err == nil && limit > 0 { + if s.session.SessionCostUSD >= limit { + handler(Event{Role: "error", Content: fmt.Sprintf("spending cap $%.2f reached (session total $%.4f)", limit, s.session.SessionCostUSD)}) + s.setState("idle") + actCancel(nil) + s.currentAction.CompareAndSwap(handle, nil) + if snapSession == nil { + s.session = nil + } else { + s.session.messages = snapMessages + } + return "" + } + } + } + if s.session != nil { s.session.resetTurnAccumulators() } - err := run(actCtx, s.cfg, s.session) - actCancel(nil) - s.currentAction.CompareAndSwap(handle, nil) + + // Run the turn, retrying once after compaction on context overflow. + var ( + overflowRetried bool + err error + ) + for { + err = run(actCtx, s.cfg, s.session) + actCancel(nil) + s.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) && s.session != nil { + overflowRetried = true + s.session.messages = snapMessages + emit(s.cfg, Event{Role: "info", Content: "context overflow — compacting and retrying...\n"}) + s.setState("compacting") + if _, cerr := s.runCompact(ctx, "overflow", handler); cerr != nil { + break + } + s.session.appendUserMessage(input) + s.setState("thinking") + s.session.resetTurnAccumulators() + replyBuf.Reset() + actCtx, actCancel = context.WithCancelCause(ctx) + handle = &actionHandle{cancel: actCancel} + s.currentAction.Store(handle) + continue + } + break + } s.mu.Lock() s.reply = replyBuf.String() diff --git a/pkg/backend/backend.go b/pkg/backend/backend.go index 723124f..1af07da 100644 --- a/pkg/backend/backend.go +++ b/pkg/backend/backend.go @@ -9,6 +9,7 @@ import ( "fmt" "io" "net/http" + "strings" "time" ) @@ -78,6 +79,14 @@ type TransientError struct { func (e *TransientError) Error() string { return e.Message } +// ContextOverflowError is returned when the request exceeds the model's +// context window. The caller should compact history and retry. +type ContextOverflowError struct { + Message string +} + +func (e *ContextOverflowError) Error() string { return e.Message } + // GenerationParams controls sampling behaviour for a single ChatStream call. // Zero values mean "use the API default". type GenerationParams struct { @@ -131,7 +140,11 @@ func streamRequest(client *http.Client, req *http.Request, label string, parseFn if resp.StatusCode != http.StatusOK { body, _ := io.ReadAll(resp.Body) resp.Body.Close() - return nil, fmt.Errorf("%s HTTP %d: %s", label, resp.StatusCode, body) + msg := string(body) + if resp.StatusCode == http.StatusBadRequest && isContextOverflow(msg) { + return nil, &ContextOverflowError{Message: fmt.Sprintf("%s HTTP %d: %s", label, resp.StatusCode, msg)} + } + return nil, fmt.Errorf("%s HTTP %d: %s", label, resp.StatusCode, msg) } ch := make(chan StreamEvent, 8) go func() { @@ -141,3 +154,21 @@ func streamRequest(client *http.Client, req *http.Request, label string, parseFn }() return ch, nil } + +// isContextOverflow returns true when a 400 response body indicates the request +// exceeded the model's context window. +func isContextOverflow(body string) bool { + markers := []string{ + "context_length_exceeded", // OpenAI/OpenRouter error code + "prompt is too long", // Anthropic + "Please reduce the length", // OpenAI prose + "too many tokens", + } + lower := strings.ToLower(body) + for _, m := range markers { + if strings.Contains(lower, strings.ToLower(m)) { + return true + } + } + return false +}