agent: replace rollback with proper interruption recording

Instead of discarding history on interruption, record what happened:
- Partial assistant content and tool calls are preserved
- Skipped tool calls get synthetic cancelled results
- Phase-specific notes (stream/tools/request) tell the model what occurred
- Side effects from already-executed tools are no longer hidden from history
This commit is contained in:
Levi Neely 2026-04-10 19:39:02 +02:00
parent 45a094f46e
commit 105b56d746
3 changed files with 87 additions and 29 deletions

View File

@ -334,7 +334,6 @@ func (s *agentCore) Submit(ctx context.Context, input string, handler EventHandl
if err != nil && !errors.Is(err, context.Canceled) && !errors.Is(err, ErrInterrupted) {
handler(Event{Role: "error", Content: err.Error()})
s.session.rollback()
}
handler(Event{Role: "newline"})

View File

@ -48,11 +48,14 @@ func run(ctx context.Context, cfg loopConfig, state state) error {
if err == nil {
break
}
if ctx.Err() != nil {
recordInterruption(state, "request")
return ctx.Err()
}
var rlErr *backend.RateLimitError
if !errors.As(err, &rlErr) || attempt >= maxRateLimitRetries {
return fmt.Errorf("step %d: %w", step, err)
}
// Exponential backoff: 5s, 10s, 20s — unless the server told us exactly.
wait := rlErr.RetryAfter
if wait == 0 {
wait = time.Duration(5<<attempt) * time.Second
@ -81,8 +84,22 @@ func run(ctx context.Context, cfg loopConfig, state state) error {
}
if !done {
return fmt.Errorf("step %d: stream ended without done event", step)
// Stream interrupted. Record whatever we got.
msg := backend.Message{Role: "assistant", Content: content.String(), ToolCalls: toolCalls}
var results []toolResult
for _, tc := range toolCalls {
results = append(results, toolResult{
ToolCallID: tc.ID,
Name: tc.Name,
Content: `{"status":"cancelled","error":"stream interrupted"}`,
IsError: true,
})
}
state.update(msg, results) //nolint:errcheck
recordInterruption(state, "stream")
return fmt.Errorf("step %d: stream interrupted: %w", step, ctx.Err())
}
switch stopReason {
case "stop", "tool_calls", "length", "":
// normal
@ -91,11 +108,12 @@ func run(ctx context.Context, cfg loopConfig, state state) error {
}
totalToolCalls += len(toolCalls)
// Announce and execute tool calls.
// Execute tool calls, handling mid-execution interruption.
msg := backend.Message{Role: "assistant", Content: content.String(), ToolCalls: toolCalls}
var results []toolResult
interrupted := false
for _, tc := range toolCalls {
for i, tc := range toolCalls {
if tc.Name == "" {
results = append(results, toolResult{
ToolCallID: tc.ID,
@ -105,6 +123,22 @@ func run(ctx context.Context, cfg loopConfig, state state) error {
})
continue
}
// Check for cancellation before executing.
if ctx.Err() != nil {
// Fill synthetic results for this and all remaining tool calls.
for _, remaining := range toolCalls[i:] {
results = append(results, toolResult{
ToolCallID: remaining.ID,
Name: remaining.Name,
Content: `{"status":"cancelled","error":"interrupted"}`,
IsError: true,
})
}
interrupted = true
break
}
emit(cfg, Event{Role: "call", Name: tc.Name, Content: string(tc.Arguments)})
var result string
@ -114,6 +148,25 @@ func run(ctx context.Context, cfg loopConfig, state state) error {
if err != nil {
result = fmt.Sprintf("error: %v", err)
isErr = true
// If cancelled during execution, fill remaining and break.
if ctx.Err() != nil {
results = append(results, toolResult{
ToolCallID: tc.ID,
Name: tc.Name,
Content: result,
IsError: true,
})
for _, remaining := range toolCalls[i+1:] {
results = append(results, toolResult{
ToolCallID: remaining.ID,
Name: remaining.Name,
Content: `{"status":"cancelled","error":"interrupted"}`,
IsError: true,
})
}
interrupted = true
break
}
} else {
result = out
}
@ -122,21 +175,26 @@ func run(ctx context.Context, cfg loopConfig, state state) error {
isErr = true
}
results = append(results, toolResult{
ToolCallID: tc.ID,
Name: tc.Name,
Content: result,
IsError: isErr,
})
emit(cfg, Event{Role: "tool", Name: tc.Name, Content: result})
if !interrupted {
results = append(results, toolResult{
ToolCallID: tc.ID,
Name: tc.Name,
Content: result,
IsError: isErr,
})
emit(cfg, Event{Role: "tool", Name: tc.Name, Content: result})
}
}
if err := state.update(msg, results); err != nil {
return fmt.Errorf("step %d update: %w", step, err)
}
if interrupted {
recordInterruption(state, "tools")
return ctx.Err()
}
if len(toolCalls) == 0 {
if err := state.markComplete(); err != nil {
return fmt.Errorf("mark complete: %w", err)
@ -149,7 +207,6 @@ func run(ctx context.Context, cfg loopConfig, state state) error {
}
}
// Surface stall: only when the step limit was hit without completing.
if hitLimit {
emit(cfg, Event{Role: "stalled", Content: "max steps"})
}
@ -157,15 +214,28 @@ func run(ctx context.Context, cfg loopConfig, state state) error {
return nil
}
// recordInterruption appends a note to history so the model knows what happened.
func recordInterruption(state state, phase string) {
var note string
switch phase {
case "stream":
note = "Note: generation was interrupted before completion. Treat the previous assistant output as partial."
case "tools":
note = "Note: tool execution was interrupted before completion. Tool outputs may be missing or cancelled."
case "request":
note = "Note: the previous assistant turn was interrupted before any response was received."
default:
note = "Note: the previous assistant turn was interrupted."
}
state.update(backend.Message{Role: "assistant", Content: note}, nil) //nolint:errcheck
}
func emit(cfg loopConfig, msg Event) {
if cfg.Output != nil {
cfg.Output(msg)
}
}
// retryCountdown emits one "retry" OutputMsg per second, counting down from
// wait, so the UI can display a live countdown. Returns ctx.Err() if the
// context is cancelled before the wait elapses.
func retryCountdown(ctx context.Context, cfg loopConfig, wait time.Duration) error {
deadline := time.Now().Add(wait)
for {

View File

@ -225,17 +225,6 @@ func flattenToolMessages(messages []backend.Message) []backend.Message {
return out
}
// Rollback removes any trailing non-user messages from history, discarding
// an incomplete assistant turn caused by an interruption.
func (s *Session) rollback() {
s.complete = false
i := len(s.messages)
for i > 0 && s.messages[i-1].Role != "user" {
i--
}
s.messages = s.messages[:i]
}
func (s *Session) appendUserMessage(content string) {
s.complete = false
s.messages = append(s.messages, backend.Message{Role: "user", Content: content})