diff --git a/pkg/agent/core.go b/pkg/agent/core.go index 40b9188..16b0ce0 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -310,6 +310,15 @@ func NewAgentCore(cfg AgentCoreConfig) Core { a.changeCond = sync.NewCond(&a.changeMu) a.pushSessionEnv() a.pushLockDir() + a.cfg.TurnError = func(ctx context.Context, errType, errMsg string) HookResult { + return a.hooks.Run(ctx, HookTurnError, map[string]string{ + "session_id": a.sessionID, + "cwd": a.CWD(), + "model": a.cfg.Backend.Model(), + "error_type": errType, + "error": errMsg, + }, a.log) + } return a } @@ -497,6 +506,7 @@ func (s *agent) spawnContext(ctx context.Context, handler EventHandler) string { "session_id": s.sessionID, "agent": s.agentName, "cwd": s.CWD(), + "model": s.cfg.Backend.Model(), }, s.log) if result.Warning != "" { handler(infoEvent(result.Warning)) @@ -807,6 +817,8 @@ func (s *agent) executeTurn(ctx context.Context, input string, handler EventHand case "tool": s.setState("thinking") s.auditLog.Debug("result: %s %s", ev.Name, auditTruncate(ev.Content)) + case "limitretry": + s.setState("limitretry") case "error": s.auditLog.Debug("error: %s", ev.Content) } diff --git a/pkg/agent/hooks.go b/pkg/agent/hooks.go index 4f7d233..39eb5e3 100644 --- a/pkg/agent/hooks.go +++ b/pkg/agent/hooks.go @@ -23,6 +23,7 @@ const ( HookPostTool = "postTool" HookPreCompact = "preCompact" HookPostCompact = "postCompact" + HookTurnError = "turnError" ) const defaultHookTimeout = 60 @@ -94,6 +95,18 @@ func runHookCmd(ctx context.Context, name, cmdStr string, payloadJSON []byte, cw cmd.Dir = paths.ExpandHome(cwd) } + // Inject payload map keys as OLLIE_* environment variables so hook + // commands can use $OLLIE_SESSION_ID, $OLLIE_MODEL, etc. without + // parsing the JSON payload on stdin. + var payloadMap map[string]string + if err := json.Unmarshal(payloadJSON, &payloadMap); err == nil { + env := cmd.Environ() + for k, v := range payloadMap { + env = append(env, "OLLIE_"+strings.ToUpper(k)+"="+v) + } + cmd.Env = env + } + var stdout, stderr bytes.Buffer cmd.Stdout = &stdout cmd.Stderr = &stderr diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 5211d55..b5e3167 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -30,6 +30,7 @@ type agentConfig struct { SaveSession func() // called after each state.update(); persists mid-turn progress PreTool func(ctx context.Context, name string, args json.RawMessage) HookResult // called before each tool; exit 2 blocks execution PostTool func(ctx context.Context, name string, args json.RawMessage, result string) HookResult // called after each tool; exit 0 appends, exit 2 replaces result + TurnError func(ctx context.Context, errType, errMsg string) HookResult // called on first backend error; if ran, skips retries } func run(ctx context.Context, cfg agentConfig, state state) error { @@ -60,10 +61,23 @@ func run(ctx context.Context, cfg agentConfig, state state) error { if ctx.Err() != nil { return ctx.Err() } + // On the first error, fire the turnError hook. If it handles + // the error (exit 0), return immediately — the hook is + // responsible for recovery (e.g. switching model and resubmitting). + if attempt == 0 && cfg.TurnError != nil { + errType := classifyError(err) + if r := cfg.TurnError(ctx, errType, err.Error()); r.Ran { + return fmt.Errorf("step %d: %w", step, err) + } + } wait, retryable := transientWait(err, attempt) if !retryable || attempt >= maxTransientRetries { return fmt.Errorf("step %d: %w", step, err) } + var rlErr *backend.RateLimitError + if errors.As(err, &rlErr) { + emit(cfg, Event{Role: "limitretry"}) + } if err := retryCountdown(ctx, cfg, wait); err != nil { return fmt.Errorf("step %d: %w", step, err) } @@ -355,6 +369,28 @@ func emit(cfg agentConfig, msg Event) { } } +// classifyError returns a short string identifying the error type for the +// turnError hook payload. +func classifyError(err error) string { + var rlErr *backend.RateLimitError + if errors.As(err, &rlErr) { + return "rate_limit" + } + var tuErr *backend.ToolUnsupportedError + if errors.As(err, &tuErr) { + return "tool_unsupported" + } + var coErr *backend.ContextOverflowError + if errors.As(err, &coErr) { + return "context_overflow" + } + var tErr *backend.TransientError + if errors.As(err, &tErr) { + return "transient" + } + return "unknown" +} + // transientWait returns the retry wait for a retryable error and whether it is // retryable. Rate limits use longer waits; transient/network errors use shorter. func transientWait(err error, attempt int) (time.Duration, bool) { diff --git a/pkg/agent/loop_turnerror_test.go b/pkg/agent/loop_turnerror_test.go new file mode 100644 index 0000000..90ff90f --- /dev/null +++ b/pkg/agent/loop_turnerror_test.go @@ -0,0 +1,158 @@ +package agent + +import ( + "context" + "sync/atomic" + "testing" + + "ollie/pkg/backend" +) + +// errStream returns a backend respond function that always returns the given error. +func errStream(err error) func(context.Context, []backend.Message, []backend.Tool, backend.GenerationParams) (<-chan backend.StreamEvent, error) { + return func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { + return nil, err + } +} + +// errThenOKStream returns an error on the first call, then a text response. +func errThenOKStream(err error) func(context.Context, []backend.Message, []backend.Tool, backend.GenerationParams) (<-chan backend.StreamEvent, error) { + var n int32 + return func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { + if atomic.AddInt32(&n, 1) == 1 { + return nil, err + } + return textStream("recovered"), nil + } +} + +// TestTurnError_HookInterceptsRateLimit verifies that a turnError hook fired +// on a RateLimitError causes the loop to skip retries and return immediately. +func TestTurnError_HookInterceptsRateLimit(t *testing.T) { + var hookCalls int32 + be := defaultBE() + be.respond = errStream(&backend.RateLimitError{Message: "quota exceeded"}) + + c := newCore(t, be, nil) + c.cfg.TurnError = func(_ context.Context, errType, _ string) HookResult { + atomic.AddInt32(&hookCalls, 1) + if errType != "rate_limit" { + t.Errorf("errType = %q; want rate_limit", errType) + } + return HookResult{Ran: true} + } + + evs := collectEvents(context.Background(), c, "hi") + + if n := atomic.LoadInt32(&hookCalls); n != 1 { + t.Errorf("hook called %d times; want 1 (no retries after hook intercept)", n) + } + errEvs := byRole(evs, "error") + if len(errEvs) == 0 { + t.Error("expected an error event") + } +} + +// TestTurnError_HookInterceptsToolUnsupported verifies the same skip-retry +// behaviour for ToolUnsupportedError. +func TestTurnError_HookInterceptsToolUnsupported(t *testing.T) { + var hookCalls int32 + be := defaultBE() + be.respond = errStream(&backend.ToolUnsupportedError{Message: "model does not support tools"}) + + c := newCore(t, be, nil) + c.cfg.TurnError = func(_ context.Context, errType, _ string) HookResult { + atomic.AddInt32(&hookCalls, 1) + if errType != "tool_unsupported" { + t.Errorf("errType = %q; want tool_unsupported", errType) + } + return HookResult{Ran: true} + } + + collectEvents(context.Background(), c, "hi") + + if n := atomic.LoadInt32(&hookCalls); n != 1 { + t.Errorf("hook called %d times; want 1", n) + } +} + +// TestTurnError_NoHookFallsThrough verifies that when no turnError hook is +// configured, normal retry behaviour proceeds for retryable errors. +func TestTurnError_NoHookFallsThrough(t *testing.T) { + var attempts int32 + be := defaultBE() + be.respond = func(_ context.Context, _ []backend.Message, _ []backend.Tool, _ backend.GenerationParams) (<-chan backend.StreamEvent, error) { + atomic.AddInt32(&attempts, 1) + return nil, &backend.RateLimitError{Message: "slow down"} + } + + c := newCore(t, be, nil) + // TurnError is nil — no hook configured. + + collectEvents(context.Background(), c, "hi") + + // Should have attempted maxTransientRetries+1 = 4 times. + if n := atomic.LoadInt32(&attempts); n != maxTransientRetries+1 { + t.Errorf("attempts = %d; want %d (full retry cycle)", n, maxTransientRetries+1) + } +} + +// TestTurnError_NonRetryableErrorNoHook verifies that a plain (non-retryable) +// error fires the hook once and does not retry. +func TestTurnError_NonRetryableErrorNoHook(t *testing.T) { + var hookCalls int32 + be := defaultBE() + be.respond = errStream(&backend.ToolUnsupportedError{Message: "no tools"}) + + c := newCore(t, be, nil) + c.cfg.TurnError = func(_ context.Context, errType, _ string) HookResult { + atomic.AddInt32(&hookCalls, 1) + return HookResult{Ran: true} + } + + collectEvents(context.Background(), c, "hi") + + if n := atomic.LoadInt32(&hookCalls); n != 1 { + t.Errorf("hook called %d times; want 1", n) + } +} + +// TestTurnError_HookNotRunOnSuccess verifies that the turnError hook is never +// called when the backend succeeds on the first attempt. +func TestTurnError_HookNotRunOnSuccess(t *testing.T) { + var hookCalls int32 + be := defaultBE() + // Default respond returns textStream("ok") — no error. + + c := newCore(t, be, nil) + c.cfg.TurnError = func(_ context.Context, _, _ string) HookResult { + atomic.AddInt32(&hookCalls, 1) + return HookResult{Ran: true} + } + + collectEvents(context.Background(), c, "hi") + + if n := atomic.LoadInt32(&hookCalls); n != 0 { + t.Errorf("hook called %d times on success; want 0", n) + } +} + +// TestTurnError_ClassifyError verifies that classifyError returns the correct +// string for each known error type. +func TestTurnError_ClassifyError(t *testing.T) { + cases := []struct { + err error + want string + }{ + {&backend.RateLimitError{Message: "x"}, "rate_limit"}, + {&backend.ToolUnsupportedError{Message: "x"}, "tool_unsupported"}, + {&backend.ContextOverflowError{Message: "x"}, "context_overflow"}, + {&backend.TransientError{Message: "x"}, "transient"}, + } + for _, tc := range cases { + got := classifyError(tc.err) + if got != tc.want { + t.Errorf("classifyError(%T) = %q; want %q", tc.err, got, tc.want) + } + } +} diff --git a/pkg/backend/backend.go b/pkg/backend/backend.go index 3613d87..4b76af2 100644 --- a/pkg/backend/backend.go +++ b/pkg/backend/backend.go @@ -94,6 +94,14 @@ type ContextOverflowError struct { func (e *ContextOverflowError) Error() string { return e.Message } +// ToolUnsupportedError is returned when the model or backend does not support +// tool/function calling. The request should not be retried with the same model. +type ToolUnsupportedError struct { + Message string +} + +func (e *ToolUnsupportedError) Error() string { return e.Message } + // GenerationParams controls sampling behaviour for a single ChatStream call. // Zero values mean "use the API default". type GenerationParams struct { @@ -151,6 +159,9 @@ func streamRequest(client *http.Client, req *http.Request, label string, parseFn if resp.StatusCode == http.StatusBadRequest && isContextOverflow(msg) { return nil, &ContextOverflowError{Message: fmt.Sprintf("%s HTTP %d: %s", label, resp.StatusCode, msg)} } + if (resp.StatusCode == http.StatusBadRequest || resp.StatusCode == http.StatusUnprocessableEntity) && isToolUnsupported(msg) { + return nil, &ToolUnsupportedError{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) @@ -162,6 +173,27 @@ func streamRequest(client *http.Client, req *http.Request, label string, parseFn return ch, nil } +// isToolUnsupported returns true when a 400/422 response body indicates the +// model does not support tool/function calling. +func isToolUnsupported(body string) bool { + markers := []string{ + "tool use is not supported", + "tools is not supported", + "function calling is not supported", + "does not support tools", + "does not support function", + "tool_use is not enabled", + "not support tool", + } + lower := strings.ToLower(body) + for _, m := range markers { + if strings.Contains(lower, m) { + return true + } + } + return false +} + // isContextOverflow returns true when a 400 response body indicates the request // exceeded the model's context window. func isContextOverflow(body string) bool { diff --git a/pkg/store/memory.go b/pkg/store/memory.go new file mode 100644 index 0000000..878149d --- /dev/null +++ b/pkg/store/memory.go @@ -0,0 +1,106 @@ +package store + +import ( + "context" + "fmt" + "os" + "sync" +) + +// MemoryStore is an in-memory key/value store implementing Store. +// Values are []byte blobs held in RAM; nothing is persisted to disk. +// All operations are safe for concurrent use. +type MemoryStore struct { + mu sync.RWMutex + entries map[string][]byte +} + +func NewMemoryStore() *MemoryStore { + return &MemoryStore{entries: make(map[string][]byte)} +} + +// Get returns the current value for key, or nil if absent. +func (m *MemoryStore) Get(key string) []byte { + m.mu.RLock() + v := m.entries[key] + m.mu.RUnlock() + return v +} + +// Set overwrites the value for key. +func (m *MemoryStore) Set(key string, value []byte) { + m.mu.Lock() + m.entries[key] = value + m.mu.Unlock() +} + +// --- Store interface --- + +func (m *MemoryStore) Stat(name string) (os.FileInfo, error) { + m.mu.RLock() + _, ok := m.entries[name] + sz := int64(len(m.entries[name])) + m.mu.RUnlock() + if !ok { + return nil, fmt.Errorf("%s: not found", name) + } + return &SyntheticFileInfo{Name_: name, Mode_: 0444, Size_: sz}, nil +} + +func (m *MemoryStore) List() ([]os.DirEntry, error) { + m.mu.RLock() + entries := make([]os.DirEntry, 0, len(m.entries)) + for name := range m.entries { + entries = append(entries, FileEntry(name, 0444)) + } + m.mu.RUnlock() + return entries, nil +} + +func (m *MemoryStore) Open(name string) (StoreEntry, error) { + notBlocking := func(context.Context, string) ([]byte, error) { + return nil, fmt.Errorf("blocking read not supported") + } + return &EntryConfig{ + StatFn: func() (os.FileInfo, error) { return m.Stat(name) }, + ReadFn: func() ([]byte, error) { + m.mu.RLock() + v := make([]byte, len(m.entries[name])) + copy(v, m.entries[name]) + m.mu.RUnlock() + return v, nil + }, + WriteFn: func([]byte) error { return fmt.Errorf("%s: read-only", name) }, + BlockingReadFn: notBlocking, + }, nil +} + +func (m *MemoryStore) Create(name string) error { + m.mu.Lock() + m.entries[name] = nil + m.mu.Unlock() + return nil +} + +func (m *MemoryStore) Delete(name string) error { + m.mu.Lock() + _, ok := m.entries[name] + delete(m.entries, name) + m.mu.Unlock() + if !ok { + return fmt.Errorf("%s: not found", name) + } + return nil +} + +func (m *MemoryStore) Rename(oldName, newName string) error { + m.mu.Lock() + defer m.mu.Unlock() + v, ok := m.entries[oldName] + if !ok { + return fmt.Errorf("%s: not found", oldName) + } + m.entries[newName] = v + delete(m.entries, oldName) + return nil +} diff --git a/pkg/store/session.go b/pkg/store/session.go index c1e7e1b..3565b43 100644 --- a/pkg/store/session.go +++ b/pkg/store/session.go @@ -28,6 +28,7 @@ type Session struct { logVers uint32 ChatOffset int plan []byte + prevPrompt []byte // last submitted prompt; overwritten on each new submission } func NewSession(id string, core agent.Core, ctx context.Context, cancel context.CancelFunc) *Session { diff --git a/pkg/store/sessionfile.go b/pkg/store/sessionfile.go index f3acea1..3551419 100644 --- a/pkg/store/sessionfile.go +++ b/pkg/store/sessionfile.go @@ -33,6 +33,7 @@ var SessionFileList = []struct { {"systemprompt", 0444}, {"env", 0444}, {"tail", 0555}, + {"prompt.prev", 0444}, } // SessionFileStore is a RunFileStore for a session directory. @@ -75,6 +76,14 @@ func (h *sessionHelper) fileSpec(name string, mode os.FileMode) FileSpec { h.sess.mu.RUnlock() return int64(n) } + case "prompt.prev": + fs.Read = func() ([]byte, error) { + h.sess.mu.RLock() + data := make([]byte, len(h.sess.prevPrompt)) + copy(data, h.sess.prevPrompt) + h.sess.mu.RUnlock() + return data, nil + } case "chat": fs.Read = func() ([]byte, error) { h.sess.mu.RLock() @@ -127,6 +136,9 @@ func (h *sessionHelper) fileSpec(name string, mode os.FileMode) FileSpec { if input == "" { return nil } + h.sess.mu.Lock() + h.sess.prevPrompt = []byte(input) + h.sess.mu.Unlock() pub := h.makePublish() go func() { h.sess.Core.Submit(h.sess.Ctx, input, pub)