Files
probo/pkg/agent/run.go
Aurélien Sibiril c4228e8e7c Drive graceful agent suspend from ctx cancellation
Collapse the dual-mechanism (ctx.Done() = abort + WithStopSignal =
graceful suspend) into a single signal: ctx.Done() now means
graceful suspend. coreLoop shadows the incoming ctx with
context.WithoutCancel(ctx) on entry and uses the shadow for every
downstream call (LLM, tools, hooks, guardrails, save), keeping the
original ctx only for the at-boundary cancellation check.
restoreNestedSuspended applies the same shadow to its
saveProgress closure so partial nested-restore writes survive a
graceful cancel. Resume and resumeNested mirror the pattern so
their pre-loop tool dispatch is non-cancellable while coreLoop
still detects the cancel at its first turn boundary. The dedicated
stop signal API (WithStopSignal / stopSignalFrom) is removed.

There is no longer an in-process hard-abort path; tool authors
who need a deadline must derive it themselves. Document the new
contract on Run, RunStreamed, Resume, and Restore.

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>
2026-05-08 12:49:03 +02:00

1563 lines
43 KiB
Go

// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted, provided that the above
// copyright notice and this permission notice appear in all copies.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
// PERFORMANCE OF THIS SOFTWARE.
package agent
import (
"context"
"errors"
"fmt"
"sync"
"go.gearno.de/kit/log"
"go.opentelemetry.io/otel"
"go.opentelemetry.io/otel/attribute"
"go.opentelemetry.io/otel/codes"
"go.opentelemetry.io/otel/trace"
"go.probo.inc/probo/pkg/llm"
)
const (
tracerName = "go.probo.inc/probo/pkg/agent"
// synthesisNudge is the static user message appended after tool
// exploration completes, asking the model to produce the final
// structured output on the next (synthesis) turn.
synthesisNudge = "Based on everything you have gathered, produce the final structured output now."
)
type (
CallLLMFunc func(ctx context.Context, agent *Agent, req *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error)
RunOption func(*runOpts)
runOpts struct {
callLLM CallLLMFunc
onEvent func(ctx context.Context, ev StreamEvent)
skipInputGuardrails bool
skipSessionLoad bool
initialUsage llm.Usage
initialTurns int
checkpointer Checkpointer
runID string
toolUsedInRun bool
}
loopState struct {
agent *Agent
toolMap map[string]ToolDescriptor
toolDefs []llm.Tool
messages []llm.Message
inputMessages []llm.Message
systemPrompt string
totalUsage llm.Usage
turns int
toolUsedInRun bool
tracer trace.Tracer
runSpan trace.Span
opts runOpts
logger *log.Logger
}
parallelToolEntry struct {
result ToolResult
err error
}
)
func WithCheckpointer(cp Checkpointer, runID string) RunOption {
return func(o *runOpts) {
o.checkpointer = cp
o.runID = runID
}
}
func noopEvent(_ context.Context, _ StreamEvent) {}
func blockingCallLLM(ctx context.Context, agent *Agent, req *llm.ChatCompletionRequest) (*llm.ChatCompletionResponse, error) {
resp, err := agent.client.ChatCompletion(ctx, req)
if err == nil {
return resp, nil
}
// Some providers (e.g. Anthropic) require streaming for large
// max_tokens or when thinking is enabled. Fall back to streaming
// transparently when the blocking call returns ErrStreamingRequired.
var streamRequired *llm.ErrStreamingRequired
if !errors.As(err, &streamRequired) {
return nil, err
}
stream, sErr := agent.client.ChatCompletionStream(ctx, req)
if sErr != nil {
return nil, err // return the original error
}
defer func() { _ = stream.Close() }()
acc := llm.NewStreamAccumulator(stream)
for acc.Next() {
}
if sErr := acc.Err(); sErr != nil {
return nil, sErr
}
return acc.Response(), nil
}
// Run executes the agent loop. ctx is a graceful-stop signal: when it is
// cancelled the loop checkpoints at its next safe boundary and returns
// a *SuspendedError. The in-flight LLM call, the in-flight tool, and
// the checkpoint write all run on a non-cancellable shadow of ctx so
// they complete normally. ctx deadlines therefore become a "max
// wall-clock budget, then suspend" rather than a hard cut-off; there
// is no in-process hard-abort path. Callers that need to truly kill a
// run must terminate the process.
func (a *Agent) Run(ctx context.Context, messages []llm.Message, opts ...RunOption) (*Result, error) {
ro := runOpts{
callLLM: blockingCallLLM,
onEvent: noopEvent,
}
for _, opt := range opts {
opt(&ro)
}
return coreLoop(ctx, a, messages, ro)
}
func (s *loopState) resolveAgentTools(ctx context.Context) error {
tools, toolMap, err := s.agent.resolveTools(ctx)
if err != nil {
return err
}
toolDefs := make([]llm.Tool, len(tools))
for i, t := range tools {
toolDefs[i] = t.Definition()
}
s.toolMap = toolMap
s.toolDefs = toolDefs
return nil
}
func (s *loopState) finishRun(ctx context.Context, result *Result, err error) (*Result, error) {
defer s.runSpan.End()
defer func() {
emitHook(s.agent, func(h RunHooks) { h.OnRunEnd(ctx, s.agent, result, err) })
}()
if err != nil {
if _, ok := errors.AsType[*SuspendedError](err); ok {
s.runSpan.SetAttributes(attribute.Bool("agent.suspended", true))
s.opts.onEvent(ctx, StreamEvent{Type: StreamEventSuspended, Agent: s.agent})
s.logger.InfoCtx(
ctx,
"agent run suspended",
log.String("agent", s.agent.name),
log.Int("turns", s.turns),
)
return result, err
}
s.runSpan.RecordError(err)
s.runSpan.SetStatus(codes.Error, err.Error())
s.opts.onEvent(ctx, StreamEvent{Type: StreamEventError, Agent: s.agent, Err: err})
s.logger.ErrorCtx(
ctx,
"agent run failed",
log.String("agent", s.agent.name),
log.Int("turns", s.turns),
log.Error(err),
)
return result, err
}
if s.agent.session != nil {
if saveErr := s.agent.session.Save(ctx, s.agent.sessionID, result.Messages); saveErr != nil {
err = fmt.Errorf("cannot save session: %w", saveErr)
result = nil
s.runSpan.RecordError(err)
s.runSpan.SetStatus(codes.Error, err.Error())
s.opts.onEvent(ctx, StreamEvent{Type: StreamEventError, Agent: s.agent, Err: err})
s.logger.ErrorCtx(
ctx,
"cannot save session",
log.String("agent", s.agent.name),
log.Error(err),
)
return result, err
}
}
s.runSpan.SetAttributes(
attribute.Int("agent.turns", result.Turns),
attribute.Int("agent.usage.input_tokens", result.Usage.InputTokens),
attribute.Int("agent.usage.output_tokens", result.Usage.OutputTokens),
)
s.logger.InfoCtx(
ctx,
"agent run completed",
log.String("agent", s.agent.name),
log.Int("turns", result.Turns),
log.Int("input_tokens", result.Usage.InputTokens),
log.Int("output_tokens", result.Usage.OutputTokens),
)
s.opts.onEvent(ctx, StreamEvent{Type: StreamEventComplete, Agent: s.agent, Result: result})
return result, err
}
func (s *loopState) buildCheckpoint(status AgentStatus) *Checkpoint {
msgsCopy := make([]llm.Message, len(s.messages))
copy(msgsCopy, s.messages)
return &Checkpoint{
Status: status,
AgentName: s.agent.name,
Config: AgentConfig{
MaxTurns: s.agent.maxTurns,
},
Messages: msgsCopy,
Usage: s.totalUsage,
Turns: s.turns,
ToolUsedInRun: s.toolUsedInRun,
}
}
func (s *loopState) applyHandoff(ctx context.Context, handoffTarget *Handoff) error {
emitHook(s.agent, func(h RunHooks) { h.OnHandoff(ctx, s.agent, handoffTarget.Agent) })
emitAgentHook(handoffTarget.Agent, func(h AgentHooks) { h.OnHandoff(ctx, handoffTarget.Agent, s.agent) })
s.opts.onEvent(ctx, StreamEvent{Type: StreamEventHandoff, Agent: handoffTarget.Agent})
s.logger.InfoCtx(
ctx,
"agent handoff",
log.String("from", s.agent.name),
log.String("to", handoffTarget.Agent.name),
)
if handoffTarget.InputFilter != nil {
s.messages = handoffTarget.InputFilter(
HandoffInputData{
InputHistory: s.inputMessages,
NewItems: s.messages[len(s.inputMessages):],
},
)
}
s.agent = handoffTarget.Agent
s.logger = s.agent.logger
if err := s.resolveAgentTools(ctx); err != nil {
return err
}
emitAgentHook(s.agent, func(h AgentHooks) { h.OnStart(ctx, s.agent) })
s.opts.onEvent(ctx, StreamEvent{Type: StreamEventAgentStart, Agent: s.agent})
s.systemPrompt = s.agent.buildSystemPrompt(ctx)
s.toolUsedInRun = false
return nil
}
func coreLoop(ctx context.Context, startAgent *Agent, inputMessages []llm.Message, opts runOpts) (*Result, error) {
// Cancellation contract: ctx.Done() means "graceful suspend" — the
// loop checkpoints at its next safe boundary and returns
// SuspendedError. Every downstream call (LLM, tools, hooks,
// guardrails, save) runs on a non-cancellable shadow so an in-flight
// turn completes naturally and the checkpoint write itself isn't
// killed by the cancel that triggered the suspend. outerCtx is
// retained only for the at-boundary cancellation check.
outerCtx := ctx
ctx = context.WithoutCancel(ctx)
s := &loopState{
agent: startAgent,
inputMessages: inputMessages,
systemPrompt: startAgent.buildSystemPrompt(ctx),
totalUsage: opts.initialUsage,
turns: opts.initialTurns,
toolUsedInRun: opts.toolUsedInRun,
tracer: otel.GetTracerProvider().Tracer(tracerName),
opts: opts,
logger: startAgent.logger,
}
s.messages = make([]llm.Message, 0, len(inputMessages))
if !opts.skipSessionLoad && s.agent.session != nil {
prev, err := s.agent.session.Load(ctx, s.agent.sessionID)
if err != nil {
return nil, fmt.Errorf("cannot load session: %w", err)
}
s.logger.InfoCtx(
ctx,
"session loaded",
log.String("session_id", s.agent.sessionID),
log.Int("message_count", len(prev)),
)
s.messages = append(s.messages, prev...)
}
s.messages = append(s.messages, inputMessages...)
if !opts.skipInputGuardrails {
if err := runInputGuardrails(ctx, s.agent, s.messages); err != nil {
opts.onEvent(ctx, StreamEvent{Type: StreamEventError, Agent: s.agent, Err: err})
s.logger.ErrorCtx(
ctx,
"input guardrail tripped",
log.Error(err),
)
return nil, err
}
}
emitHook(s.agent, func(h RunHooks) { h.OnRunStart(ctx, s.agent, s.messages) })
emitAgentHook(s.agent, func(h AgentHooks) { h.OnStart(ctx, s.agent) })
opts.onEvent(ctx, StreamEvent{Type: StreamEventAgentStart, Agent: s.agent})
ctx, s.runSpan = s.tracer.Start(
ctx,
fmt.Sprintf("agent.run %s", s.agent.name),
trace.WithSpanKind(trace.SpanKindInternal),
trace.WithAttributes(
attribute.String("agent.name", s.agent.name),
attribute.String("agent.model", s.agent.model),
),
)
if err := s.resolveAgentTools(ctx); err != nil {
s.runSpan.End()
return nil, err
}
s.logger.InfoCtx(
ctx,
"agent run started",
log.String("model", s.agent.model),
log.Int("max_turns", s.agent.maxTurns),
log.Int("tool_count", len(s.toolDefs)),
)
emptyOutputRetries := 0
structuredFormat := resolveStructuredFormat(s.agent)
// When the agent has both tools and a structured output request,
// we delay structured output enforcement until a dedicated
// synthesis turn. Enforcing the schema during tool exploration
// causes models with extended thinking to stuff planning prose
// into the first text field of the schema as a scratchpad,
// burning the entire max_tokens budget on thinking-inside-JSON
// before ever producing a valid object. Instead, we let the
// model freely call tools without a schema, then force one final
// synthesis turn with ToolChoice=none + schema enforced once the
// model signals it has enough information (finish_reason=stop).
// Agents without tools or without a structured output request
// do not need this dance and enforce the schema immediately.
exploring := structuredFormat != nil && len(s.toolDefs) > 0
for {
if err := outerCtx.Err(); err != nil {
cp := s.buildCheckpoint(AgentStatusSuspended)
se := &SuspendedError{RunID: s.opts.runID}
if s.opts.checkpointer != nil && s.opts.runID != "" {
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save suspension checkpoint", log.Error(saveErr))
se.Checkpoint = cp
} else {
emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) })
}
} else {
se.Checkpoint = cp
}
return s.finishRun(ctx, nil, se)
}
if s.turns >= s.agent.maxTurns {
return s.finishRun(ctx, nil, &MaxTurnsExceededError{MaxTurns: s.agent.maxTurns})
}
fullMessages := buildFullMessages(s.systemPrompt, s.messages)
var responseFormat *llm.ResponseFormat
if !exploring {
responseFormat = structuredFormat
}
toolChoice := s.agent.modelSettings.ToolChoice
if s.toolUsedInRun && s.agent.resetToolChoice && toolChoice != nil {
toolChoice = nil
}
if !exploring && structuredFormat != nil && len(s.toolDefs) > 0 {
// On the synthesis turn, forbid further tool calls so the
// model is forced to convert what it has into JSON.
none := llm.ToolChoice{Type: llm.ToolChoiceNone}
toolChoice = &none
}
req := &llm.ChatCompletionRequest{
Model: s.agent.model,
Messages: fullMessages,
Tools: s.toolDefs,
Temperature: s.agent.modelSettings.Temperature,
TopP: s.agent.modelSettings.TopP,
FrequencyPenalty: s.agent.modelSettings.FrequencyPenalty,
PresencePenalty: s.agent.modelSettings.PresencePenalty,
MaxTokens: s.agent.modelSettings.MaxTokens,
ToolChoice: toolChoice,
ParallelToolCalls: s.agent.modelSettings.ParallelToolCalls,
ResponseFormat: responseFormat,
Thinking: s.agent.modelSettings.Thinking,
}
s.logger.InfoCtx(
ctx,
"calling LLM",
log.Int("turn", s.turns+1),
log.Int("message_count", len(fullMessages)),
)
resp, err := callLLMWithHooks(ctx, s.agent, req, opts)
if err != nil {
return s.finishRun(ctx, nil, fmt.Errorf("cannot call LLM: %w", err))
}
s.totalUsage = s.totalUsage.Add(resp.Usage)
s.turns++
s.logger.InfoCtx(
ctx,
"LLM response received",
log.Int("turn", s.turns),
log.String("finish_reason", string(resp.FinishReason)),
log.Int("input_tokens", resp.Usage.InputTokens),
log.Int("output_tokens", resp.Usage.OutputTokens),
)
s.messages = append(s.messages, resp.Message)
switch resp.FinishReason {
case llm.FinishReasonStop, llm.FinishReasonLength:
// Model signalled it has nothing more to do with tools.
// If we have a structured output request but haven't
// enforced the schema yet, promote this turn to the
// synthesis turn: the next iteration runs with
// ToolChoice=none and the schema enforced, so the model
// converts what it has gathered into JSON in one shot.
//
// Anthropic requires the last message in the conversation
// to be a user message, so we cannot simply continue after
// an assistant stop turn. Drop empty (thinking-only) turns
// from history and append a user nudge that asks for the
// final structured output. Non-empty assistant turns stay
// in history so the model can reference its own
// conclusions during synthesis.
if exploring && s.turns < s.agent.maxTurns {
exploring = false
if resp.Message.Text() == "" {
s.messages = s.messages[:len(s.messages)-1]
}
s.messages = append(
s.messages,
llm.Message{
Role: llm.RoleUser,
Parts: []llm.Part{llm.TextPart{Text: synthesisNudge}},
},
)
s.logger.WarnCtx(
ctx,
"entering synthesis turn: forcing structured output with tool_choice=none",
log.Int("turn", s.turns),
log.Int("output_tokens", resp.Usage.OutputTokens),
)
continue
}
// Anthropic extended-thinking models can return a synthesis turn
// that contains only thinking blocks and no text part, leaving us
// with no structured output to validate. Retry the same turn a
// bounded number of times so the model gets another chance to
// emit the required JSON output. The empty assistant turn must be
// dropped from history because Anthropic rejects requests where
// the last message is a thinking-only assistant turn.
if structuredFormat != nil && resp.Message.Text() == "" && emptyOutputRetries < s.agent.maxEmptyOutputRetries && s.turns < s.agent.maxTurns {
emptyOutputRetries++
s.messages = s.messages[:len(s.messages)-1]
s.logger.WarnCtx(
ctx,
"retrying turn: structured output expected but got empty text",
log.Int("turn", s.turns),
log.Int("retry", emptyOutputRetries),
log.Int("output_tokens", resp.Usage.OutputTokens),
)
continue
}
if err := runOutputGuardrails(ctx, s.agent, resp.Message); err != nil {
return s.finishRun(ctx, nil, err)
}
result := &Result{
Messages: s.messages,
Usage: s.totalUsage,
Turns: s.turns,
LastAgent: s.agent,
}
emitAgentHook(s.agent, func(h AgentHooks) { h.OnEnd(ctx, s.agent, resp.Message.Text()) })
opts.onEvent(ctx, StreamEvent{Type: StreamEventAgentEnd, Agent: s.agent})
return s.finishRun(ctx, result, nil)
case llm.FinishReasonToolCalls:
s.toolUsedInRun = true
emptyOutputRetries = 0
s.logger.InfoCtx(
ctx,
"dispatching tool calls",
log.Int("count", len(resp.Message.ToolCalls)),
)
handoffTarget, toolResults, toolMsgs, err := dispatchToolCalls(
ctx,
s.tracer,
s.agent,
resp.Message.ToolCalls,
s.toolMap,
opts.onEvent,
s.logger,
)
s.messages = append(s.messages, toolMsgs...)
if err != nil {
if se, ok := errors.AsType[*SuspendedError](err); ok {
outerCP := s.buildCheckpoint(AgentStatusSuspended)
if se.Checkpoint != nil {
outerCP.AllToolCalls = se.Checkpoint.AllToolCalls
outerCP.InnerCheckpoints = se.Checkpoint.InnerCheckpoints
outerCP.CompletedCalls = se.Checkpoint.CompletedCalls
}
if s.opts.checkpointer != nil && s.opts.runID != "" {
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, outerCP); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr))
} else {
emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, outerCP) })
}
}
return s.finishRun(ctx, nil, &SuspendedError{RunID: s.opts.runID, Checkpoint: outerCP})
}
if nae, ok := errors.AsType[*needsApprovalError](err); ok {
s.logger.InfoCtx(
ctx,
"run interrupted, approval required",
log.Int("pending_count", len(nae.pendingApprovals)),
)
msgsCopy := make([]llm.Message, len(s.messages))
copy(msgsCopy, s.messages)
if s.opts.checkpointer != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(AgentStatusAwaitingApproval)
cp.PendingToolCalls = nae.allToolCalls
cp.PendingApprovals = nae.pendingApprovals
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save approval checkpoint", log.Error(saveErr))
} else {
emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) })
}
}
return s.finishRun(
ctx,
nil,
&InterruptedError{
ToolCalls: nae.allToolCalls,
PendingApprovals: nae.pendingApprovals,
Agent: s.agent,
Messages: msgsCopy,
Usage: s.totalUsage,
Turns: s.turns,
},
)
}
if nie, ok := errors.AsType[*nestedInterruptionError](err); ok {
s.logger.InfoCtx(
ctx,
"nested agent interrupted, approval required",
log.Int("pending_count", len(nie.inner.PendingApprovals)),
log.String("nested_agent", nie.inner.Agent.name),
)
msgsCopy := make([]llm.Message, len(s.messages))
copy(msgsCopy, s.messages)
if s.opts.checkpointer != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(AgentStatusAwaitingApproval)
cp.PendingToolCalls = nie.inner.ToolCalls
cp.PendingApprovals = nie.inner.PendingApprovals
cp.AllToolCalls = nie.allToolCalls
cp.CompletedCalls = nie.completedCalls
cp.InnerCheckpoints = map[string]*Checkpoint{
nie.toolCallID: {
Status: AgentStatusAwaitingApproval,
AgentName: nie.inner.Agent.name,
Messages: nie.inner.Messages,
Usage: nie.inner.Usage,
Turns: nie.inner.Turns,
PendingToolCalls: nie.inner.ToolCalls,
PendingApprovals: nie.inner.PendingApprovals,
},
}
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save nested approval checkpoint", log.Error(saveErr))
} else {
emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) })
}
}
return s.finishRun(
ctx,
nil,
&InterruptedError{
ToolCalls: nie.inner.ToolCalls,
PendingApprovals: nie.inner.PendingApprovals,
Agent: nie.inner.Agent,
Messages: nie.inner.Messages,
Usage: nie.inner.Usage,
Turns: nie.inner.Turns,
outerState: &outerLoopState{
agent: s.agent,
messages: msgsCopy,
usage: s.totalUsage,
turns: s.turns,
allToolCalls: nie.allToolCalls,
toolCallID: nie.toolCallID,
completedCalls: nie.completedCalls,
innerInterrupt: nie.inner,
},
},
)
}
return s.finishRun(ctx, nil, err)
}
finalOutput, isFinal, behaviorErr := s.agent.toolUseBehavior(ctx, toolResults)
if behaviorErr != nil {
return s.finishRun(ctx, nil, fmt.Errorf("cannot evaluate tool use behavior: %w", behaviorErr))
}
if isFinal {
s.messages = append(
s.messages,
llm.Message{
Role: llm.RoleAssistant,
Parts: []llm.Part{llm.TextPart{Text: finalOutput}},
},
)
result := &Result{
Messages: s.messages,
Usage: s.totalUsage,
Turns: s.turns,
LastAgent: s.agent,
}
emitAgentHook(s.agent, func(h AgentHooks) { h.OnEnd(ctx, s.agent, finalOutput) })
opts.onEvent(ctx, StreamEvent{Type: StreamEventAgentEnd, Agent: s.agent})
return s.finishRun(ctx, result, nil)
}
if handoffTarget != nil {
if handoffErr := s.applyHandoff(ctx, handoffTarget); handoffErr != nil {
return s.finishRun(ctx, nil, handoffErr)
}
}
// Save incremental checkpoint after completed tool-call turn.
if s.opts.checkpointer != nil && s.opts.runID != "" {
cp := s.buildCheckpoint(AgentStatusSuspended)
if saveErr := s.opts.checkpointer.Save(ctx, s.opts.runID, cp); saveErr != nil {
s.logger.ErrorCtx(ctx, "cannot save checkpoint", log.Error(saveErr))
} else {
emitHook(s.agent, func(h RunHooks) { h.OnRunSnapshot(ctx, s.agent, cp) })
}
}
case llm.FinishReasonContentFilter:
return s.finishRun(ctx, nil, fmt.Errorf("cannot complete: content was filtered by the provider"))
default:
return s.finishRun(ctx, nil, fmt.Errorf("cannot complete: unexpected finish reason %q", resp.FinishReason))
}
}
}
func buildFullMessages(systemPrompt string, messages []llm.Message) []llm.Message {
fullMessages := make([]llm.Message, 0, len(messages)+1)
if systemPrompt != "" {
fullMessages = append(
fullMessages,
llm.Message{
Role: llm.RoleSystem,
Parts: []llm.Part{llm.TextPart{Text: systemPrompt}},
},
)
}
return append(fullMessages, messages...)
}
func callLLMWithHooks(
ctx context.Context,
agent *Agent,
req *llm.ChatCompletionRequest,
opts runOpts,
) (*llm.ChatCompletionResponse, error) {
emitHook(agent, func(h RunHooks) { h.OnLLMStart(ctx, agent, req.Messages) })
emitAgentHook(agent, func(h AgentHooks) { h.OnLLMStart(ctx, agent, req.Messages) })
resp, err := opts.callLLM(ctx, agent, req)
if err != nil {
emitHook(agent, func(h RunHooks) { h.OnLLMEnd(ctx, agent, nil, err) })
emitAgentHook(agent, func(h AgentHooks) { h.OnLLMEnd(ctx, agent, nil, err) })
return nil, err
}
emitHook(agent, func(h RunHooks) { h.OnLLMEnd(ctx, agent, resp, nil) })
emitAgentHook(agent, func(h AgentHooks) { h.OnLLMEnd(ctx, agent, resp, nil) })
return resp, nil
}
func dispatchToolCalls(
ctx context.Context,
tracer trace.Tracer,
agent *Agent,
toolCalls []llm.ToolCall,
toolMap map[string]ToolDescriptor,
onEvent func(context.Context, StreamEvent),
logger *log.Logger,
) (*Handoff, []ToolCallResult, []llm.Message, error) {
if err := checkApproval(ctx, agent, toolCalls); err != nil {
return nil, nil, nil, err
}
return executeToolCalls(ctx, tracer, agent, toolCalls, toolMap, onEvent, logger)
}
func executeToolCalls(
ctx context.Context,
tracer trace.Tracer,
agent *Agent,
toolCalls []llm.ToolCall,
toolMap map[string]ToolDescriptor,
onEvent func(context.Context, StreamEvent),
logger *log.Logger,
) (*Handoff, []ToolCallResult, []llm.Message, error) {
descriptors := make([]ToolDescriptor, len(toolCalls))
handoffIdx := -1
for i, tc := range toolCalls {
desc, ok := toolMap[tc.Function.Name]
if !ok {
return nil, nil, nil, fmt.Errorf("cannot dispatch tool call: unknown tool %q", tc.Function.Name)
}
descriptors[i] = desc
if _, isHandoff := desc.(*handoffToolAdapter); isHandoff && handoffIdx == -1 {
handoffIdx = i
}
}
if handoffIdx >= 0 {
return executeWithHandoff(ctx, tracer, agent, toolCalls, descriptors, handoffIdx, onEvent, logger)
}
tools := make([]Tool, len(descriptors))
for i, d := range descriptors {
tools[i] = d.(Tool)
}
results, msgs, err := executeParallel(ctx, tracer, agent, toolCalls, tools, onEvent, logger)
return nil, results, msgs, err
}
func executeWithHandoff(
ctx context.Context,
tracer trace.Tracer,
agent *Agent,
toolCalls []llm.ToolCall,
descriptors []ToolDescriptor,
handoffIdx int,
onEvent func(context.Context, StreamEvent),
logger *log.Logger,
) (*Handoff, []ToolCallResult, []llm.Message, error) {
var (
results []ToolCallResult
msgs []llm.Message
)
for i := range handoffIdx {
logger.InfoCtx(
ctx,
"executing tool before handoff",
log.String("tool", toolCalls[i].Function.Name),
)
tr, err := executeSingleTool(ctx, tracer, agent, toolCalls[i], descriptors[i].(Tool), onEvent, logger)
if err != nil {
if ie, ok := errors.AsType[*InterruptedError](err); ok {
var completed []CompletedCall
for j := range results {
completed = append(
completed,
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: results[j].Result,
},
)
}
return nil, nil, msgs, &nestedInterruptionError{
inner: ie,
toolCallID: toolCalls[i].ID,
allToolCalls: toolCalls,
completedCalls: completed,
}
}
return nil, nil, msgs, err
}
msgs = append(
msgs,
llm.Message{
Role: llm.RoleTool,
ToolCallID: toolCalls[i].ID,
Parts: []llm.Part{llm.TextPart{Text: tr.Content}},
},
)
results = append(
results,
ToolCallResult{
ToolName: toolCalls[i].Function.Name,
Arguments: toolCalls[i].Function.Arguments,
Result: tr,
},
)
}
ht := descriptors[handoffIdx].(*handoffToolAdapter)
if ht.handoff.OnHandoff != nil {
if err := ht.handoff.OnHandoff(ctx); err != nil {
return nil, nil, msgs, fmt.Errorf("cannot execute handoff callback for %q: %w", ht.handoff.Agent.name, err)
}
}
msgs = append(
msgs,
llm.Message{
Role: llm.RoleTool,
ToolCallID: toolCalls[handoffIdx].ID,
Parts: []llm.Part{llm.TextPart{Text: fmt.Sprintf("Transferred to %s", ht.handoff.Agent.name)}},
},
)
for i := handoffIdx + 1; i < len(toolCalls); i++ {
msgs = append(
msgs,
llm.Message{
Role: llm.RoleTool,
ToolCallID: toolCalls[i].ID,
Parts: []llm.Part{
llm.TextPart{Text: "Tool call was not executed because a handoff occurred."},
},
},
)
}
return ht.handoff, results, msgs, nil
}
func executeParallel(
ctx context.Context,
tracer trace.Tracer,
agent *Agent,
toolCalls []llm.ToolCall,
tools []Tool,
onEvent func(context.Context, StreamEvent),
logger *log.Logger,
) ([]ToolCallResult, []llm.Message, error) {
entries := make([]parallelToolEntry, len(toolCalls))
var wg sync.WaitGroup
wg.Add(len(toolCalls))
for i := range toolCalls {
go func(idx int, tc llm.ToolCall, tool Tool) {
defer wg.Done()
tr, err := executeSingleTool(ctx, tracer, agent, tc, tool, onEvent, logger)
if err != nil {
entries[idx] = parallelToolEntry{err: err}
return
}
entries[idx] = parallelToolEntry{result: tr}
}(i, toolCalls[i], tools[i])
}
wg.Wait()
for i, entry := range entries {
ie, ok := errors.AsType[*InterruptedError](entry.err)
if !ok {
continue
}
var completed []CompletedCall
for j, other := range entries {
if j == i {
continue
}
if other.err != nil {
completed = append(
completed,
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: ToolResult{
Content: fmt.Sprintf("Error: %s", other.err.Error()),
IsError: true,
},
},
)
continue
}
completed = append(
completed,
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: other.result,
},
)
}
return nil, nil, &nestedInterruptionError{
inner: ie,
toolCallID: toolCalls[i].ID,
allToolCalls: toolCalls,
completedCalls: completed,
}
}
// Check for suspended inner agents (stop signal propagated).
for i, entry := range entries {
if entry.err == nil {
continue
}
se, ok := errors.AsType[*SuspendedError](entry.err)
if ok && se.Checkpoint != nil {
innerCheckpoints := make(map[string]*Checkpoint)
var completed []CompletedCall
for j, other := range entries {
if j == i {
continue
}
if other.err == nil {
completed = append(
completed,
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: other.result,
},
)
continue
}
otherSE, ok := errors.AsType[*SuspendedError](other.err)
if ok && otherSE.Checkpoint != nil {
innerCheckpoints[toolCalls[j].ID] = otherSE.Checkpoint
} else {
completed = append(
completed,
CompletedCall{
ToolCallID: toolCalls[j].ID,
Result: ToolResult{
Content: fmt.Sprintf("Error: %s", other.err.Error()),
IsError: true,
},
},
)
}
}
innerCheckpoints[toolCalls[i].ID] = se.Checkpoint
outerSE := &SuspendedError{
Checkpoint: &Checkpoint{
Status: AgentStatusSuspended,
AllToolCalls: toolCalls,
InnerCheckpoints: innerCheckpoints,
CompletedCalls: completed,
},
}
return nil, nil, outerSE
}
}
var (
results []ToolCallResult
msgs []llm.Message
)
for i, tc := range toolCalls {
entry := entries[i]
if entry.err != nil {
msgs = append(
msgs,
llm.Message{
Role: llm.RoleTool,
ToolCallID: tc.ID,
Parts: []llm.Part{
llm.TextPart{
Text: fmt.Sprintf("Error: %s", entry.err.Error()),
},
},
},
)
results = append(
results,
ToolCallResult{
ToolName: tc.Function.Name,
Arguments: tc.Function.Arguments,
Result: ToolResult{
Content: fmt.Sprintf("Error: %s", entry.err.Error()),
IsError: true,
},
},
)
continue
}
msgs = append(
msgs,
llm.Message{
Role: llm.RoleTool,
ToolCallID: tc.ID,
Parts: []llm.Part{llm.TextPart{Text: entry.result.Content}},
},
)
results = append(
results,
ToolCallResult{
ToolName: tc.Function.Name,
Arguments: tc.Function.Arguments,
Result: entry.result,
},
)
}
return results, msgs, nil
}
func executeSingleTool(
ctx context.Context,
tracer trace.Tracer,
agent *Agent,
tc llm.ToolCall,
tool Tool,
onEvent func(context.Context, StreamEvent),
logger *log.Logger,
) (ToolResult, error) {
onEvent(ctx, StreamEvent{Type: StreamEventToolStart, Agent: agent, Tool: tool})
emitHook(agent, func(h RunHooks) { h.OnToolStart(ctx, agent, tool, tc.Function.Arguments) })
emitAgentHook(agent, func(h AgentHooks) { h.OnToolStart(ctx, agent, tool) })
toolCtx, toolSpan := tracer.Start(
ctx,
fmt.Sprintf("agent.tool %s", tool.Name()),
trace.WithSpanKind(trace.SpanKindInternal),
trace.WithAttributes(
attribute.String("tool.name", tool.Name()),
),
)
logger.InfoCtx(
ctx,
"executing tool",
log.String("tool", tool.Name()),
)
result, err := tool.Execute(toolCtx, tc.Function.Arguments)
if err != nil {
if _, ok := errors.AsType[*InterruptedError](err); ok {
toolSpan.SetAttributes(attribute.Bool("tool.interrupted", true))
toolSpan.End()
onEvent(ctx, StreamEvent{Type: StreamEventToolEnd, Agent: agent, Tool: tool, Err: err})
emitHook(agent, func(h RunHooks) { h.OnToolEnd(ctx, agent, tool, ToolResult{}, err) })
emitAgentHook(agent, func(h AgentHooks) { h.OnToolEnd(ctx, agent, tool, ToolResult{}) })
return ToolResult{}, err
}
if _, ok := errors.AsType[*SuspendedError](err); ok {
toolSpan.SetAttributes(attribute.Bool("tool.suspended", true))
toolSpan.End()
onEvent(ctx, StreamEvent{Type: StreamEventToolEnd, Agent: agent, Tool: tool, Err: err})
emitHook(agent, func(h RunHooks) { h.OnToolEnd(ctx, agent, tool, ToolResult{}, err) })
emitAgentHook(agent, func(h AgentHooks) { h.OnToolEnd(ctx, agent, tool, ToolResult{}) })
return ToolResult{}, err
}
toolSpan.RecordError(err)
toolSpan.SetStatus(codes.Error, err.Error())
toolSpan.End()
onEvent(ctx, StreamEvent{Type: StreamEventToolEnd, Agent: agent, Tool: tool, Err: err})
emitHook(agent, func(h RunHooks) { h.OnToolEnd(ctx, agent, tool, result, err) })
emitAgentHook(agent, func(h AgentHooks) { h.OnToolEnd(ctx, agent, tool, result) })
logger.ErrorCtx(
ctx,
"tool execution failed",
log.String("tool", tool.Name()),
log.Error(err),
)
return ToolResult{}, fmt.Errorf("cannot execute tool %q: %w", tool.Name(), err)
}
toolSpan.SetAttributes(attribute.Bool("tool.is_error", result.IsError))
toolSpan.End()
onEvent(ctx, StreamEvent{Type: StreamEventToolEnd, Agent: agent, Tool: tool, ToolResult: &result})
emitHook(agent, func(h RunHooks) { h.OnToolEnd(ctx, agent, tool, result, nil) })
emitAgentHook(agent, func(h AgentHooks) { h.OnToolEnd(ctx, agent, tool, result) })
if result.IsError {
content := result.Content
if len(content) > 200 {
content = content[:200] + "... (truncated)"
}
logger.WarnCtx(
ctx,
"tool returned error",
log.String("tool", tool.Name()),
log.String("content", content),
)
} else {
logger.InfoCtx(
ctx,
"tool execution completed",
log.String("tool", tool.Name()),
)
}
return result, nil
}
func checkApproval(ctx context.Context, a *Agent, toolCalls []llm.ToolCall) error {
if a.approval == nil {
return nil
}
var pending []llm.ToolCall
for _, tc := range toolCalls {
if a.approval.requiresApproval(ctx, tc) {
pending = append(pending, tc)
}
}
if len(pending) > 0 {
return &needsApprovalError{
allToolCalls: toolCalls,
pendingApprovals: pending,
}
}
return nil
}
func runInputGuardrails(ctx context.Context, agent *Agent, messages []llm.Message) error {
for _, g := range agent.inputGuardrails {
result, err := g.Check(ctx, messages)
if err != nil {
return fmt.Errorf("cannot run input guardrail %q: %w", g.Name(), err)
}
if result != nil && result.Tripwire {
emitHook(agent, func(h RunHooks) { h.OnGuardrailTripped(ctx, agent, g.Name(), result) })
return &InputGuardrailTrippedError{
Guardrail: g.Name(),
Message: result.Message,
}
}
}
return nil
}
func runOutputGuardrails(ctx context.Context, agent *Agent, message llm.Message) error {
for _, g := range agent.outputGuardrails {
result, err := g.Check(ctx, message)
if err != nil {
return fmt.Errorf("cannot run output guardrail %q: %w", g.Name(), err)
}
if result != nil && result.Tripwire {
emitHook(agent, func(h RunHooks) { h.OnGuardrailTripped(ctx, agent, g.Name(), result) })
return &OutputGuardrailTrippedError{
Guardrail: g.Name(),
Message: result.Message,
}
}
}
return nil
}
// Resume continues an interrupted run after human approval decisions have
// been collected. It executes or denies each pending tool call according to
// the provided ResumeInput, then re-enters the agent loop. Input guardrails
// are not re-evaluated because the messages were already validated in the
// original Run call.
//
// ctx follows the same graceful-suspend contract as Run: cancellation
// produces a SuspendedError with a checkpoint covering progress up to
// the next safe boundary. See Run for details.
func Resume(ctx context.Context, interrupted *InterruptedError, input ResumeInput, opts ...RunOption) (*Result, error) {
ro := runOpts{
callLLM: blockingCallLLM,
onEvent: noopEvent,
}
for _, opt := range opts {
opt(&ro)
}
return resumeWithOpts(ctx, interrupted, input, ro)
}
func resumeWithOpts(ctx context.Context, interrupted *InterruptedError, input ResumeInput, ro runOpts) (*Result, error) {
// Same cancellation contract as Run: ctx.Done() means graceful
// suspend. Pre-loop tool execution and handoff callbacks run on the
// non-cancellable shadow so an in-flight tool completes; coreLoop
// receives the original ctx so it can detect cancellation at its
// next turn boundary and trigger the suspend path.
outerCtx := ctx
ctx = context.WithoutCancel(ctx)
if interrupted.outerState != nil {
return resumeNested(outerCtx, interrupted, input, ro)
}
tracer := otel.GetTracerProvider().Tracer(tracerName)
logger := interrupted.Agent.logger
agent := interrupted.Agent
messages := interrupted.Messages
logger.InfoCtx(
ctx,
"resuming interrupted run",
log.String("agent", agent.name),
log.Int("pending_approvals", len(interrupted.PendingApprovals)),
log.Int("total_tool_calls", len(interrupted.ToolCalls)),
)
_, toolMap, err := agent.resolveTools(ctx)
if err != nil {
return nil, fmt.Errorf("cannot resolve tools for resume: %w", err)
}
pendingSet := make(map[string]struct{}, len(interrupted.PendingApprovals))
for _, tc := range interrupted.PendingApprovals {
pendingSet[tc.ID] = struct{}{}
}
var handoffTarget *Handoff
for _, tc := range interrupted.ToolCalls {
if _, needsApproval := pendingSet[tc.ID]; needsApproval {
approval, ok := input.Approvals[tc.ID]
if !ok || !approval.Approved {
reason := "Tool call was denied by human review."
if ok && approval.Message != "" {
reason = approval.Message
}
logger.InfoCtx(
ctx,
"tool call denied",
log.String("tool", tc.Function.Name),
log.String("tool_call_id", tc.ID),
log.String("reason", reason),
)
messages = append(
messages,
llm.Message{
Role: llm.RoleTool,
ToolCallID: tc.ID,
Parts: []llm.Part{
llm.TextPart{Text: reason},
},
},
)
continue
}
logger.InfoCtx(
ctx,
"tool call approved",
log.String("tool", tc.Function.Name),
log.String("tool_call_id", tc.ID),
)
}
desc, toolOK := toolMap[tc.Function.Name]
if !toolOK {
return nil, fmt.Errorf("cannot dispatch unknown tool %q", tc.Function.Name)
}
if ht, ok := desc.(*handoffToolAdapter); ok {
if ht.handoff.OnHandoff != nil {
if err := ht.handoff.OnHandoff(ctx); err != nil {
return nil, fmt.Errorf("cannot execute handoff callback for %q: %w", ht.handoff.Agent.name, err)
}
}
messages = append(
messages,
llm.Message{
Role: llm.RoleTool,
ToolCallID: tc.ID,
Parts: []llm.Part{
llm.TextPart{
Text: fmt.Sprintf("Transferred to %s", ht.handoff.Agent.name),
},
},
},
)
handoffTarget = ht.handoff
break
}
tr, execErr := executeSingleTool(ctx, tracer, agent, tc, desc.(Tool), ro.onEvent, logger)
if execErr != nil {
return nil, execErr
}
messages = append(
messages,
llm.Message{
Role: llm.RoleTool,
ToolCallID: tc.ID,
Parts: []llm.Part{
llm.TextPart{
Text: tr.Content,
},
},
},
)
}
resumeAgent := agent
if handoffTarget != nil {
logger.InfoCtx(
ctx,
"resume handoff",
log.String("from", agent.name),
log.String("to", handoffTarget.Agent.name),
)
emitHook(agent, func(h RunHooks) { h.OnHandoff(ctx, agent, handoffTarget.Agent) })
emitAgentHook(handoffTarget.Agent, func(h AgentHooks) { h.OnHandoff(ctx, handoffTarget.Agent, agent) })
if handoffTarget.InputFilter != nil {
filtered := handoffTarget.InputFilter(
HandoffInputData{
InputHistory: interrupted.Messages,
NewItems: messages[len(interrupted.Messages):],
},
)
messages = filtered
}
resumeAgent = handoffTarget.Agent
}
return coreLoop(
outerCtx,
resumeAgent,
messages,
runOpts{
callLLM: ro.callLLM,
onEvent: ro.onEvent,
skipInputGuardrails: true,
skipSessionLoad: true,
initialUsage: interrupted.Usage,
initialTurns: interrupted.Turns,
checkpointer: ro.checkpointer,
runID: ro.runID,
toolUsedInRun: ro.toolUsedInRun,
},
)
}
func resumeNested(ctx context.Context, interrupted *InterruptedError, input ResumeInput, ro runOpts) (*Result, error) {
// Same shadow as resumeWithOpts: pre-loop bridging work runs on
// the non-cancellable shadow; coreLoop receives outerCtx so it can
// detect a graceful-suspend cancel at its turn boundary.
outerCtx := ctx
ctx = context.WithoutCancel(ctx)
outer := interrupted.outerState
logger := outer.agent.logger
logger.InfoCtx(
ctx,
"resuming nested agent interruption",
log.String("outer_agent", outer.agent.name),
log.String("inner_agent", interrupted.Agent.name),
)
innerResult, err := resumeWithOpts(outerCtx, outer.innerInterrupt, input, ro)
if err != nil {
innerIE, ok := errors.AsType[*InterruptedError](err)
if ok {
return nil, &InterruptedError{
ToolCalls: innerIE.ToolCalls,
PendingApprovals: innerIE.PendingApprovals,
Agent: innerIE.Agent,
Messages: innerIE.Messages,
Usage: innerIE.Usage,
Turns: innerIE.Turns,
outerState: &outerLoopState{
agent: outer.agent,
messages: outer.messages,
usage: outer.usage,
turns: outer.turns,
allToolCalls: outer.allToolCalls,
toolCallID: outer.toolCallID,
completedCalls: outer.completedCalls,
innerInterrupt: innerIE,
},
}
}
return nil, fmt.Errorf("cannot resume nested agent: %w", err)
}
completedMap := make(map[string]ToolResult, len(outer.completedCalls))
for _, cc := range outer.completedCalls {
completedMap[cc.ToolCallID] = cc.Result
}
messages := make([]llm.Message, len(outer.messages))
copy(messages, outer.messages)
for _, tc := range outer.allToolCalls {
var content string
if tc.ID == outer.toolCallID {
content = innerResult.FinalMessage().Text()
} else if cr, ok := completedMap[tc.ID]; ok {
content = cr.Content
} else {
content = "Error: tool execution was interrupted"
}
messages = append(
messages,
llm.Message{
Role: llm.RoleTool,
ToolCallID: tc.ID,
Parts: []llm.Part{llm.TextPart{Text: content}},
},
)
}
return coreLoop(
outerCtx,
outer.agent,
messages,
runOpts{
callLLM: ro.callLLM,
onEvent: ro.onEvent,
skipInputGuardrails: true,
skipSessionLoad: true,
initialUsage: outer.usage,
initialTurns: outer.turns,
checkpointer: ro.checkpointer,
runID: ro.runID,
toolUsedInRun: ro.toolUsedInRun,
},
)
}
func emitHook(agent *Agent, fn func(RunHooks)) {
for _, h := range agent.hooks {
fn(h)
}
}
func emitAgentHook(agent *Agent, fn func(AgentHooks)) {
if agent.agentHooks != nil {
fn(agent.agentHooks)
}
}
// resolveStructuredFormat returns the structured output request the
// agent wants enforced on its final turn, or nil if none. An agent can
// declare structured output through either WithOutputType (typed
// sub-agents) or a directly-set responseFormat (the RunTyped
// convenience wrapper).
func resolveStructuredFormat(a *Agent) *llm.ResponseFormat {
if a.responseFormat != nil {
return a.responseFormat
}
if a.outputType != nil {
return a.outputType.responseFormat()
}
return nil
}