Skip to content

Persist turn closeout messages during shutdown #261

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
anabellabuckvar wants to merge 1 commit into from
Closed

Persist turn closeout messages during shutdown #261

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 32 additions & 9 deletions loop/loop.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,14 @@ import (
// heartbeats/keepalives (which reset the idle timer) from hanging a turn
// indefinitely. It is deliberately far larger than the idle window so that
// genuinely long, steadily-streaming turns are unaffected.
const maxTurnDuration = 15 * time.Minute
const (
maxTurnDuration = 15 * time.Minute
durableMessageTimeout = 5 * time.Second
)

// ErrShutdown marks loop cancellation caused by graceful server shutdown.
// Final tool results and closeout errors use a detached context for this cause.
var ErrShutdown = errors.New("loop stopped by server shutdown")

// MessageRecordFunc is called to record new messages to persistent storage.
// otherUsage carries the usage of indirect LLM calls affiliated with the
Expand Down Expand Up @@ -138,6 +145,21 @@ func NewLoop(config Config) *Loop {
}
}

// recordMessageDurably detaches a bounded write from the turn context while
// preserving its values.
func (l *Loop) recordMessageDurably(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error {
durableCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), durableMessageTimeout)
defer cancel()
return l.recordMessage(durableCtx, message, usage, otherUsage)
}

func (l *Loop) recordToolResult(ctx context.Context, message llm.Message, otherUsage []llm.PurposedUsage) error {
if ctx.Err() == nil || (errors.Is(ctx.Err(), context.Canceled) && !errors.Is(context.Cause(ctx), ErrShutdown)) {
return l.recordMessage(ctx, message, llm.Usage{}, otherUsage)
}
return l.recordMessageDurably(ctx, message, llm.Usage{}, otherUsage)
}

// Retry signals the loop to re-attempt the next LLM request without queueing
// a new user message. The loop's in-memory history is unchanged (failed
// requests don't append anything to history, and error messages are persisted
Expand Down Expand Up @@ -267,6 +289,9 @@ func (l *Loop) Go(ctx context.Context) error {
l.logger.Debug("processing queued messages", "count", 1)
if err := l.processLLMRequest(ctx); err != nil {
l.logger.Error("failed to process LLM request", "error", err)
if ctx.Err() != nil {
return ctx.Err()
}
time.Sleep(time.Second) // Wait before retrying
continue
}
Expand Down Expand Up @@ -458,17 +483,15 @@ func (l *Loop) processLLMRequest(ctx context.Context) error {
// duplicate that — and, because ctx is already dead, the write
// failed anyway ("failed to create message: Tx: context canceled"),
// producing a scary log line on every user cancel. Skip it.
if errors.Is(ctx.Err(), context.Canceled) {
if errors.Is(ctx.Err(), context.Canceled) && !errors.Is(context.Cause(ctx), ErrShutdown) {
l.logger.Info("LLM request aborted by loop cancellation", "error", err)
return fmt.Errorf("LLM request failed: %w", err)
}
// Record the error as a message so it can be displayed in the UI.
// EndOfTurn must be true so the agent working state is properly
// updated. WithoutCancel: this row's MarkAgentDone is what clears
// the persisted agent_working flag; losing the write to a context
// that expired (e.g. the 12-hour conversation-loop ceiling) or was
// cancelled between the failure above and this write would wedge
// the conversation in "Agent working..." forever.
// updated. recordMessageDurably detaches this write when the turn
// expired or graceful shutdown cancelled it; losing the row would
// wedge the conversation in "Agent working..." forever.
errorMessage := llm.Message{
Role: llm.MessageRoleAssistant,
Content: []llm.Content{
Expand All @@ -481,7 +504,7 @@ func (l *Loop) processLLMRequest(ctx context.Context) error {
ErrorType: llm.ErrorTypeLLMRequest,
ErrorRetryable: IsRetryableLLMError(err),
}
if recordErr := l.recordMessage(context.WithoutCancel(ctx), errorMessage, llm.Usage{}, nil); recordErr != nil {
if recordErr := l.recordMessageDurably(ctx, errorMessage, llm.Usage{}, nil); recordErr != nil {
l.logger.Error("failed to record error message", "error", recordErr)
}
return fmt.Errorf("LLM request failed: %w", err)
Expand Down Expand Up @@ -870,7 +893,7 @@ func (l *Loop) executeToolCalls(ctx context.Context, content []llm.Content) erro
l.mu.Unlock()

// Record tool result message
if err := l.recordMessage(ctx, toolMessage, llm.Usage{}, otherUsage.Take()); err != nil {
if err := l.recordToolResult(ctx, toolMessage, otherUsage.Take()); err != nil {
l.logger.Error("failed to record tool result message", "error", err)
}
}
Expand Down
117 changes: 117 additions & 0 deletions loop/loop_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package loop
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
Expand Down Expand Up @@ -1494,6 +1495,49 @@ func TestProcessLLMRequestError(t *testing.T) {
}
}

func TestProcessLLMRequestShutdownPersistsError(t *testing.T) {
ctx, cancel := context.WithCancelCause(context.WithValue(context.Background(), "test-key", "test-value"))
service := &cancelingLLMService{cancel: cancel}

var recorded llm.Message
var recordCount int
var recordCtxErr error
var recordCtxValue any
var recordCtxHasDeadline bool
agentLoop := NewLoop(Config{
LLM: service,
RecordMessage: func(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error {
recordCount++
recorded = message
recordCtxErr = ctx.Err()
recordCtxValue = ctx.Value("test-key")
_, recordCtxHasDeadline = ctx.Deadline()
return nil
},
})
agentLoop.QueueUserMessage(llm.UserStringMessage("cancel this request"))

err := agentLoop.ProcessOneTurn(ctx)
if !errors.Is(err, context.Canceled) {
t.Fatalf("ProcessOneTurn error = %v, want context.Canceled", err)
}
if recordCount != 1 {
t.Fatalf("record count = %d, want 1", recordCount)
}
if recorded.Role != llm.MessageRoleAssistant || !recorded.EndOfTurn {
t.Fatalf("recorded message = %+v, want end-of-turn assistant error", recorded)
}
if recordCtxErr != nil {
t.Fatalf("record context error = %v, want nil", recordCtxErr)
}
if recordCtxValue != "test-value" {
t.Fatalf("record context value = %v, want test-value", recordCtxValue)
}
if !recordCtxHasDeadline {
t.Fatal("record context has no durability deadline")
}
}

// errorLLMService is a test LLM service that always returns an error
type errorLLMService struct {
err error
Expand All @@ -1517,6 +1561,16 @@ func (e *errorLLMService) MaxImageBytes() int {
return 5 * 1024 * 1024
}

type cancelingLLMService struct {
errorLLMService
cancel context.CancelCauseFunc
}

func (s *cancelingLLMService) Do(ctx context.Context, req *llm.Request) (*llm.Response, error) {
s.cancel(ErrShutdown)
return nil, ctx.Err()
}

// retryableLLMService fails with a retryable error a specified number of times, then succeeds
type retryableLLMService struct {
failuresRemaining int
Expand Down Expand Up @@ -1872,6 +1926,69 @@ func TestExecuteToolCallsWithMissingTool(t *testing.T) {
}
}

func TestExecuteToolCallsShutdownPersistsResult(t *testing.T) {
var recorded llm.Message
var recordCount int
agentLoop := NewLoop(Config{
Tools: []*llm.Tool{},
RecordMessage: func(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error {
if err := ctx.Err(); err != nil {
return err
}
recordCount++
recorded = message
return nil
},
})

ctx, cancel := context.WithCancelCause(context.Background())
cancel(ErrShutdown)
err := agentLoop.executeToolCalls(ctx, []llm.Content{{
ID: "cancelled-tool",
Type: llm.ContentTypeToolUse,
ToolName: "missing-tool",
}})
if err != nil {
t.Fatalf("executeToolCalls error = %v", err)
}
if recordCount != 1 {
t.Fatalf("record count = %d, want 1", recordCount)
}
if recorded.Role != llm.MessageRoleUser || len(recorded.Content) != 1 {
t.Fatalf("recorded message = %+v, want one user tool result", recorded)
}
if got := recorded.Content[0]; got.Type != llm.ContentTypeToolResult || got.ToolUseID != "cancelled-tool" {
t.Fatalf("recorded content = %+v, want cancelled-tool result", got)
}
}

func TestExecuteToolCallsUserCancellationDoesNotPersistLateResult(t *testing.T) {
var persisted bool
agentLoop := NewLoop(Config{
Tools: []*llm.Tool{},
RecordMessage: func(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error {
if err := ctx.Err(); err != nil {
return err
}
persisted = true
return nil
},
})

ctx, cancel := context.WithCancel(context.Background())
cancel()
if err := agentLoop.executeToolCalls(ctx, []llm.Content{{
ID: "cancelled-tool",
Type: llm.ContentTypeToolUse,
ToolName: "missing-tool",
}}); err != nil {
t.Fatalf("executeToolCalls error = %v", err)
}
if persisted {
t.Fatal("late tool result persisted after user cancellation")
}
}

func TestExecuteToolCallsWithErrorTool(t *testing.T) {
var recordedMessages []llm.Message
recordFunc := func(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error {
Expand Down
Loading